diff --git a/.gitignore b/.gitignore index 820833f..a6d25e3 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ models/ crates/adapters/el-ffi/src/frb_generated.rs # Claude Code local state (worktrees, settings.local.json) .claude/ +benchmarks/data/ diff --git a/Cargo.lock b/Cargo.lock index 55937bc..9f3edfa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -793,7 +793,7 @@ checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" [[package]] name = "el-bench" -version = "0.3.7" +version = "0.3.8" dependencies = [ "el-core", "el-engine-candle", diff --git a/crates/adapters/el-engine-candle/src/lib.rs b/crates/adapters/el-engine-candle/src/lib.rs index ab9a32a..756ea6d 100644 --- a/crates/adapters/el-engine-candle/src/lib.rs +++ b/crates/adapters/el-engine-candle/src/lib.rs @@ -170,6 +170,12 @@ impl InferenceEngine for CandleEngine { fn eos_token(&self) -> Token { self.eos } + + /// Stateless: this engine's forward is `embed[committed.last()] · w_out`, so + /// it holds no KV cache to restore — a rollback is a no-op. + fn rollback(&mut self, _keep_committed: u32) -> Result<()> { + Ok(()) + } } // ── LlmProvider (text-level) wrapper (ADR-010) ─────────────────────────────── @@ -343,9 +349,15 @@ mod bench { /// /// Holds candle's stateful KV cache. Within one generation it is fed /// incrementally (prefill, then one new token per `next_logits` call); candle -/// exposes no public cache reset, so a fresh conversation builds a new engine. -/// Float logits are quantised to integer milli-logits at the seam, exactly like -/// [`CandleEngine`], so the runtime stays float-free. +/// exposes no public cache reset, so a *fresh conversation* builds a new engine. +/// +/// A *within-generation* safety backtrack (ADR-012) is supported via +/// [`InferenceEngine::rollback`]: candle's attention discards its cache when a +/// forward runs at `index_pos == 0`, so we retain the prompt and replay it from +/// position 0 to rebuild the cache for the safe prefix (the session then +/// re-feeds the retained committed tokens). Float logits are quantised to +/// integer milli-logits at the seam, exactly like [`CandleEngine`], so the +/// runtime stays float-free. pub struct QwenEngine { model: Qwen2Weights, device: Device, @@ -353,6 +365,9 @@ pub struct QwenEngine { index_pos: usize, /// How many of the runtime-`committed` tokens have already been fed. fed: usize, + /// The prefill prompt, retained so a rollback can replay it from position 0 + /// to rebuild candle's KV cache (which has no public truncation). + prompt: Vec, /// Milli-logits produced after the most recent forward. last_logits: Vec, vocab: usize, @@ -375,6 +390,7 @@ impl QwenEngine { device, index_pos: 0, fed: 0, + prompt: Vec::new(), last_logits: Vec::new(), vocab: 0, eos, @@ -416,6 +432,7 @@ impl InferenceEngine for QwenEngine { fn prefill(&mut self, tokens: &[Token]) -> Result { self.index_pos = 0; self.fed = 0; + self.prompt = tokens.to_vec(); // retained for rollback replay for &t in tokens { self.last_logits = self.forward_one(t)?; } @@ -441,6 +458,24 @@ impl InferenceEngine for QwenEngine { fn eos_token(&self) -> Token { self.eos } + + fn rollback(&mut self, _keep_committed: u32) -> Result<()> { + // candle's KV cache is append-only with no public truncation, but its + // attention discards the cache on a forward at `index_pos == 0` (see + // quantized_qwen2). So rebuild deterministically: replay the prompt from + // position 0 — the first forward resets the cache, the rest re-append it — + // leaving the engine in its exact post-prefill state. We reset `fed` to 0 + // so the session's next `next_logits` re-feeds the retained committed + // prefix (already truncated to `keep_committed`) on top. Cost is bounded + // by `max_rollbacks` (ADR-012). + self.index_pos = 0; + self.fed = 0; + for i in 0..self.prompt.len() { + let t = self.prompt[i]; + self.last_logits = self.forward_one(t)?; + } + Ok(()) + } } /// A real local chat backend: a Qwen2 GGUF model + its tokenizer, driven diff --git a/crates/el-grammar/src/lib.rs b/crates/el-grammar/src/lib.rs index 2ae42fb..2a081ab 100644 --- a/crates/el-grammar/src/lib.rs +++ b/crates/el-grammar/src/lib.rs @@ -162,6 +162,9 @@ mod tests { fn eos_token(&self) -> Token { 99 } + fn rollback(&mut self, _keep: u32) -> el_core::Result<()> { + Ok(()) // stateless + } } struct OkVerifier; diff --git a/crates/el-memory/src/kv.rs b/crates/el-memory/src/kv.rs index 7be8872..b4f20d1 100644 --- a/crates/el-memory/src/kv.rs +++ b/crates/el-memory/src/kv.rs @@ -52,6 +52,14 @@ impl KvRegion { &self.slots } + /// Drop tail descriptors so the region retains only the first `len` slots. + /// `O(dropped)`; survivors keep their `offset` (no data copy). This is the + /// rollback primitive for the safety control loop (ADR-012): restoring a + /// checkpoint rewinds committed KV without replaying prefill. + pub fn truncate(&mut self, len: u32) { + self.slots.truncate(len as usize); + } + /// Remove pruned descriptors and re-index survivors. Returns how many were /// reclaimed. **Survivors keep their original `offset`** — proof that data /// is not moved, only descriptors are reshuffled. @@ -87,4 +95,18 @@ mod tests { // ...but its logical index was shuffled down. assert_eq!(kv.slots()[1].token_index, 1); } + + #[test] + fn truncate_rewinds_to_len_without_touching_survivors() { + let mut kv = KvRegion::new(); + kv.push(0); + kv.push(64); + kv.push(128); + let off1 = kv.slots()[1].offset; + kv.truncate(2); + assert_eq!(kv.len(), 2); + assert_eq!(kv.slots()[1].offset, off1); // survivor untouched (no data copy) + kv.truncate(0); + assert!(kv.is_empty()); + } } diff --git a/crates/el-runtime/src/defaults.rs b/crates/el-runtime/src/defaults.rs index fc71712..c1bc075 100644 --- a/crates/el-runtime/src/defaults.rs +++ b/crates/el-runtime/src/defaults.rs @@ -52,4 +52,9 @@ impl InferenceEngine for NullEngine { fn eos_token(&self) -> Token { self.eos } + + /// Stateless: `next_logits` ignores `committed`, so there is nothing to undo. + fn rollback(&mut self, _keep_committed: u32) -> Result<()> { + Ok(()) + } } diff --git a/crates/el-runtime/src/lib.rs b/crates/el-runtime/src/lib.rs index 2933378..aa55ee6 100644 --- a/crates/el-runtime/src/lib.rs +++ b/crates/el-runtime/src/lib.rs @@ -17,5 +17,8 @@ pub use defaults::{AllowAllMasker, IdentityCompressor, NullEngine}; pub use ports::{GrammarMasker, HybridRelay, InferenceEngine, Ports, PromptCompressor}; pub use session::InferenceSession; -// Re-export the safety port so callers wire one type system. -pub use el_safety::{LogitAdjustment, SafetySteerer}; +// Re-export the safety ports so callers wire one type system. +pub use el_safety::{ + Checkpoint, CheckpointManager, ChunkGuard, LogitAdjustment, RollbackPolicy, SafetyScore, + SafetySteerer, +}; diff --git a/crates/el-runtime/src/ports.rs b/crates/el-runtime/src/ports.rs index 341862d..cd7a4e4 100644 --- a/crates/el-runtime/src/ports.rs +++ b/crates/el-runtime/src/ports.rs @@ -1,7 +1,7 @@ //! Port traits the Inference Runtime depends on (the collaborator contexts). use el_core::{Result, Token}; -use el_safety::SafetySteerer; +use el_safety::{ChunkGuard, SafetySteerer}; /// The inference engine adapter (`RuntimeAcl`). Implemented for real by Candle /// in the excluded adapter `el-engine-candle` (ADR-002). @@ -15,6 +15,28 @@ pub trait InferenceEngine { fn next_logits(&mut self, committed: &[Token]) -> Vec; /// The end-of-sequence token id. fn eos_token(&self) -> Token; + + /// Roll the engine's internal state back so its context is exactly the + /// prompt plus `keep_committed` generated tokens. + /// + /// The ADR-012 control loop truncates the session's committed output and KV + /// descriptors on a safety backtrack. A **stateful** engine (one holding a + /// real KV cache and position counters, e.g. a transformer) must mirror that + /// truncation here — otherwise it keeps serving logits from the abandoned + /// (unsafe) branch and never re-feeds the replacement tokens, so the rollback + /// is silently a no-op at the engine level. + /// + /// After `Ok(())`, the next [`next_logits`](Self::next_logits) call — passed a + /// `committed` slice of length `keep_committed` — must produce logits + /// consistent with that prefix. Returning `Err` makes the loop fail closed + /// rather than resume on an inconsistent cache. + /// + /// This method is **required, with no default**, deliberately: a default + /// no-op would let a stateful adapter that forgot to override silently resume + /// on a stale KV cache — a safety bug that fails *open*. Every engine must + /// make the choice explicit. A stateless engine whose `next_logits` + /// recomputes purely from `committed` implements it as `Ok(())`. + fn rollback(&mut self, keep_committed: u32) -> Result<()>; } /// Prompt Compression port (LLMLingua-2 — context 2). @@ -40,17 +62,21 @@ pub struct Ports { pub compressor: Box, pub grammar: Box, pub safety: Box, + /// Optional chunk guard for the checkpointed-rollback control loop + /// (ADR-012). `None` runs the plain single-pass decode. + pub guard: Option>, pub relay: Option>, } impl Ports { /// Defaults: identity compression, all-tokens-allowed grammar, no safety - /// steering, no relay (air-gapped). + /// steering, no guard, no relay (air-gapped). pub fn permissive() -> Self { Self { compressor: Box::new(super::defaults::IdentityCompressor), grammar: Box::new(super::defaults::AllowAllMasker), safety: Box::new(el_safety::NoSafety), + guard: None, relay: None, } } diff --git a/crates/el-runtime/src/session.rs b/crates/el-runtime/src/session.rs index 6d64985..93e1ef1 100644 --- a/crates/el-runtime/src/session.rs +++ b/crates/el-runtime/src/session.rs @@ -2,12 +2,17 @@ use crate::ports::{InferenceEngine, Ports}; use el_core::{ - DomainEvent, EdgeError, EventEnvelope, Phase, Result, SessionConfig, SessionId, StopReason, - Token, + DegradeReason, DomainEvent, EdgeError, EventEnvelope, Phase, Result, SessionConfig, SessionId, + StopReason, Token, }; use el_memory::KvRegion; use el_provenance::LoadPermit; -use el_safety::LogitAdjustment; +use el_safety::{Checkpoint, CheckpointManager, LogitAdjustment, RollbackPolicy}; + +/// Below this static-memory budget there is no room to retain rollback +/// checkpoints, so the control loop degrades to guard-only (no rollback) and +/// emits `SafetyDisabled` (ADR-012 tier-aware degradation; ADR-003 budget). +const MIN_CHECKPOINT_BUDGET_BYTES: u64 = 64 * 1024 * 1024; /// One live generation. Constructing it requires a [`LoadPermit`], so a model /// that has not passed the provenance gate (ADR-006) cannot reach the runtime — @@ -24,6 +29,28 @@ pub struct InferenceSession { events: Vec, } +/// Outcome of one chunk-guard evaluation in the control loop (ADR-012). +enum GuardVerdict { + /// Below the hard threshold — safe (safe checkpoint advanced) or tolerated. + Pass, + /// Hard breach rolled back to a safe checkpoint; decoding should resume. + RolledBack, + /// Hard breach with no rollback budget or target — refuse (fail closed). + FailClosed, +} + +/// Mutable bookkeeping for the checkpointed-rollback loop, threaded through both +/// the cadence guard check and the mandatory final guard check. +struct GuardState { + checkpoints: CheckpointManager, + rollback_count: u8, + banned: Vec, + /// The post-prefill safe baseline `(output_len, kv_len)` — the fail-closed + /// restore target when no checkpoint exists, so prompt prefill KV survives. + start_out: u32, + start_kv: u32, +} + impl InferenceSession { pub fn new(id: SessionId, config: SessionConfig, engine: E, permit: LoadPermit) -> Self { let mut s = Self { @@ -118,10 +145,38 @@ impl InferenceSession { Ok(()) } - /// Run the decode loop until EOS or `max_tokens`. Each step composes - /// collaborators in the invariant order: grammar mask → safety adjust → - /// sample → commit. + /// Run the decode loop until EOS or `max_tokens`, deriving the rollback + /// policy from the session's device tier and safety mode (ADR-005/ADR-012). pub fn generate(&mut self, ports: &Ports, max_tokens: u32) -> Result { + let policy = RollbackPolicy::for_device(self.config.device, self.config.safety); + self.generate_with_policy(ports, max_tokens, policy) + } + + /// The checkpointed-rollback safety control loop (ADR-012). + /// + /// Every step preserves the invariant order **grammar mask → safety adjust → + /// sample → commit** (ADR-005). When a [`ChunkGuard`](el_safety::ChunkGuard) + /// is wired and the policy enables guarding, the loop additionally captures a + /// checkpoint at each guard-verified-safe boundary, scores recent output + /// every `guard_every` tokens, and on a hard-threshold breach rolls the KV + /// cache *and* the committed output back to the last safe checkpoint — + /// banning the offending token through the grammar mask so the resumed + /// decode necessarily diverges. Rollbacks are bounded by `max_rollbacks`; on + /// exhaustion — or with no checkpoint (e.g. under memory pressure) — the loop + /// fails closed with a deterministic refusal (`StopReason::Stopped`). + /// + /// Termination (EOS or `max_tokens`) is gated behind a **mandatory final + /// guard check**: the loop scores the trailing chunk before honouring either + /// stop condition, so no completion is ever returned unscored — including one + /// shorter than `guard_every` or whose unsafe tail ends in EOS. A final check + /// coincident with a cadence boundary is idempotent (re-scoring identical + /// output yields the same verdict). + pub fn generate_with_policy( + &mut self, + ports: &Ports, + max_tokens: u32, + policy: RollbackPolicy, + ) -> Result { if self.phase != Phase::Decoding { return Err(EdgeError::InvalidPhase { expected: "Decoding", @@ -129,43 +184,125 @@ impl InferenceSession { }); } let eos = self.engine.eos_token(); - - let stop = loop { - if self.output.len() as u32 >= max_tokens { - break StopReason::MaxTokens; + let guarding = policy.guards() && ports.guard.is_some(); + + // Tier-aware degradation (ADR-003/ADR-012): without budget for + // checkpoints, run guard-only with no rollback capability. + let checkpoints = if guarding { + if self.config.memory_budget_bytes < MIN_CHECKPOINT_BUDGET_BYTES { + self.emit(DomainEvent::SafetyDisabled { + reason: DegradeReason::MemoryPressure, + }); + CheckpointManager::new(0) + } else { + CheckpointManager::new(policy.max_checkpoints) } + } else { + CheckpointManager::new(0) + }; + let mut state = GuardState { + checkpoints, + rollback_count: 0, + banned: Vec::new(), + // The post-prefill baseline: the safe prefix to restore to when no + // checkpoint exists (e.g. checkpointing disabled under memory + // pressure). Captured as (output, KV) so fail-closed never drops + // prompt prefill KV. + start_out: self.output.len() as u32, + start_kv: self.kv.len(), + }; + // Seed the safe prefix at generation start (an empty continuation is safe). + if state.checkpoints.enabled() { + state.checkpoints.push(Checkpoint { + output_len: state.start_out, + kv_len: state.start_kv, + }); + } - // 2. verify / next-token logits (1. drafting is off by default) - let logits = self.engine.next_logits(&self.output); - let vocab = logits.len(); - - // 3. grammar mask (BEFORE safety) - let mask = ports.grammar.mask(&self.output, vocab); - let allowed = mask.iter().filter(|b| **b).count() as u32; - self.emit(DomainEvent::TokenMaskApplied { allowed }); + let stop = loop { + // A *candidate* termination for this iteration: the token cap is + // reached (checked before generating), or — set below — the model + // emitted EOS. With a guard active, neither is honoured until the + // final chunk passes the mandatory guard check, so a short or + // EOS-terminated tail cannot bypass scoring. + let mut terminating: Option = None; - // 4. safety adjust (AFTER mask, BEFORE sampling) - let adj = ports.safety.adjust(&self.output); - if !adj.is_empty() { - self.emit(DomainEvent::LogitsSteered { - adjustment_norm_milli: adj.l1_norm_milli(), + if self.output.len() as u32 >= max_tokens { + terminating = Some(StopReason::MaxTokens); + } else { + // 2. next-token logits (drafting off by default). + let logits = self.engine.next_logits(&self.output); + let vocab = logits.len(); + + // 3. grammar mask (BEFORE safety). Rollback bans ride the mask so + // the resumed decode cannot re-pick the off-trajectory token. + let mut mask = ports.grammar.mask(&self.output, vocab); + for &t in &state.banned { + if let Some(slot) = mask.get_mut(t as usize) { + *slot = false; + } + } + let allowed = mask.iter().filter(|b| **b).count() as u32; + self.emit(DomainEvent::TokenMaskApplied { allowed }); + + // 4. safety adjust (AFTER mask, BEFORE sampling). + let adj = ports.safety.adjust(&self.output); + if !adj.is_empty() { + self.emit(DomainEvent::LogitsSteered { + adjustment_norm_milli: adj.l1_norm_milli(), + }); + } + + // 5. sample (greedy argmax over legal, steered logits). If grammar + // + rollback bans leave no legal token, fail closed rather than + // emit a masked/banned token. + let token = match pick(&logits, &mask, &adj) { + Some(t) => t, + None => { + self.emit(DomainEvent::GrammarViolationBlocked); + break StopReason::Stopped; + } + }; + self.emit(DomainEvent::TokenGenerated { sampled: false }); + + // 6. commit. + self.output.push(token); + self.kv.push(self.output.len() as u64); + self.step += 1; + self.emit(DomainEvent::TokenCommitted { + kv_len: self.kv.len(), }); - } - // 5. sample (greedy argmax over legal, steered logits) - let token = pick(&logits, &mask, &adj); - self.emit(DomainEvent::TokenGenerated { sampled: false }); + if token == eos { + terminating = Some(StopReason::Eos); + } + } - // 6. commit - self.output.push(token); - self.kv.push(self.output.len() as u64); - self.step += 1; - self.emit(DomainEvent::TokenCommitted { - kv_len: self.kv.len(), - }); + // ---- chunk guard + checkpointed rollback (ADR-012) ---- + // Score at each `guard_every` cadence boundary AND before any + // termination (the mandatory final check). This closes the bypass + // where EOS or the token cap returned a tail shorter than + // `guard_every` unscored. + if guarding { + let guard = ports + .guard + .as_deref() + .expect("guarding implies a guard is wired"); + let at_boundary = (self.output.len() as u32).is_multiple_of(policy.guard_every); + if terminating.is_some() || at_boundary { + match self.guard_chunk(guard, &policy, &mut state) { + // Fail closed: no checkpoint, or rollback budget spent. + GuardVerdict::FailClosed => break StopReason::Stopped, + // Rolled back: the candidate termination (if any) was + // undone with it, so resume decoding from the safe prefix. + GuardVerdict::RolledBack => continue, + GuardVerdict::Pass => {} + } + } + } - if token == eos { - break StopReason::Eos; + if let Some(reason) = terminating { + break reason; } }; @@ -177,6 +314,81 @@ impl InferenceSession { Ok(stop) } + /// Score the committed output and apply the ADR-012 rollback policy: + /// advance the safe checkpoint when verified safe, roll back (banning the + /// divergence token) on a hard breach within budget, or fail closed. + /// + /// Invoked both at `guard_every` cadence boundaries and as the mandatory + /// final check before termination, so no completion is returned unscored. + /// On [`GuardVerdict::RolledBack`]/[`GuardVerdict::FailClosed`] the output + /// **and** KV are truncated together to the safe prefix (or the post-prefill + /// baseline) so prompt prefill descriptors are never dropped (AC-5). A + /// rollback also restores the *engine's* internal state via + /// [`InferenceEngine::rollback`] — a stateful engine (real KV cache) that + /// kept the abandoned branch would otherwise serve logits from the unsafe + /// path and skip the replacement tokens. + fn guard_chunk( + &mut self, + guard: &dyn el_safety::ChunkGuard, + policy: &RollbackPolicy, + state: &mut GuardState, + ) -> GuardVerdict { + let score = guard.score(&self.output); + if score >= policy.hard_threshold { + self.emit(DomainEvent::SafetyViolationDetected { + score_milli: score.milli(), + threshold_milli: policy.hard_threshold.milli(), + }); + match state.checkpoints.last() { + Some(cp) if state.rollback_count < policy.max_rollbacks => { + // Restore the engine's internal state (real KV cache / + // position) to the checkpoint too. If it cannot, fail closed + // rather than resume decoding on an inconsistent cache. + if self.engine.rollback(cp.output_len).is_err() { + self.output.truncate(cp.output_len as usize); + self.kv.truncate(cp.kv_len); + return GuardVerdict::FailClosed; + } + // Ban the token that began the unsafe span → divergence. + if let Some(&bad) = self.output.get(cp.output_len as usize) { + state.banned.push(bad); + } + self.output.truncate(cp.output_len as usize); + self.kv.truncate(cp.kv_len); + state.rollback_count += 1; + self.emit(DomainEvent::ClaimBacktracked { + claim_index: cp.output_len, + }); + GuardVerdict::RolledBack + } + _ => { + let (safe_out, safe_kv) = state + .checkpoints + .last() + .map_or((state.start_out, state.start_kv), |c| { + (c.output_len, c.kv_len) + }); + self.output.truncate(safe_out as usize); + self.kv.truncate(safe_kv); + GuardVerdict::FailClosed + } + } + } else if score < policy.soft_threshold { + // Verified safe: advance the last-safe checkpoint, drop bans. + if state.checkpoints.enabled() { + state.checkpoints.push(Checkpoint { + output_len: self.output.len() as u32, + kv_len: self.kv.len(), + }); + } + state.banned.clear(); + GuardVerdict::Pass + } else { + // soft ≤ score < hard: tolerated but not checkpointed (still risky). + GuardVerdict::Pass + } + } + /// Clear KV/output for a fresh conversation (volatile memory only). pub fn reset(&mut self) { self.kv = KvRegion::new(); @@ -205,7 +417,9 @@ impl InferenceSession { /// Greedy pick over legal, safety-steered logits. Masked-out tokens are skipped /// entirely; the safety delta is added to surviving logits before argmax. -fn pick(logits: &[i32], mask: &[bool], adj: &LogitAdjustment) -> Token { +/// Returns `None` when no token is legal (every token masked out or banned), so +/// the caller fails closed instead of emitting a rejected token. +fn pick(logits: &[i32], mask: &[bool], adj: &LogitAdjustment) -> Option { let mut best: Option = None; let mut best_val = i32::MIN; for (i, &l) in logits.iter().enumerate() { @@ -218,7 +432,7 @@ fn pick(logits: &[i32], mask: &[bool], adj: &LogitAdjustment) -> Token { best = Some(i as Token); } } - best.unwrap_or(0) + best } #[cfg(test)] @@ -258,6 +472,9 @@ mod tests { fn eos_token(&self) -> Token { 9999 } + fn rollback(&mut self, _keep: u32) -> Result<()> { + Ok(()) // stateless + } } // Grammar masker that disallows specific token ids. @@ -304,6 +521,7 @@ mod tests { compressor: Box::new(crate::defaults::IdentityCompressor), grammar: Box::new(DisallowMasker(vec![0])), // grammar removes the top token safety: Box::new(LightweightFilter::new(vec![1])), // safety bans the next-best + guard: None, relay: None, }; s.load_prompt(&ports, &[1]).unwrap(); @@ -384,4 +602,408 @@ mod tests { )); assert!(matches!(evs[1].event, DomainEvent::ModelLoaded { .. })); } + + // ----- ADR-012 checkpointed-rollback control loop ----- + + use el_safety::{ChunkGuard, SafetyScore}; + + /// Hard-unsafe whenever the given token appears in the output. + struct BanToken(Token); + impl ChunkGuard for BanToken { + fn score(&self, recent: &[Token]) -> SafetyScore { + if recent.contains(&self.0) { + SafetyScore::MAX + } else { + SafetyScore::SAFE + } + } + } + + /// Always hard-unsafe — exercises the rollback bound and fail-closed path. + struct AlwaysHot; + impl ChunkGuard for AlwaysHot { + fn score(&self, _recent: &[Token]) -> SafetyScore { + SafetyScore::MAX + } + } + + fn tiny_policy(max_rollbacks: u8) -> RollbackPolicy { + RollbackPolicy { + guard_every: 1, + soft_threshold: SafetyScore::from_milli(500), + hard_threshold: SafetyScore::from_milli(800), + max_rollbacks, + max_checkpoints: 8, + } + } + + /// Grammar masker that disallows every token — exercises the no-legal-token + /// fail-closed path. + struct DenyAllMasker; + impl GrammarMasker for DenyAllMasker { + fn mask(&self, _recent: &[Token], vocab: usize) -> Vec { + vec![false; vocab] + } + } + + /// Engine that always prefers token 0, then 1, 2, 3 — so banning the + /// chosen token forces the next-best, giving deterministic divergence. + fn descending_engine() -> FixedEngine { + FixedEngine { + logits: vec![5, 4, 3, 2], + } + } + + /// Emits the unsafe token `0` first, then EOS — a completion shorter than a + /// large `guard_every`, so only the *final* mandatory guard check can catch + /// it. Banning token `0` forces the next-best (a safe token), then EOS. + struct UnsafeThenEos { + eos: Token, + vocab: usize, + } + impl InferenceEngine for UnsafeThenEos { + fn prefill(&mut self, t: &[Token]) -> Result { + Ok(t.len() as u32) + } + fn next_logits(&mut self, ctx: &[Token]) -> Vec { + let mut v = vec![0i32; self.vocab]; + if ctx.is_empty() { + v[0] = 10; // unsafe token 0 wins the first step + } else { + v[self.eos as usize] = 10; // then terminate with EOS + } + v + } + fn eos_token(&self) -> Token { + self.eos + } + fn rollback(&mut self, _keep: u32) -> Result<()> { + Ok(()) // stateless: logits depend only on the passed ctx + } + } + + /// `guard_every` larger than any completion here, so the *cadence* check + /// never fires — isolating the mandatory final guard check. + fn coarse_policy(max_rollbacks: u8) -> RollbackPolicy { + RollbackPolicy { + guard_every: 16, + soft_threshold: SafetyScore::from_milli(500), + hard_threshold: SafetyScore::from_milli(800), + max_rollbacks, + max_checkpoints: 8, + } + } + + #[test] + fn eos_terminated_short_completion_is_scored_not_bypassed() { + // Regression (P1): EOS was handled before guard evaluation, so an unsafe + // tail ending in EOS within < guard_every tokens escaped scoring. + let mut s = InferenceSession::new( + SessionId(26), + SessionConfig::default(), + UnsafeThenEos { eos: 5, vocab: 8 }, + permit(), + ); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + // No rollback budget → the final check must refuse, not return EOS. + let stop = s.generate_with_policy(&ports, 8, coarse_policy(0)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + assert!(s.output().is_empty()); // unsafe EOS-terminated tail not emitted + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))); + } + + #[test] + fn max_tokens_partial_chunk_is_scored_not_bypassed() { + // Regression (P1): hitting max_tokens exited before flushing a partial + // chunk, so a completion shorter than guard_every was never scored. + let mut s = InferenceSession::new( + SessionId(27), + SessionConfig::default(), + descending_engine(), // always prefers unsafe token 0 + permit(), + ); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + // 2 tokens < guard_every (16): only the final check can catch the breach. + let stop = s.generate_with_policy(&ports, 2, coarse_policy(0)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + assert!(s.output().is_empty()); + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))); + } + + #[test] + fn eos_unsafe_tail_rolls_back_and_recovers() { + // With rollback budget, an unsafe EOS-terminated tail is rolled back to + // the seeded safe prefix, the offending token banned, and decoding + // resumes to a safe completion that then terminates cleanly. + let mut s = InferenceSession::new( + SessionId(28), + SessionConfig::default(), + UnsafeThenEos { eos: 5, vocab: 8 }, + permit(), + ); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + let stop = s.generate_with_policy(&ports, 8, coarse_policy(1)).unwrap(); + + assert_eq!(stop, StopReason::Eos); + assert!(!s.output().contains(&0)); // unsafe token banned out of the result + assert_eq!(s.output().last(), Some(&5)); // ends on EOS, scored safe + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. }))); + } + + /// Mirrors `QwenEngine`'s statefulness: tracks how many committed tokens it + /// has "fed" into its (mock) KV cache. If the session truncated its own + /// output without telling the engine, `committed` would shrink below `fed` — + /// the desync the real engine hits as serving stale logits and never + /// re-feeding. Shared cells let the test observe behaviour after the engine + /// is moved into the session. + struct StatefulEngine { + fed: usize, + logits: Vec, + eos: Token, + rollbacks: std::rc::Rc>, + last_keep: std::rc::Rc>, + desynced: std::rc::Rc>, + } + impl InferenceEngine for StatefulEngine { + fn prefill(&mut self, t: &[Token]) -> Result { + self.fed = 0; + Ok(t.len() as u32) + } + fn next_logits(&mut self, committed: &[Token]) -> Vec { + // The engine should never be "ahead" of the committed context; if it + // is, a rollback was not propagated here (the bug). + if self.fed > committed.len() { + self.desynced.set(true); + } + while self.fed < committed.len() { + self.fed += 1; + } + self.logits.clone() + } + fn eos_token(&self) -> Token { + self.eos + } + fn rollback(&mut self, keep_committed: u32) -> Result<()> { + self.fed = keep_committed as usize; // re-sync to the retained prefix + self.rollbacks.set(self.rollbacks.get() + 1); + self.last_keep.set(keep_committed); + Ok(()) + } + } + + #[test] + fn rollback_restores_engine_state_not_just_session_metadata() { + // Regression (P1): the loop truncated only the session's output + KV + // descriptors; a stateful engine kept the abandoned branch. The session + // must drive `InferenceEngine::rollback` on every backtrack. + let rollbacks = std::rc::Rc::new(std::cell::Cell::new(0u32)); + let last_keep = std::rc::Rc::new(std::cell::Cell::new(u32::MAX)); + let desynced = std::rc::Rc::new(std::cell::Cell::new(false)); + let engine = StatefulEngine { + fed: 0, + logits: vec![5, 4, 3, 2], // prefers the unsafe token 0 + eos: 9999, + rollbacks: rollbacks.clone(), + last_keep: last_keep.clone(), + desynced: desynced.clone(), + }; + let mut s = + InferenceSession::new(SessionId(29), SessionConfig::default(), engine, permit()); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[7, 8]).unwrap(); // non-empty prompt + + let stop = s.generate_with_policy(&ports, 3, tiny_policy(3)).unwrap(); + + assert_eq!(stop, StopReason::MaxTokens); + assert_eq!(s.output(), &[1, 1, 1]); // recovered safe completion + // The engine was told to roll back — not just the session metadata... + assert!( + rollbacks.get() >= 1, + "session must propagate the backtrack to the engine" + ); + // ...to a real prefix, and the cache never desynced from `committed`. + assert!(last_keep.get() < 3); + assert!( + !desynced.get(), + "engine cache must track the session rollback" + ); + } + + fn guarded_ports(guard: Box) -> Ports { + Ports { + compressor: Box::new(crate::defaults::IdentityCompressor), + grammar: Box::new(crate::defaults::AllowAllMasker), + safety: Box::new(el_safety::NoSafety), + guard: Some(guard), + relay: None, + } + } + + #[test] + fn hard_breach_rolls_back_kv_and_recovers() { + let mut s = InferenceSession::new( + SessionId(20), + SessionConfig::default(), + descending_engine(), + permit(), + ); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + let stop = s.generate_with_policy(&ports, 3, tiny_policy(3)).unwrap(); + + assert_eq!(stop, StopReason::MaxTokens); + // Token 0 is unsafe; each occurrence is rolled back and banned, so the + // recovered output contains only the safe next-best token. + assert_eq!(s.output(), &[1, 1, 1]); + assert!(!s.output().contains(&0)); + // KV rewound in lock-step with the committed output (AC-5). + assert_eq!(s.kv_len(), s.output().len() as u32); + + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. }))); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))); + } + + #[test] + fn fail_closed_refusal_when_no_rollback_budget() { + let mut s = InferenceSession::new( + SessionId(21), + SessionConfig::default(), + descending_engine(), + permit(), + ); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + // No rollback budget → first hard breach refuses deterministically. + let stop = s.generate_with_policy(&ports, 5, tiny_policy(0)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + assert!(s.output().is_empty()); + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::SafetyViolationDetected { .. }))); + assert!(!evs + .iter() + .any(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. }))); + } + + #[test] + fn rollbacks_are_bounded_then_fail_closed() { + let mut s = InferenceSession::new( + SessionId(22), + SessionConfig::default(), + descending_engine(), + permit(), + ); + let ports = guarded_ports(Box::new(AlwaysHot)); + s.load_prompt(&ports, &[]).unwrap(); + + let stop = s.generate_with_policy(&ports, 8, tiny_policy(2)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + let evs = s.drain_events(); + let rollbacks = evs + .iter() + .filter(|e| matches!(e.event, DomainEvent::ClaimBacktracked { .. })) + .count(); + // Exactly max_rollbacks attempts, then a deterministic refusal (AC-6). + assert_eq!(rollbacks, 2); + } + + #[test] + fn memory_pressure_disables_checkpoints_and_fails_closed() { + let cfg = SessionConfig { + memory_budget_bytes: 1024, // far below MIN_CHECKPOINT_BUDGET_BYTES + ..SessionConfig::default() + }; + let mut s = InferenceSession::new(SessionId(23), cfg, descending_engine(), permit()); + let ports = guarded_ports(Box::new(BanToken(0))); + s.load_prompt(&ports, &[]).unwrap(); + + let stop = s.generate_with_policy(&ports, 5, tiny_policy(3)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); // no checkpoint to roll back to + let evs = s.drain_events(); + assert!(evs.iter().any(|e| matches!( + e.event, + DomainEvent::SafetyDisabled { + reason: DegradeReason::MemoryPressure + } + ))); + } + + #[test] + fn no_legal_token_fails_closed() { + // Grammar disallows everything → pick() has no legal token. The loop must + // fail closed, not commit token 0 (which could be EOS). + let mut s = InferenceSession::new( + SessionId(24), + SessionConfig::default(), + descending_engine(), + permit(), + ); + let ports = Ports { + compressor: Box::new(crate::defaults::IdentityCompressor), + grammar: Box::new(DenyAllMasker), + safety: Box::new(el_safety::NoSafety), + guard: None, + relay: None, + }; + s.load_prompt(&ports, &[]).unwrap(); + + let stop = s.generate(&ports, 4).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + assert!(s.output().is_empty()); // no illegal token committed + let evs = s.drain_events(); + assert!(evs + .iter() + .any(|e| matches!(e.event, DomainEvent::GrammarViolationBlocked))); + } + + #[test] + fn fail_closed_preserves_prefill_kv() { + // Non-empty prompt → prefill KV descriptors must survive a + // budget-exhausted refusal (regression: fail-closed once truncated KV to + // output_len, dropping prefill). + let mut s = InferenceSession::new( + SessionId(25), + SessionConfig::default(), + descending_engine(), + permit(), + ); + let ports = guarded_ports(Box::new(AlwaysHot)); + s.load_prompt(&ports, &[7, 8, 9]).unwrap(); // prefill KV length = 3 + assert_eq!(s.kv_len(), 3); + + let stop = s.generate_with_policy(&ports, 8, tiny_policy(1)).unwrap(); + + assert_eq!(stop, StopReason::Stopped); + assert!(s.output().is_empty()); // refused back to the post-prefill prefix + assert_eq!(s.kv_len(), 3); // prefill KV intact, not truncated to 0 + } } diff --git a/crates/el-safety/README.md b/crates/el-safety/README.md index 8790b7e..8b572fc 100644 --- a/crates/el-safety/README.md +++ b/crates/el-safety/README.md @@ -27,6 +27,15 @@ Depends only on `el-core`. No `unsafe` (`#![forbid(unsafe_code)]`). **Scaffolded** follow-up: it requires two ~1B models on Candle, so until the assets are wired it returns no adjustment while honestly reporting its mode (so callers can select it without it silently mis-steering). +- **ADR-012 control-loop primitives** (consumed by `el-runtime`'s decode loop): + - `ChunkGuard::score(recent) -> SafetyScore` — risk of the recent output window. + - `SafetyScore` — integer milli-units `0..=1000` (deterministic, float-free). + - `RollbackPolicy::for_device(device, mode)` — tier-aware bounds (guard cadence + `guard_every`, soft/hard thresholds, `max_rollbacks`, checkpoint-ring size). + The early-token soft-steering window ships with the SecDecoding steerer + (follow-up); hard bans apply every step. + - `CheckpointManager` / `Checkpoint` — a bounded ring of guard-verified + safe-prefix snapshots (offsets only; KV payload is never copied). ## Usage @@ -47,18 +56,28 @@ assert_eq!(adj.delta_for(7), 0); ## Place in the workspace -Re-exported by `el-runtime` (`el_runtime::{SafetySteerer, LogitAdjustment}`) so +Re-exported by `el-runtime` (`el_runtime::{SafetySteerer, LogitAdjustment, +ChunkGuard, SafetyScore, RollbackPolicy, Checkpoint, CheckpointManager}`) so callers wire a single type system. The session applies the chosen steerer in the -invariant decode order `grammar mask → safety adjust → sample → commit`. +invariant decode order `grammar mask → safety adjust → sample → commit`, wrapped +by the ADR-012 checkpointed-rollback control loop. ## Status Partial by design: the `Lightweight` blacklist path is real and tested; `SecDecoding`/`Csd` model-backed steering is a tracked follow-up that needs -model assets. +model assets. The ADR-012 checkpointed-rollback control-loop primitives +(`ChunkGuard`, `RollbackPolicy`, `CheckpointManager`) are implemented and tested, +and the `el-runtime` session drives them (hard bans every step, chunk-guard +cadence with a mandatory final check, safe-prefix checkpoints, bounded +fail-closed rollback). Selective soft-steering over an early-token window is part +of the deferred `SecDecoding` follow-up above, not current behavior. A trained +`ChunkGuard` model is the remaining follow-up — the loop runs against any +`ChunkGuard` implementation. --- Part of the [Edge Intelligence](../../README.md) workspace. Realizes -[ADR-005](../../docs/adr/ADR-005-on-device-only-tiered-decoder-time-safety.md); +[ADR-005](../../docs/adr/ADR-005-on-device-only-tiered-decoder-time-safety.md) and +[ADR-012](../../docs/adr/ADR-012-layered-decode-time-safety-control-loop-with-checkpointed-rollback.md); see the [Safety](../../docs/ddd/bounded-contexts/05-safety.md) context. diff --git a/crates/el-safety/src/lib.rs b/crates/el-safety/src/lib.rs index 756652e..ca5789b 100644 --- a/crates/el-safety/src/lib.rs +++ b/crates/el-safety/src/lib.rs @@ -5,6 +5,12 @@ //! implemented here. `SecDecoding` (two ~1B models) and `Csd` (claim //! backtracking) require model assets and are scaffolded as follow-ups //! ([`SecDecodingSteerer`]). **No safety path touches the network.** +//! +//! ADR-012 adds the runtime-backtracking primitives consumed by the Inference +//! Runtime's decode-time control loop: [`ChunkGuard`]/[`SafetyScore`] scoring, +//! the tier-aware [`RollbackPolicy`] (cadence + bounds), and +//! [`CheckpointManager`]/[`Checkpoint`] safe-prefix snapshots (offsets only — +//! KV payload is never copied). #![forbid(unsafe_code)] @@ -128,6 +134,170 @@ impl SafetySteerer for SecDecodingSteerer { } } +// --------------------------------------------------------------------------- +// ADR-012 — checkpointed-rollback control-loop primitives +// --------------------------------------------------------------------------- + +/// Risk score in milli-units, `0` (safe) ..= `1000` (max). Integer for +/// deterministic, float-free safety decisions (ADR-008). +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Default)] +pub struct SafetyScore(u16); + +impl SafetyScore { + pub const SAFE: SafetyScore = SafetyScore(0); + pub const MAX: SafetyScore = SafetyScore(1000); + + /// Clamp a milli-unit value into `[0, 1000]`. + pub fn from_milli(milli: u16) -> Self { + Self(milli.min(1000)) + } + + /// The score in milli-units, as reported on `SafetyViolationDetected`. + pub fn milli(self) -> u16 { + self.0 + } +} + +/// Scores recent generated output for risk (ADR-012 chunk guard). Reuses the +/// active tier's safety model; like every safety path it is deterministic and +/// **never touches the network** (ADR-004). +pub trait ChunkGuard { + /// Risk of the recent output window. Higher is riskier. + fn score(&self, recent: &[Token]) -> SafetyScore; +} + +/// Cadence and bounds for the checkpointed-rollback control loop (ADR-012), +/// chosen per device tier so cost scales with the hardware budget (ADR-003). +/// +/// The early-token *soft-steering window* is intentionally absent: hard bans +/// apply every step, and the windowed SecDecoding-style steering arrives with +/// its model (follow-up). Checkpoint spacing is `guard_every` — checkpoints are +/// only taken at guard-verified-safe boundaries, so there is no separate +/// checkpoint cadence to misconfigure. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RollbackPolicy { + /// Score the output every `guard_every` tokens (`0` = never); also the + /// spacing of safe-prefix checkpoints. + pub guard_every: u32, + /// At/above this score, escalate; do not advance the safe checkpoint. + pub soft_threshold: SafetyScore, + /// At/above this score, roll back (or fail closed). + pub hard_threshold: SafetyScore, + /// Hard cap on rollbacks before a deterministic refusal (DoS bound). + pub max_rollbacks: u8, + /// Bounded checkpoint ring size — fixed memory (ADR-003). + pub max_checkpoints: u8, +} + +impl RollbackPolicy { + /// Tier-aware policy. `SafetyMode::Off` disables the loop entirely. + pub fn for_device(device: DeviceTarget, mode: SafetyMode) -> Self { + if matches!(mode, SafetyMode::Off) { + return Self::disabled(); + } + match device { + DeviceTarget::MidRange => Self { + guard_every: 16, + soft_threshold: SafetyScore(600), + hard_threshold: SafetyScore(800), + max_rollbacks: 2, + max_checkpoints: 4, + }, + DeviceTarget::HighEnd | DeviceTarget::Auto => Self { + guard_every: 4, + soft_threshold: SafetyScore(500), + hard_threshold: SafetyScore(750), + max_rollbacks: 4, + max_checkpoints: 8, + }, + } + } + + /// A policy that performs no steering, checkpointing, or guarding. + pub fn disabled() -> Self { + Self { + guard_every: 0, + soft_threshold: SafetyScore::MAX, + hard_threshold: SafetyScore::MAX, + max_rollbacks: 0, + max_checkpoints: 0, + } + } + + /// Whether the guard/rollback machinery is active under this policy. + pub fn guards(&self) -> bool { + self.guard_every > 0 + } +} + +/// A safe-prefix snapshot for rollback (ADR-012). Stores only indices: rollback +/// truncates KV descriptors (`KvRegion::truncate`) and never replays prefill, +/// and the KV payload is never copied (ADR-002/ADR-003). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Checkpoint { + /// Committed-token count at the checkpoint. + pub output_len: u32, + /// KV-cache length to restore to. + pub kv_len: u32, +} + +/// A bounded ring of guard-verified safe-prefix checkpoints. Fixed memory: the +/// oldest checkpoint is dropped once `cap` is reached (ADR-003). A `cap` of `0` +/// (or [`disable`](Self::disable)) retains nothing — the loop then has no +/// rollback target and fails closed on a breach. +#[derive(Debug, Default)] +pub struct CheckpointManager { + ring: Vec, + cap: usize, + enabled: bool, +} + +impl CheckpointManager { + pub fn new(cap: u8) -> Self { + Self { + ring: Vec::new(), + cap: cap as usize, + enabled: cap > 0, + } + } + + pub fn enabled(&self) -> bool { + self.enabled + } + + /// Drop all checkpoints and stop retaining new ones (memory-pressure + /// degradation — ADR-012/ADR-003). + pub fn disable(&mut self) { + self.enabled = false; + self.ring.clear(); + } + + /// Record a safe prefix; evicts the oldest if the ring is full. No-op when + /// disabled. + pub fn push(&mut self, checkpoint: Checkpoint) { + if !self.enabled { + return; + } + if self.ring.len() == self.cap { + self.ring.remove(0); + } + self.ring.push(checkpoint); + } + + /// The most recent safe prefix, if any. + pub fn last(&self) -> Option { + self.ring.last().copied() + } + + pub fn len(&self) -> usize { + self.ring.len() + } + + pub fn is_empty(&self) -> bool { + self.ring.is_empty() + } +} + #[cfg(test)] mod tests { use super::*; @@ -152,4 +322,60 @@ mod tests { assert_eq!(adj.delta_for(7), 0); assert!(adj.l1_norm_milli() > 0); } + + #[test] + fn safety_score_clamps_and_orders() { + assert_eq!(SafetyScore::from_milli(5000), SafetyScore::MAX); + assert!(SafetyScore::SAFE < SafetyScore::MAX); + assert_eq!(SafetyScore::from_milli(750).milli(), 750); + } + + #[test] + fn policy_off_is_disabled() { + let p = RollbackPolicy::for_device(DeviceTarget::HighEnd, SafetyMode::Off); + assert!(!p.guards()); + assert_eq!(p.max_rollbacks, 0); + } + + #[test] + fn policy_is_tier_aware() { + let mid = RollbackPolicy::for_device(DeviceTarget::MidRange, SafetyMode::Lightweight); + let high = RollbackPolicy::for_device(DeviceTarget::HighEnd, SafetyMode::SecDecoding); + assert!(mid.guards() && high.guards()); + // Stronger hardware guards more often and tolerates more rollbacks. + assert!(high.guard_every < mid.guard_every); + assert!(high.max_rollbacks >= mid.max_rollbacks); + } + + #[test] + fn checkpoint_ring_is_bounded_and_disablable() { + let mut m = CheckpointManager::new(2); + m.push(Checkpoint { + output_len: 1, + kv_len: 1, + }); + m.push(Checkpoint { + output_len: 2, + kv_len: 2, + }); + m.push(Checkpoint { + output_len: 3, + kv_len: 3, + }); // evicts the oldest + assert_eq!(m.len(), 2); + assert_eq!( + m.last(), + Some(Checkpoint { + output_len: 3, + kv_len: 3, + }) + ); + m.disable(); + assert!(!m.enabled() && m.is_empty()); + m.push(Checkpoint { + output_len: 9, + kv_len: 9, + }); // no-op when disabled + assert!(m.last().is_none()); + } } diff --git a/docs/adr/ADR-012-layered-decode-time-safety-control-loop-with-checkpointed-rollback.md b/docs/adr/ADR-012-layered-decode-time-safety-control-loop-with-checkpointed-rollback.md new file mode 100644 index 0000000..416cc24 --- /dev/null +++ b/docs/adr/ADR-012-layered-decode-time-safety-control-loop-with-checkpointed-rollback.md @@ -0,0 +1,141 @@ +# ADR-012: Layered decode-time safety control loop with checkpointed rollback + +- **Status**: accepted +- **Date**: 2026-06-16 +- **Deciders**: +- **Tags**: safety, security, on-device, runtime, supporting + +## Context + +[ADR-005](./ADR-005-on-device-only-tiered-decoder-time-safety.md) decided *which* +safety mode runs (`Off | Lightweight | SecDecoding | CSD`) and *where* the +`LogitAdjustment` sits in a decode step (after the grammar mask, before +sampling). It did not decide *how* steering recovers once generation has already +drifted into an unsafe trajectory. Treating safety as a one-shot per-token gate +is the weak point: once an autoregressive model commits to a bad prefix, +post-hoc repair is expensive and often too late. + +The `SecDecoding alternatives` research +([docs/research/SecDecoding alternatives research.md](../research/SecDecoding%20alternatives%20research.md)) +surveys the field and concludes that the strongest practical edge architecture is +*not* any single technique but a **hybrid control loop**: selective +model-backed steering (SafeDecoding-style, concentrated on the early-token +window) wrapped in **checkpointed rollback to the last safe prefix**, with a +chunk-level guard and deterministic ingress rules, degrading gracefully across +device tiers. Runtime backtracking ([RESET]-token, RoCode, Hydra) is identified +as the missing piece that turns steering into a recoverable loop. The research +also warns that a single safety judge is itself an attack surface, so the monitor +must be heterogeneous. + +## Decision + +Wrap the ADR-005 steering in a **recoverable decode-time control loop** with four +on-device, air-gapped stages (no stage may reach the network, per +[ADR-004](./ADR-004-air-gapped-by-default-with-opt-in-hybrid-mode.md)): + +1. **Selective steering window.** Apply model-backed `LogitAdjustment` on prompt + ingress and the **first 8–32 output tokens**, then fall back to normal decode + unless the guard re-escalates. This captures most early-trajectory risk at the + lowest sidecar cost. +2. **Checkpoint the KV-cache, not just text.** Snapshot generation state + (emitted token IDs, sampler RNG state, safety-control state, and KV-cache + handle / block references) every **K = 8–16 tokens** and at semantic + boundaries (newline, sentence break, closing brace, tool-call delimiter). + Rollback **restores cache state** rather than replaying text *where the engine + supports in-place cache restoration*; the realized per-engine cost (some + engines must replay) is recorded under Consequences. +3. **Chunk-guard cadence.** Score recent output every **4–16 tokens** (not every + token); where hardware allows, run the guard **asynchronously** with + generation (Hydra-style) so valid output pays near-zero cost. +4. **Bounded rollback, fail-closed.** On a hard-threshold breach, restore the last + safe checkpoint, raise `alpha` / lower temperature / inject refusal bias, and + retry. After `max_rollbacks` (or with no checkpoint), emit a **deterministic + hard refusal**. Rollback loops are always bounded. + +The monitor stays **heterogeneous**: the model-backed signal is one layer +alongside deterministic ingress controls and CSD/grammar constraints on any +machine-consumable channel — never a single judge. Cadence, checkpoint density, +and steering width are **tier-aware**, selected from the `DeviceProfile` and +subject to the [ADR-003](./ADR-003-static-memory-planning-with-zero-allocation-arena.md) +memory budget. Safety adapters, guard weights, and policy files are integrity-gated +on load per [ADR-006](./ADR-006-mandatory-ed25519-model-signature-verification-load-gate.md). + +### Model inventory + +One generator is mandatory; the safety signal's form follows the active +`SafetyMode`. The CSD layer is a grammar engine, not a model. + +| Mode | On-device model artifacts | +|---|---| +| `Off` | Main generator only | +| `Lightweight` | Main **+ one of**: same-family safety **LoRA adapter** (shares base weights), a ~0.1B refusal classifier, or token-anchor heuristics (no weights) | +| `SecDecoding` | Main **+** small **base** **+** safety-tuned **expert** (contrastive pair) — accelerator-class HW only | +| `CSD` | Main **+** grammar/completion engine (**llguidance** — code, not weights) | + +Defaults and invariants: +- **Recommended baseline:** a quantized `Qwen2.5-0.5B-Instruct` / `Qwen3-0.6B` + generator (mixed-INT4 on GPU/NPU, INT8 on CPU) **plus one safety LoRA adapter** + — two artifacts cover the `Lightweight` default. The contrastive second model is + shipped only with accelerator headroom. +- **Single shared tokenizer.** Steering intersects base/expert token sets, so the + adapter / expert / classifier must use the **identical** tokenizer (the reason + for a same-family Qwen choice). +- The chunk guard and prompt-risk triage **reuse** the active tier's safety model; + the backtracking loop adds **no** new weights. +- Every weight file (base, adapter, expert, classifier) is ED25519-signed and + integrity-gated on load (ADR-006). + +## Consequences + +### Positive +- Safety becomes recoverable, not one-shot: drift is caught mid-generation and + rewound to a safe prefix instead of repaired after the fact. +- Worst-case behaviour is bounded and deterministic (capped rollbacks → hard + refusal), which suits an air-gapped device with no field re-tuning. +- Defence-in-depth: no single guard model is the whole safety case. + +### Negative +- Rollback and replay add **latency variance**; pathological inputs can trigger + repeated rewinds (a local DoS vector that the rollback cap must contain). On + engines without in-place cache truncation each rewind replays prompt + kept + prefix (see *Realized* below), so a single capped rollback can cost up to a full + prefill — `max_rollbacks` is what keeps the worst case finite. +- Checkpoints cost memory (KV-cache handles/block refs + sampler/control state), + competing with the ADR-003 budget on low-end devices. +- Loop, checkpoint, and async-guard state machine is materially more complex than + a stateless per-token adjustment. + +### Neutral +- Under memory pressure the loop degrades gracefully: checkpoint density and guard + cadence thin out (and ultimately `SafetyDisabled`) per the ADR-003 policy. +- **As built:** hard bans apply every step; the early-token *soft-steering window* + ships with the SecDecoding steerer (follow-up), and checkpoint spacing is the + guard cadence (`guard_every`) — checkpoints are taken only at guard-verified-safe + boundaries, so there is no separate checkpoint-cadence knob to misconfigure. +- Feasibility of cheap cache-handle checkpoints depends on the inference engine + exposing KV-cache references; on [ADR-002](./ADR-002-candle-as-rust-native-inference-engine.md) + (Candle) this constrains how checkpoints are represented. **Realized — in two + layers, with different cost models:** + - *Session layer.* The descriptor-only `el-memory::KvRegion` (ADR-003) makes the + session's rollback an `O(dropped)` `truncate` of KV descriptors — no payload + copy — so the token-buffer fallback was not needed. + - *Engine layer.* Restoring the engine's own KV cache is delegated to + `InferenceEngine::rollback`, whose cost is adapter-dependent. A stateless + engine restores in `O(1)`. But an engine whose cache is **append-only with no + truncation API** cannot restore in place — candle 0.8.4's `quantized_qwen2` + holds a *private* per-layer cache and exposes no `clear`/`truncate`, so + `QwenEngine` rebuilds the safe prefix by **replaying prompt + kept tokens** + (`O(prompt + kept_prefix)` per rollback). For that engine the "restores cache + state, not a full replay" goal does **not** hold; the replay-free model is + contingent on an engine that exposes true cache truncation. The replay is + bounded by `max_rollbacks` (each capped rewind costs up to a prefill — see + *Negative*), and a future engine with a truncation API can override + `rollback` to recover the `O(dropped)` behaviour without any change to the + control loop. + +## Links +- Implementation: `crates/el-safety` (`ChunkGuard`, `SafetyScore`, `RollbackPolicy`, `CheckpointManager`/`Checkpoint`), `crates/el-runtime` (`InferenceSession::generate_with_policy` control loop), `crates/el-memory` (`KvRegion::truncate`) +- Research: [docs/research/SecDecoding alternatives research.md](../research/SecDecoding%20alternatives%20research.md) +- PRD: `docs/prd.md` §"Safety Guardrails (SecDecoding, CSD)", §Decoding Loop → Safety Adjustment +- DDD: [Safety context](../ddd/bounded-contexts/05-safety.md), [domain-events](../ddd/domain-events.md) +- Related: [ADR-005](./ADR-005-on-device-only-tiered-decoder-time-safety.md), [ADR-004](./ADR-004-air-gapped-by-default-with-opt-in-hybrid-mode.md), [ADR-006](./ADR-006-mandatory-ed25519-model-signature-verification-load-gate.md), [ADR-003](./ADR-003-static-memory-planning-with-zero-allocation-arena.md), [ADR-002](./ADR-002-candle-as-rust-native-inference-engine.md) diff --git a/docs/ddd/bounded-contexts/05-safety.md b/docs/ddd/bounded-contexts/05-safety.md index ed3e5cd..c92fc93 100644 --- a/docs/ddd/bounded-contexts/05-safety.md +++ b/docs/ddd/bounded-contexts/05-safety.md @@ -94,3 +94,26 @@ a `LogitAdjustment` applied **after** the grammar mask and **before** sampling; in CSD mode the orchestrator honors `Backtrack` signals at claim boundaries. Reads `DeviceProfile` from context 7. Independent of Grammar (4) and Speculative Decoding (3). + +## Runtime control loop (ADR-012) + +The decoder-time steering above runs inside a **recoverable control loop** (not a +one-shot gate). Implemented across `el-safety` (`ChunkGuard`, `SafetyScore`, +`RollbackPolicy`, `CheckpointManager`/`Checkpoint`) and the `el-runtime` session +(`generate_with_policy`): + +- **Hard bans every step** (the `Lightweight` blacklist applies on every token). + *Selective soft-steering* over an early-token window, with guard-gated + re-escalation, is **deferred** to the `SecDecoding` follow-up: it needs the + base+expert model pair, so `SecDecodingSteerer` is a placeholder today and the + early-token window is intentionally absent (see `el-safety` `RollbackPolicy`). +- **Safe-prefix checkpoints** captured at guard-verified boundaries — offsets + only; KV payload is never copied (`KvRegion::truncate`, ADR-003). +- **Chunk guard** scores recent output every few tokens; on a hard-threshold + breach the loop rolls KV **and** output back to the last safe checkpoint and + bans the offending token so the resume diverges (`ClaimBacktracked`). +- **Bounded, fail-closed:** rollbacks are capped (`max_rollbacks`); on exhaustion + — or under memory pressure with no checkpoint — the loop refuses + deterministically and emits `SafetyViolationDetected` / `SafetyDisabled`. + +See [ADR-012](../../adr/ADR-012-layered-decode-time-safety-control-loop-with-checkpointed-rollback.md). diff --git a/docs/ddd/domain-events.md b/docs/ddd/domain-events.md index 9677b1b..5ac5caa 100644 --- a/docs/ddd/domain-events.md +++ b/docs/ddd/domain-events.md @@ -63,7 +63,7 @@ | `SafetyModeSelected` | mode, downgradedFrom? | affordable mode chosen | | `LogitsSteered` | stepIndex, adjustmentNorm | SecDecoding adjustment applied | | `SafetyViolationDetected` | stepIndex, score, threshold | content exceeded threshold | -| `ClaimBacktracked` | claimIndex, score | CSD rewind + resample | +| `ClaimBacktracked` | claimIndex | CSD / ADR-012 rollback to last safe checkpoint | | `SafetyDisabled` | reason | downgraded for memory | ### 6. Memory Management (Supporting / Shared Kernel) diff --git a/docs/followups.md b/docs/followups.md new file mode 100644 index 0000000..91c8163 --- /dev/null +++ b/docs/followups.md @@ -0,0 +1,56 @@ + 1. Trained ChunkGuard model — do this first + + The trait already exists; today only test stubs implement it, so in production the loop never fires. This is the real safety gap and your memory flags the base model as crisis-unsafe, so + it's the priority. + + - Where: a new adapter (crates/adapters/el-guard-candle, or fold into el-engine-candle to reuse the tokenizer/runtime) implementing el_safety::ChunkGuard. Not in el-safety (keeps it + dep-free). + - Model: start with a small (~0.1B) distilled classifier over the recent token window (avoids detokenization, stays deterministic), INT8-quantized, emitting an integer → + SafetyScore::from_milli. The DDD already names this (LightweightFilter/safety-expert) and the SecDecodingAcl to keep vendor tensors out of the domain. + - Key change: the adapter's constructor takes a LoadPermit (mirror InferenceSession::new) so the guard weights pass the ED25519 gate — this turns AC-10 from "inherited" into "enforced for + the guard." + - Calibration: the soft/hard milli-thresholds in RollbackPolicy should be set empirically, not guessed. Wire a calibration pass through apps/el-bench (CounselBench/MindEval/VERA-MH) and + report ASR / over-refusal so the thresholds map to a real FPR/FNR. + - Heterogeneity (ADR-012): keep the deterministic blacklist + ingress rules alongside the trained guard — it's one signal, not the whole case. + - Seam choice: prefer the guard running its own small forward over recent (modular, seam unchanged). Only if latency forces it, extend InferenceEngine with last_hidden() and use a cheap + probe head (couples guard to engine internals — the optimization, not the start). + + 2. Soft-steering window gate — heaviest; do last + + Today safety.adjust runs every step (correct — LightweightFilter is a hard ban). The window (R1/AC-1) only matters for a soft steerer, which is still a placeholder. + + - Split the port into two channels: keep safety as the always-on hard constraint, and add a soft steerer (a real SecDecodingSteerer in an adapter running the base+expert contrastive pair, + or a safety LoRA) whose adjustment is applied only when step < steering_window || guard_mode. + - Wiring: add a guard_mode bit to the loop's control state, set it on the soft ≤ score < hard branch (which is currently a no-op) — that's what makes the window re-open on elevated risk. + Add LogitAdjustment::merge so the hard + soft adjustments combine before pick. + - ACL: the steerer holds its own small model and returns a sparse milli-logit delta from base/expert divergence — quantize at the SecDecodingAcl boundary so the core stays float-free. + - Sequence after #1 because both need the safety-expert asset and the Candle ACL. + + 3. Semantic-boundary checkpoints — light; high value + + The loop only sees token ids, so boundary detection is tokenizer-specific. + + - Add a port: trait BoundaryDetector { fn is_boundary(&self, t: Token) -> bool }, provided by the tokenizer-owning adapter; default = never (current behavior). Add boundary: Option> to Ports (same pattern as guard). + - Precompute the boundary token-id set once at load (scan vocab for ids whose surface contains \n, sentence terminators, }, tool-call delimiters) → O(1) lookup in the hot loop, + deterministic. + - Integrate: trigger a guard check at boundaries in addition to cadence, and checkpoint on the pass — so rollback targets land on clause/sentence starts. E3 (boundary+cadence coincide) is + already idempotent. + + 4. Async (Hydra-style) guard — perf optimization on top of #1 + + Only worth it once #1 exists and on accelerator-class hardware (Jetson/Core Ultra per ADR-012); on the default tier a sync guard every 4–16 tokens is fine. + + - Keep the core sync and dep-free: push threading into the adapter. ChunkGuard gains a non-blocking score_async(recent) -> Ticket; a worker thread holds the model; the loop try_recv()s + each step and continues decoding meanwhile. + - Add CheckpointManager::last_safe_at_or_before(index) (we only have last()); on a flagged breach roll back to the newest checkpoint with output_len ≤ scored_index (E4 is already specified + for this). + - Determinism guard-rail (the key constraint): make the rollback target depend only on the logical chunk index, never on real-time arrival, so output stays a deterministic function of + (input, model, policy) even though compute overlaps. Size the checkpoint ring to span the in-flight window. + + Suggested sequence & tracking + + #1 → #3 → #4 → #2. #1 is the safety substance; #3 is cheap precision; #4 optimizes #1; #2 is a separate, heavier capability. + + Each deserves its own increment under the same SPARC/GOAP tracking (and #1 and #2 each warrant a short ADR — new model assets + a provenance/threat surface). I'd reuse apps/el-bench as the + gate for #1's thresholds. \ No newline at end of file diff --git a/docs/research/Building Edge LLM SDK.md b/docs/research/Building Edge LLM SDK.md new file mode 100644 index 0000000..7766563 --- /dev/null +++ b/docs/research/Building Edge LLM SDK.md @@ -0,0 +1,301 @@ +**Prompt**: Research how to build pipeline/flow that utilize local llm models and optimize the prompt and response to achieve more accurate results. Find SOTA solutions and explore latest peer review academic work. + +The goal is to create sdk to use in an edge device like smartphone + +**Response**: +# **Technical Design of an Edge-Native LLM SDK Pipeline: Optimizing Execution Runtimes, Prompt Constraints, Speculative Verification, and Structured Generation on Smartphone Hardware** + +Deploying large language models (LLMs) directly on resource-constrained smartphone hardware represents a paradigm shift away from cloud-based machine learning serving. This shift is motivated by the demands for strict user data privacy, guaranteed offline operation, reduced latency, and eliminated server-side API operational expenses.1 To achieve these goals, an edge-native Software Development Kit (SDK) must resolve the severe physical limitations of smartphone platforms. + +## **Core Hardware Bottlenecks and SDK Optimization Requirements** + +Modern smartphone System-on-Chips (SoCs) possess highly capable Neural Processing Units (NPUs) and Graphics Processing Units (GPUs).4 For example, the Apple A19 Pro Neural Engine delivers approximately 35 Tera Operations Per Second (TOPS), while the Qualcomm Snapdragon 8 Elite Gen 5 offers up to 60 TOPS.5 However, these processing capabilities are severely throttled by mobile memory architectures.5 +While datacenter accelerators feature memory bandwidths ranging from 2 to 3 Terabytes per second (TB/s), modern smartphone SoCs operate within a constrained envelope of 50 to 90 Gigabytes per second (GB/s).5 This 30x to 50x bandwidth deficit creates a profound bottleneck during the autoregressive decoding phase of language model inference.5 During decoding, the entire weight tensor of the model must be loaded from physical memory to the processor core to generate a single token, leaving high-performance tensor cores idle while waiting for memory transfers.5 +An edge-native SDK must therefore optimize the entire computational pipeline, from prompt ingestion to final token output. This report details the design of a unified, high-performance edge execution pipeline, synthesizing state-of-the-art (SOTA) research across model compilation, dynamic context pruning, hardware-cooperative speculative decoding, grammar-constrained generation, and real-time safety enforcement. + +## **Co-Design and Performance Evaluation of Mobile Runtimes** + +The foundation of the edge SDK is the selection and configuration of a low-level execution engine. The engine must compile model graphs, manage memory, and map operators directly to the heterogeneous compute units of the device. + +| Runtime Framework | Model Support and Portability | Memory Allocation Strategy | Accelerator Delegation | Key Engineering Trade-offs | +| :---- | :---- | :---- | :---- | :---- | +| **ExecuTorch** | PyTorch-exclusive via torch.export; generates unified .pte binaries 4 | Static ahead-of-time (AOT) memory planning; zero-heap runtime 4 | Pluggable Delegates (Apple CoreML, Qualcomm QNN, Arm Vulkan/XNNPACK) 4 | Minimal runtime binary footprint (KB core); steep learning curve; strict export compliance 6 | +| **ONNX Runtime Mobile** | Universal format; accepts PyTorch, TensorFlow, and Scikit-Learn 6 | Dynamic allocation with execution-provider-specific optimizations 6 | Execution Providers (CoreML, NNAPI, QNN, DirectML) 6 | Broad model compatibility; heavier runtime footprint; potential semantic conversion gaps 4 | +| **llama.cpp** | Focused on GGUF format; highly optimized CPU implementations 6 | Contiguous buffer pooling; direct memory mapping (mmap) 9 | GGML Backend supporting Armv9 KleidiAI, SVE2, and SME 9 | Ultra-fast single-batch CPU execution; lacks sophisticated multi-accelerator partitioning 6 | +| **Google MediaPipe** | Curated task pipelines and Gemma models via Tasks API 6 | Solution-specific managed pipeline buffers 6 | GPU/NPU acceleration via Google LiteRT (formerly TensorFlow Lite) 4 | Extremely fast integration for standard vision/text tasks; highly rigid; limited custom LLM flexibility 6 | +| **Cactus SDK** | Unified API wrapper over native engines 6 | Managed abstract buffers with automatic host recycling 6 | Native framework delegates with automated fallback 6 | Simplest developer experience; unique Hybrid Cloud \+ On-Device routing; proprietary 6 | + +### **ExecuTorch Memory and Delegate Architecture** + +ExecuTorch eliminates the execution overhead associated with dynamic interpreters by utilizing ahead-of-time (AOT) graph-level compilation.4 Models are exported into an intermediate representation called Edge Dialect, which converts scalar types into tensors and flattens operators into the standardized Core ATen operator set.11 This standardization minimizes the operator surface that third-party hardware compilers must implement, enabling robust execution across diverse mobile chipsets.11 +Memory optimization within ExecuTorch is handled by the MemoryManager and MemoryAllocator abstractions.7 During AOT compilation, a static memory planner maps the lifespans of all intermediate tensors and packs them into a contiguous, user-allocated buffer.4 This approach eliminates the need for dynamic heap allocations (malloc or new) during the inference loop, preventing memory fragmentation and potential application crashes on mobile operating systems.7 +ExecuTorch supports memory hierarchy planning, enabling developers to map high-frequency, mutable intermediate tensors (such as Key-Value caches) to fast on-chip Static Random-Access Memory (SRAM), while constant weights are stored in slower Dynamic Random-Access Memory (DRAM).7 Constant weights are memory-mapped directly from the compiled .pte file using the DataLoader interface, avoiding redundant copying and minimizing the system memory footprint.7 +To leverage mobile NPUs, ExecuTorch uses a delegate system that partitions the computational graph.4 Subgraphs containing compatible operators are delegated to hardware-specific compilation binaries (e.g., Qualcomm QNN or Apple CoreML), while incompatible operators fall back to ExecuTorch's portable C++ operator library.4 +On Arm Cortex-A architectures, ExecuTorch integrates the Arm KleidiAI library via the XNNPACK delegate.10 This integration provides highly optimized kernel routines for quantized models executing on Armv9 CPUs with the i8mm ISA extension.10 On consumer mobile hardware, this co-designed software stack achieves speeds exceeding 350 tokens per second during the prefill stage for 4-bit block-quantized Llama models.10 + +### **llama.cpp and Snapdragon Compilation Architecture** + +For deployments targeting CPU-bound execution, llama.cpp remains a highly optimized framework.9 It achieves high efficiency by utilizing direct C/C++ implementations, custom quantization schemes, and assembly-level hardware instructions.8 +When targeting ARMv9-compatible mobile processors (such as the Snapdragon 8 Gen 5), compilation flags must be explicitly configured to leverage KleidiAI routines and specialized vector extensions.9 The table below outlines the compilation variables required to compile llama.cpp for modern Android platforms using the Android NDK.9 + +| Compilation Parameter | Value Configuration | Target Functionality | +| :---- | :---- | :---- | +| **CMAKE\_TOOLCHAIN\_FILE** | $NDK\_PATH/build/cmake/android.toolchain.cmake | Configures the cross-compilation environment for Android platforms 9 | +| **ANDROID\_ABI** | arm64-v8a | Restricts the target architecture to 64-bit ARM mobile processors 9 | +| **ANDROID\_PLATFORM** | android-29 | Sets the minimum Android API level (Android 10\) to balance compatibility and features 9 | +| **GGML\_CPU\_KLEIDIAI** | ON | Integrates Arm KleidiAI optimized mathematical kernel routines 9 | +| **GGML\_SYSTEM\_ARCH** | ARM | Directs the build system to generate ARM-specific instruction sets 9 | +| **GGML\_CPU\_AARCH64** | ON | Enables 64-bit execution optimizations for ARM architectures 9 | +| **GGML\_CPU\_ARM\_ARCH** | armv9.2-a+sve2+sme+dotprod+i8mm | Targets ARMv9 extensions, including Scalable Vector Extension (SVE2), Scalable Matrix Extension (SME), dot-product instructions, and 8-bit integer matrix multiplication 9 | +| **CMAKE\_C\_FLAGS / CMAKE\_CXX\_FLAGS** | \-march=armv9.2-a+sve2+sme+dotprod+i8mm | Passes machine-level architecture parameters to the compiler to generate optimized machine code 9 | + +## **SOTA On-Device Foundation Model Architectures** + +In addition to runtime optimizations, deploying language models on edge devices requires utilizing foundation models specifically co-designed for mobile hardware constraints.14 Recent peer-reviewed academic work has introduced architectural designs that maximize performance-per-parameter under tight latency constraints.14 + +### **MobileLLM-Flash** + +The MobileLLM-Flash family of foundation models (available in 350M, 650M, and 1.4B parameter variants) is designed using a hardware-in-the-loop architecture search under strict mobile latency constraints.14 Rather than relying on specialized attention mechanisms that require custom hardware kernels, MobileLLM-Flash uses standard attention mechanisms combined with attention skipping for long-context acceleration.14 +By treating candidate architectures as pruned variants of a larger pretrained backbone with inherited weights, the search space jointly optimizes layer dimensions, depth, and attention patterns.14 This optimization yields models compatible with standard mobile runtimes, such as ExecuTorch, delivering up to a 1.8x speedup in prefill latency and a 1.6x speedup in decoding throughput on mobile CPUs while supporting context windows up to 8k tokens.14 + +### **MobileLLM-Pro** + +The MobileLLM-Pro architecture addresses the challenge of running highly capable 1-billion-parameter models on mobile hardware while supporting context windows up to 128,000 tokens.15 To achieve this capability without exceeding memory limits, the model integrates two main innovations 15: + +* **Implicit Positional Distillation**: This technique transfers long-context capabilities from a larger teacher model to the compact student model during training, minimizing performance degradation at extended sequence lengths.15 +* **Specialist Model Merging**: This framework merges multiple domain-specific expert models into a single, compact foundation model without increasing overall parameter size.15 + +Furthermore, MobileLLM-Pro shows minimal performance degradation when subjected to 4-bit quantization, allowing it to maintain reasoning and comprehension capabilities within a small physical memory footprint.15 + +## **Context Pruning and Prompt Compression Frameworks** + +To further reduce latency and minimize the memory footprint of the Key-Value (KV) cache, prompt compression is employed as a key pre-processing step.16 Prompt compression is defined as mapping an original input prompt of length to a compressed prompt of length , minimizing semantic divergence under a fixed budget constraint.16 + +In soft compression paradigms, segments of the prompt are mapped to continuous token vectors via a frozen encoder and a trainable bridge projection 16: + +The model is optimized using a joint reconstruction and alignment loss 16: + +For mobile pipelines, hard compression is preferred because it produces standard text tokens that are fully compatible with black-box and localized edge-native LLMs.17 + +### **The LLMLingua Family** + +The LLMLingua framework uses a small, well-trained auxiliary language model (such as GPT-2 Small or a 1B-parameter mobile model) to compute the perplexity and conditional probabilities of sentences and individual tokens within a long prompt.17 Since natural language contains significant redundancy, tokens that contribute minimally to the overall information density have low perplexity and can be safely pruned.17 + +Raw Prompt \---\> \---\> Coarse-Grained Sentence Pruning (PPL-based) + | + v + \---\> Fine-Grained Segment Compression + | + v + \---\> Target LLM Semantic Mapping \---\> Compressed Prompt + +1. **Budget Controller**: This module dynamically distributes the target compression ratio across different segments of the prompt.18 It prioritizes instructions and final questions, forcing higher compression rates onto intermediate few-shot demonstrations.20 +2. **Iterative Token-Level Compression**: Simple token pruning ignores the changing conditional probabilities of downstream tokens as upstream tokens are removed.18 LLMLingua resolves this by dividing the prompt into segments and iteratively updating conditional probabilities, preserving linguistic coherence and preventing critical context loss.18 +3. **Distribution Alignment**: Small compression models exhibit semantic gaps when compared to massive target black-box or local edge models.17 LLMLingua introduces an alignment mechanism that forces the perplexity estimation of the small model to closely match the probability distribution of the target model, preserving reasoning and in-context learning capabilities up to a 20x compression ratio with only a 1.5% drop in performance.17 + +| Prompt Compression Engine | Primary Compression Mechanism | Optimization Focus | Speed and Throughput Characteristics | +| :---- | :---- | :---- | :---- | +| **LLMLingua** | Iterative token-level perplexity estimation and distribution alignment 17 | Retains in-context learning and multi-step reasoning capabilities 17 | Moderate latency; requires multiple forward passes of the small model 17 | +| **LongLLMLingua** | Query-aware token routing and key information reorganization 18 | Long-context retrieval-augmented generation (RAG) tasks 18 | Reduces context size up to 4x while improving key information retrieval by 17.1% 19 | +| **LLMLingua-2** | Token classification using a bidirectional BERT encoder 19 | Task-agnostic, low-latency, real-time edge processing 19 | **3x to 6x faster** than original LLMLingua, making it ideal for edge devices 19 | + +LLMLingua-2 is the most suitable compression engine for mobile pipelines.19 By formulating compression as a binary token classification task, it determines which tokens to preserve in a single forward pass.19 This approach eliminates the computational overhead of iterative perplexity calculations, allowing on-device prompt optimization to run in real time.19 + +## **Heterogeneous and Speculative Decoding Systems** + +Speculative decoding is an optimization technique that accelerates autoregressive generation by predicting and verifying multiple tokens simultaneously.22 A small, fast draft model proposes a sequence of candidate tokens (typically 3 to 12 tokens), which are then verified by a larger target model in a single, parallel forward pass.22 Rejection sampling determines which candidates are accepted or rejected.22 This shift from a memory-bound sequential bottleneck to a compute-bound parallel operation dramatically reduces inference latency on modern hardware.5 + +Draft Stage (SSM on DRAM/NPU) \---\> Proposes 3-12 Candidate Tokens + | + v +Verification Stage (Target on NPU) \---\> Parallel Single-Pass Validation + | + v +Rejection Sampling \---\> Commits Accepted Prefix / Prunes Rejected Tail + +While speculative decoding works well on high-power server accelerators, standard speculative frameworks degrade on mobile hardware because of the overheads associated with running multiple active models in RAM and switching graphs on the NPU.23 Two state-of-the-art frameworks, sd.npu and Lever, address these limitations through hardware-cooperative optimizations.24 + +### **sd.npu (CoordGen): NPU-Centric Speculative Orchestration** + +Mobile Neural Processing Units (NPUs) are domain-specific architectures optimized for massive tensor operations.24 They operate with high efficiency when execution graphs have static, predictable shapes and input batches are large.23 However, speculative decoding introduces dynamic, fragmented shapes that leave up to 75% of the NPU's compute capacity underutilized.24 +The sd.npu framework (also known as CoordGen) maximizes NPU efficiency through three synergistic components 1: + +* **Progressive Graph Scheduling**: Traditional setups keep only one static graph resident in the NPU memory, forcing a high-latency teardown-and-load cycle when transitioning from the long-sequence prefill graph to the short-sequence decoding graph.23 sd.npu partitions the model and progressively switches prefill graphs to decoding graphs block-by-block.23 This transition is overlapped with chunked prefill operations, eliminating graph-switching latency entirely.24 +* **In-Context Distribution Calibration**: Retrieval-based drafts (which build drafts from local historical context or database documents) often suffer from lexical divergence, causing high rejection rates during verification.1 sd.npu leverages the logits generated during the prefill phase to construct a model-calibrated token tree using a lightweight depth-first search (DFS).1 Retrieving from this calibrated tree ensures that draft sequences match the target model's output distribution without requiring additional training or runtime execution.1 +* **NPU-Optimized Draft Reuse**: NPUs are weight-stationary architectures that operate efficiently only when sequence lengths are sufficiently large (tokens) to saturate parallel compute lanes.23 Since retrieval drafts are often short (tokens), over 70% of NPU capacity sits idle.24 sd.npu selectively identifies and reuses high-confidence rejected tokens from prior steps, extending draft length and converting underutilized NPU capacity into valid generation throughput.23 + +### **Lever: Speculative Decoding with DRAM-Flash Heterogeneous Storage** + +When target models exceed smartphone DRAM limits, they must reside in flash storage, leaving only a tiny speculative draft model resident in DRAM.2 Flash-backed autoregressive decoding is slow because loading target weights for every sequential step triggers sustained, high-latency flash I/O operations.27 The Lever architecture utilizes the DRAM-resident draft model to generate candidate token trees, ensuring that the target model in flash is invoked only once per multiple tokens, thereby amortizing I/O costs.2 +To optimize this process, Lever implements three main components 2: + +* **Mobile-Optimized Drafting Construction**: Instead of building linear sequences, the DRAM-resident draft model constructs token trees using an I/O- and compute-aware gain-cost objective.2 It estimates the marginal gain of candidate tokens and balances this against mobile-specific verification costs, prioritizing high-value branches.2 +* **Predictor-Based Verification Pruning**: To prevent the flash-backed model from performing unnecessary computations, Lever inserts a lightweight predictor at an intermediate target layer.2 This predictor evaluates early hidden states to predict and prune low-value draft branches before they traverse the remaining target layers, saving compute cycles.2 +* **Hardware-Hybrid Execution**: Lever maps computations across hardware based on workload profiles.2 Draft tree generation is scheduled in batches on the NPU to maximize parallel efficiency.25 During verification, the target transformer layers are executed on the NPU, but final output projection and logit calculation are offloaded to the CPU to be executed on demand only along the accepted path, reducing redundant NPU verification calculations.2 + +## **Grammar-Constrained Token Masking and Agentic Tool Routing** + +On-device agent applications require language models to produce highly structured output formats (such as JSON schemas or dynamic tool calls) to reliably interact with local system APIs, databases, and UI components.28 Freeform natural language introductions introduce the risk of syntax and parsing failures, which can break downstream application execution.28 +Grammar-constrained decoding solves this formatting challenge at the generation layer.28 At each decoding step, the generation engine uses a formal grammar (such as GBNF or a JSON schema) to identify valid next tokens and masks out invalid tokens by setting their logits to negative infinity () before sampling.28 This approach guarantees 100% compliance with the specified schema, eliminating the need for post-generation parsing or retry-on-invalid-JSON routines.28 + +Target Model Logits \---\> \[ Grammar-Constrained Mask Filter \] \---\> \[ Masked Logits \] \---\> Token Sampler + ^ + (Dynamic FSM State) + +Traditional grammar engines are static, requiring substantial compilation time (often several seconds) to preprocess schemas and generate state-machine lookup tables.32 This latency is unacceptable for mobile agents, where schemas change dynamically on a per-request basis.32 + +### **XGrammar-2: High-Performance Grammar Engine** + +XGrammar-2 is a high-performance structured generation engine designed specifically for dynamic agentic workloads.32 It addresses the computational overhead of dynamic grammars through four key optimizations 32: + +* **TagDispatch**: Instead of compiling monolithic grammars that represent all possible tool schemas simultaneously, TagDispatch treats structured generation as a dynamic, tag-triggered dispatching process.32 Parameterized as a triple containing tag strings, sub-grammars, and stop strings, it runs in a lightweight "Dispatching Mode" using an Aho-Corasick automaton to scan for tags in the output stream.33 Once a tag is matched, the engine switches to "Dispatched Mode" and applies the corresponding sub-grammar constraint until a stop string is reached, avoiding the memory and latency overhead of compiling and tracking monolithic grammars.32 +* **Cross-Grammar Cache**: XGrammar-2 observes that different JSON schemas share identical sub-structures (such as generic string patterns, nested objects, or arrays).32 The engine compiles grammars into Finite State Machines (FSMs), hashes their acyclic and cyclic portions, and stores them in a global cross-grammar cache.35 By reusing these precompiled FSM sub-structures across different requests, the engine drops compilation latency to single-digit milliseconds.33 +* **Repetition State Compression**: Large arrays or repeating structures (e.g., {"type": "array", "maxItems": 1000000}) typically scale preprocessing time linearly with the repetition count.36 XGrammar-2 introduces a specialized repetition grammar primitive that compresses the state space.36 This primitive reduces preprocessing complexity to scaling time, preventing runtime execution pauses.32 +* **Partial Just-in-Time (JIT) Mask Compilation**: Building a complete token-prefix-to-allowed-token mask cache for large schemas can block inference for several seconds.33 XGrammar-2 employs a partial-JIT approach.33 During the prefill phase, the engine analyzes grammar states and compiles only the top\-most computationally expensive states within a set time budget.33 The remaining states are compiled on-demand during decoding.33 Because compile-on-demand steps are lightweight, their execution is masked by the model's forward inference pass, eliminating runtime latency spikes.33 + +## **Localized Decoding-Time Safety Guardrails and Content Policy Enforcement** + +Deploying models directly on user devices requires robust, localized safety guardrails to prevent the generation of harmful, toxic, or policy-violating content.37 Calling cloud-based moderation APIs is not feasible due to offline requirements and latency constraints, while hosting full-sized classification models on-device introduces prohibitive memory overhead.40 SOTA solutions integrate safety enforcement directly into the decoding loop, altering token selection at runtime.37 + +### **SecDecoding: Dual-Contrastive Logit Steering** + +SecDecoding is a modular, decoding-time defense framework that blocks jailbreaks and adversarial prompt-injection attacks without degrading downstream model performance.37 Rather than running an expensive classifier, SecDecoding utilizes a pair of lightweight, auxiliary 1B-parameter models: a standard base model and a safety-fine-tuned expert.37 +During generation, the divergence between the output distributions of these two auxiliary models is calculated to isolate a token-level safety signal.37 This signal is applied as a dynamic probabilistic constraint, modifying the target model's logits and steering the generation trajectory away from unsafe paths while preserving the helpfulness of safe responses.37 Because it operates on hidden states, SecDecoding is compatible with speculative decoding pipelines, providing up to a 1.5x inference speedup by overlapping safety checks with speculative verification.37 + + + | + v (Measure Divergence) +Target Model Logits \---\> \---\> Dynamic Logit Adjustment \---\> Safe Output + ^ + | + + +### **Claim-Based Stream Decoding (CSD)** + +For strict regulatory compliance, Claim-Based Stream Decoding (CSD) provides provable, certifiable safety bounds via conformal analysis.39 CSD divides the target model’s sequential output stream into discrete, semantic claims using a lightweight parsing model.39 A streaming guardrail model then evaluates each completed claim, calculating safety risk.39 +If a claim's safety risk exceeds a mathematically certified threshold, the system initiates a backtracking routine.39 By rewinding the KV-cache to the start of the flagged claim, the system blocks the unsafe output and resamples a safe alternative sequence, ensuring grammatical fluency while maintaining theoretical safety guarantees.39 + +### **Training-Free and Multi-Task Guardrails** + +When auxiliary models are unavailable, developers can implement training-free methods like **Gradient-Controlled Decoding** (GCD), which utilizes fixed anchor tokens—such as the acceptance anchor "Sure" and the refusal anchor "Sorry"—to dynamically tighten decision boundaries, significantly reducing false-positive over-refusals on benign queries.44 +For applications requiring explainable moderation, the **Lightweight Explainable Guardrail** (LEG) utilizes a multi-task learning architecture to classify prompt safety while simultaneously highlighting the specific words triggering the decision.41 This design matches or exceeds the detection accuracy of larger models while operating at minimal computational cost.41 + +## **Architectural Implementation Synthesis for the SDK Pipeline** + +To build a high-performance on-device LLM SDK, developers must integrate these optimization layers into a unified pipeline. The pipeline should be written in clean, compile-optimized C++ wrapping native mobile delegates, ensuring seamless operation across both iOS and Android platforms.4 + +\[ User Input / Context \] + | + v +\========================================= PRE-PROCESSING PHASE \========================================= + + \--\> Compresses context/history in a single bidirectional pass (3x-6x faster than original LLMLingua) + \--\> Reduces prefill latency and initial KV-cache footprint + | + v +\=========================================== PREFILL PHASE \=========================================== + + \--\> Progressive Graph Scheduling loads the decoding graph block-by-block + \--\> Switches graphs during chunked prefill computation to eliminate switching latency + \--\> Generates initial logits from the compressed prompt + | + v +\======================================== CALIBRATION PHASE \======================================== + + \--\> Uses prefill logits to construct a model-calibrated DFS token tree + \--\> Aligns context semantics with model expectations to reduce draft rejections + | + v +\=========================================== DECODING PHASE \=========================================== +Loop until End-of-Sequence (EOS) or stop criteria met: + + 1\. + \--\> Generates draft sequences from the calibrated DFS tree on the NPU + \--\> Lengthens drafts via NPU-Optimized Draft Reuse, reclaiming prior rejected tokens + + 2\. \[ Parallel Verification \] + \--\> Verifies draft sequences in parallel on the NPU + \--\> Applies intermediate predictor pruning to abort unpromising branches early + + 3\. + \--\> Employs TagDispatch to run in lightweight Aho-Corasick dispatching mode \[32, 33, 35\] + \--\> On tag match, switches to Dispatched Mode, applying the schema constraint + \--\> Uses Cross-Grammar Cache & partial-JIT to generate token masks with zero latency overhead \[33, 36\] + + 4\. + \--\> Adjusts target model logits using dual-contrastive auxiliary steering (Base vs. Expert 1B) + \--\> Redirects output trajectories away from jailbreaks in real-time + + 5\. + \--\> Commits accepted, safe, schema-compliant tokens to the active KV-cache \[22, 24, 39\] +\======================================================================================================== + | + v + +### **Actionable Integration Directives for SDK Construction** + +To execute this architectural design, developers must configure the following low-level system settings within the C++ library implementation: + +#### **1\. Implement a Zero-Heap Allocation Memory Strategy** + +Within the ExecuTorch runtime configuration, standard memory allocation mechanisms must be completely bypassed during the active inference loop.7 Developers should instantiate a contiguous, pre-planned memory buffer using a static hierarchy, mapping fast-changing tensors (such as KV-cache blocks) directly to the SoC's high-speed SRAM.7 Constant weights must be mapped using mmap to allow multiple processes to reference the model binary directly from storage without causing memory inflation or triggering Out-of-Memory (OOM) process termination.7 + +#### **2\. Optimize CPU-NPU Threading and Queue Management** + +To minimize execution bubbles, speculative drafting on the host CPU and verification on the NPU must be managed asynchronously.25 The C++ SDK should maintain parallel circular queues where the NPU processes verification batches for sequence segment , while the CPU concurrently runs draft tree generation for segment .25 Thread priorities should be bound to native platform scheduling engines (such as Android's pthread\_setschedparam or iOS's Grand Central Dispatch QoS) to prevent background UI rendering threads from interrupting inference cycles.4 + +#### **3\. Dynamically Compact and Re-index the KV-Cache** + +When using LLMLingua-2 for context pruning, the SDK must dynamically compact and re-index active KV-cache allocations.7 Pruning tokens creates fragmented holes in the cache memory.4 The C++ pipeline must execute a quick, memory-copy-free tensor re-indexing operation at the start of the prefill stage, ensuring that contiguous memory blocks are presented to the NPU's static verification graphs.7 + +#### **4\. Compile Target-Specific Hardware Instructions** + +When building GGUF or portable operator libraries, compilation must target ARMv9 extensions to leverage SVE2 and SME instruction sets.9 On devices with the Snapdragon 8 Gen 5 or equivalent SoCs, these vectorized kernels execute matrix operations up to 20% faster than standard Neon vector code, allowing the CPU to execute auxiliary safety and grammar engines with minimal latency overhead.9 +By implementing this integrated architectural pipeline, the edge-native SDK overcomes the memory and processing limitations of smartphone hardware.4 It delivers a private, reliable, and highly responsive user experience, enabling complex agentic behaviors directly on user devices.1 + +#### **עבודות שצוטטו** + +1. Accelerating Mobile Language Model Generation via Hybrid Context and Hardware Coordination \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2510.15312v1](https://arxiv.org/html/2510.15312v1) +2. Lever: Speculative LLM Inference on Smartphones \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/pdf/2605.16786](https://arxiv.org/pdf/2605.16786) +3. \[2510.15312\] Accelerating Mobile Language Model via Speculative Decoding and NPU-Coordinated Execution \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/abs/2510.15312](https://arxiv.org/abs/2510.15312) +4. ExecuTorch \-- A Unified PyTorch Solution to Run AI Models On-Device \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/pdf/2605.08195](https://arxiv.org/pdf/2605.08195) +5. On-Device LLMs: State of the Union, 2026 \- Vikas Chandra, נרשמה גישה בתאריך יוני 10, 2026, [https://v-chandra.github.io/on-device-llms/](https://v-chandra.github.io/on-device-llms/) +6. ExecuTorch vs ONNX Runtime: PyTorch Native vs Universal Model ..., נרשמה גישה בתאריך יוני 10, 2026, [https://cactuscompute.com/compare/executorch-vs-onnx-runtime](https://cactuscompute.com/compare/executorch-vs-onnx-runtime) +7. ExecuTorch Runtime Overview \- PyTorch documentation, נרשמה גישה בתאריך יוני 10, 2026, [https://docs.pytorch.org/executorch/0.4/runtime-overview.html](https://docs.pytorch.org/executorch/0.4/runtime-overview.html) +8. Accelerating Phi-2, CodeLlama, Gemma and other Gen AI models with ONNX Runtime, נרשמה גישה בתאריך יוני 10, 2026, [https://onnxruntime.ai/blogs/accelerating-phi-2](https://onnxruntime.ai/blogs/accelerating-phi-2) +9. Accelerating LLAMA inference on mobile CPUs using Qualcomm Matrix Extensions, נרשמה גישה בתאריך יוני 10, 2026, [https://www.qualcomm.com/developer/blog/2026/04/llama-models-acceleration-on-cpu-qmx](https://www.qualcomm.com/developer/blog/2026/04/llama-models-acceleration-on-cpu-qmx) +10. Unleashing the Power of AI on Mobile: LLM Inference for Llama 3.2 Quantized Models with ExecuTorch and KleidiAI \- Arm Developer, נרשמה גישה בתאריך יוני 10, 2026, [https://developer.arm.com/community/arm-community-blogs/b/ai-blog/posts/llm-inference-llama-quantized-models-executorch-kleidiai](https://developer.arm.com/community/arm-community-blogs/b/ai-blog/posts/llm-inference-llama-quantized-models-executorch-kleidiai) +11. How ExecuTorch Works \- PyTorch documentation, נרשמה גישה בתאריך יוני 10, 2026, [https://docs.pytorch.org/executorch/stable/intro-how-it-works.html](https://docs.pytorch.org/executorch/stable/intro-how-it-works.html) +12. Concepts — ExecuTorch 1.0 documentation, נרשמה גישה בתאריך יוני 10, 2026, [https://docs.pytorch.org/executorch/1.0/concepts.html](https://docs.pytorch.org/executorch/1.0/concepts.html) +13. Understanding Backends and Delegates — ExecuTorch 1.3 documentation, נרשמה גישה בתאריך יוני 10, 2026, [https://docs.pytorch.org/executorch/stable/compiler-delegate-and-partitioner.html](https://docs.pytorch.org/executorch/stable/compiler-delegate-and-partitioner.html) +14. \[2603.15954\] MobileLLM-Flash: Latency-Guided On-Device LLM Design for Industry Scale Deployment \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/abs/2603.15954](https://arxiv.org/abs/2603.15954) +15. \[2511.06719\] MobileLLM-Pro Technical Report \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/abs/2511.06719](https://arxiv.org/abs/2511.06719) +16. Prompt Compression for LLMs \- Emergent Mind, נרשמה גישה בתאריך יוני 10, 2026, [https://www.emergentmind.com/topics/prompt-compression-for-large-language-models](https://www.emergentmind.com/topics/prompt-compression-for-large-language-models) +17. Compressing Prompts for Accelerated Inference of Large Language Models \- LLMLingua, נרשמה גישה בתאריך יוני 10, 2026, [https://llmlingua.com/llmlingua.html](https://llmlingua.com/llmlingua.html) +18. LLMLingua: Innovating LLM efficiency with prompt compression \- Microsoft Research, נרשמה גישה בתאריך יוני 10, 2026, [https://www.microsoft.com/en-us/research/blog/llmlingua-innovating-llm-efficiency-with-prompt-compression/](https://www.microsoft.com/en-us/research/blog/llmlingua-innovating-llm-efficiency-with-prompt-compression/) +19. Prompt Compression in Large Language Models (LLMs): Making Every Token Count | by Sahin Ahmed(Data Scientist/MLE) | Medium, נרשמה גישה בתאריך יוני 10, 2026, [https://medium.com/@sahin.samia/prompt-compression-in-large-language-models-llms-making-every-token-count-078a2d1c7e03](https://medium.com/@sahin.samia/prompt-compression-in-large-language-models-llms-making-every-token-count-078a2d1c7e03) +20. Compressing Prompts with LLMLingua: Reduce Costs, Retain Performance \- PromptHub, נרשמה גישה בתאריך יוני 10, 2026, [https://www.prompthub.us/blog/compressing-prompts-with-llmlingua-reduce-costs-retain-performance](https://www.prompthub.us/blog/compressing-prompts-with-llmlingua-reduce-costs-retain-performance) +21. LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2310.05736v2](https://arxiv.org/html/2310.05736v2) +22. An Introduction to Speculative Decoding for Reducing Latency in AI Inference | NVIDIA Technical Blog, נרשמה גישה בתאריך יוני 10, 2026, [https://developer.nvidia.com/blog/an-introduction-to-speculative-decoding-for-reducing-latency-in-ai-inference/](https://developer.nvidia.com/blog/an-introduction-to-speculative-decoding-for-reducing-latency-in-ai-inference/) +23. Accelerating Mobile Language Model via Speculative Decoding and NPU-Coordinated Execution \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2510.15312v3](https://arxiv.org/html/2510.15312v3) +24. Accelerating Mobile Language Model via Speculative Decoding and NPU-Coordinated Execution \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2510.15312v4](https://arxiv.org/html/2510.15312v4) +25. Lever: Speculative LLM Inference on Smartphones \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2605.16786v1](https://arxiv.org/html/2605.16786v1) +26. Accelerating Mobile Language Model via Speculative Decoding and NPU-Coordinated Execution \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/pdf/2510.15312](https://arxiv.org/pdf/2510.15312) +27. \[2605.16786\] Lever: Speculative LLM Inference on Smartphones \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/abs/2605.16786](https://arxiv.org/abs/2605.16786) +28. LM-Kit.NET Structured Output: Constrained JSON Generation in C\# .NET, נרשמה גישה בתאריך יוני 10, 2026, [https://docs.lm-kit.com/lm-kit-net/guides/glossary/structured-output.html](https://docs.lm-kit.com/lm-kit-net/guides/glossary/structured-output.html) +29. \[RFC\] Constrained decoding for extension/llm — is it worth doing? · Issue \#19215 · pytorch/executorch \- GitHub, נרשמה גישה בתאריך יוני 10, 2026, [https://github.com/pytorch/executorch/issues/19215](https://github.com/pytorch/executorch/issues/19215) +30. Structured Output Generation in LLMs: JSON Schema and Grammar-Based Decoding | by Emre Karatas | Medium, נרשמה גישה בתאריך יוני 10, 2026, [https://medium.com/@emrekaratas-ai/structured-output-generation-in-llms-json-schema-and-grammar-based-decoding-6a5c58b698a6](https://medium.com/@emrekaratas-ai/structured-output-generation-in-llms-json-schema-and-grammar-based-decoding-6a5c58b698a6) +31. Structured Output of Large Language Models | Niklas Heidloff, נרשמה גישה בתאריך יוני 10, 2026, [https://heidloff.net/article/llm-structured-output/](https://heidloff.net/article/llm-structured-output/) +32. XGrammar-2: Efficient Dynamic Structured Generation Engine for Agentic LLMs \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2601.04426v2](https://arxiv.org/html/2601.04426v2) +33. XGrammar 2: High-Performance Grammar Systems \- Emergent Mind, נרשמה גישה בתאריך יוני 10, 2026, [https://www.emergentmind.com/topics/xgrammar-2](https://www.emergentmind.com/topics/xgrammar-2) +34. 1 Introduction \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2601.04426v1](https://arxiv.org/html/2601.04426v1) +35. XGrammar-2: Efficient Dynamic Structured Generation Engine for Agentic LLMs \- arXiv, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2601.04426v3](https://arxiv.org/html/2601.04426v3) +36. XGrammar-2: Fast and Customizable Structured Generation for Tool Calling and Agents, נרשמה גישה בתאריך יוני 10, 2026, [https://blog.mlc.ai/2026/05/04/xgrammar-2-fast-customizable-structured-generation](https://blog.mlc.ai/2026/05/04/xgrammar-2-fast-customizable-structured-generation) +37. SecDecoding: Steerable Decoding for Safer LLM Generation \- ACL Anthology, נרשמה גישה בתאריך יוני 10, 2026, [https://aclanthology.org/2025.findings-emnlp.1118/](https://aclanthology.org/2025.findings-emnlp.1118/) +38. Strengthening LLM guardrails with synthetic data generation \- JPMorganChase, נרשמה גישה בתאריך יוני 10, 2026, [https://www.jpmorganchase.com/about/technology/blog/fence-framework](https://www.jpmorganchase.com/about/technology/blog/fence-framework) +39. C-SafeGen: Certified Safe LLM Generation with Claim-Based Streaming Guardrails, נרשמה גישה בתאריך יוני 10, 2026, [https://neurips.cc/virtual/2025/poster/116139](https://neurips.cc/virtual/2025/poster/116139) +40. LLM Guardrails in Production: Input, Output, and Runtime Checks That Actually Work, נרשמה גישה בתאריך יוני 10, 2026, [https://www.kalviumlabs.ai/blog/guardrails-for-llm-applications/](https://www.kalviumlabs.ai/blog/guardrails-for-llm-applications/) +41. A Lightweight Explainable Guardrail for Prompt Safety \- OpenReview, נרשמה גישה בתאריך יוני 10, 2026, [https://openreview.net/forum?id=M4He5YzG44](https://openreview.net/forum?id=M4He5YzG44) +42. SecDecoding: Steerable Decoding for Safer LLM Generation \- ACL Anthology, נרשמה גישה בתאריך יוני 10, 2026, [https://aclanthology.org/2025.findings-emnlp.1118.pdf](https://aclanthology.org/2025.findings-emnlp.1118.pdf) +43. C-SafeGen: Certified Safe LLM Generation with Claim-Based Streaming Guardrails | OpenReview, נרשמה גישה בתאריך יוני 10, 2026, [https://openreview.net/forum?id=nOsEyBGk1I](https://openreview.net/forum?id=nOsEyBGk1I) +44. Gradient-Controlled Decoding: A Safety Guardrail for LLMs with Dual-Anchor Steering, נרשמה גישה בתאריך יוני 10, 2026, [https://arxiv.org/html/2604.05179v1](https://arxiv.org/html/2604.05179v1) diff --git a/docs/research/SecDecoding alternatives research.md b/docs/research/SecDecoding alternatives research.md new file mode 100644 index 0000000..a6c172f --- /dev/null +++ b/docs/research/SecDecoding alternatives research.md @@ -0,0 +1,355 @@ +# Implementing SecDecoding and CSD Style Model Backed Safety With Runtime Backtracking on Offline Air Gapped Edge Devices + +## Executive summary + +This report assumes no fixed hardware platform, operating system, latency target, or power budget beyond the stated requirement that the system be fully offline and air-gapped. I also interpret **CSD** in the sense used by **Synchromesh**—**Constrained Semantic Decoding**—because that is the CSD variant most directly tied to runtime validators, semantic constraints, and rollback/backtracking mechanics. Under those assumptions, the strongest practical architecture for a **~0.5B open-weight edge deployment** is usually **not** a pure implementation of SecDecoding or a pure implementation of CSD. It is a **hybrid**: a quantized small instruct model for primary generation, a small safety sidecar or safety adapter for decoding-time steering, a **checkpointed token/KV-cache rollback loop**, and a **deterministic policy boundary** for any high-impact outputs or actions. This combines the low operational footprint of decoding-time steering with the stronger recoverability of rollback and the stronger assurance of symbolic constraints where the output space is structured. + +The central distinction is this. **SecDecoding** uses a pair of small contrastive models—a base model and a safety fine-tuned expert—to estimate a token-level safety signal from their output-distribution divergence, then adds that signal to the target model’s logits; in the published EMNLP 2025 form it is modular, resource-efficient, and compatible with speculative decoding, but it still assumes continuous access to auxiliary models during decoding. **SafeDecoding**, by contrast, intersects top-token sets from an original model and a safety-tuned expert model, then linearly reweights token probabilities; it is cheaper to implement on edge devices, especially if only applied for the first few tokens. **CSD** is different again: it is not mainly “safe steering,” but **constraint-preserving decoding** against a completion engine that defines the valid continuation set; if the completion engine is exhaustive and correct, the output satisfies those constraints by construction. + +For **0.5B-class edge models**, the most attractive deployment pattern is usually: **Qwen2.5-0.5B-Instruct or Qwen3-0.6B** as the main model; **4-bit or mixed-INT4** for GPU/NPU paths, **INT8** where tooling is simpler; **memory-mapped** model loading when available; **chunk-level** guard evaluation every few tokens instead of every token; **KV-cache checkpoints** every few tokens and at semantic boundaries; and **rollback to the last safe checkpoint** when a guard score crosses threshold. If the device must emit structured outputs—commands, JSON, configuration, SQL, or code snippets used for actuation—then add a **CSD-style grammar/semantic layer** around the final emitted channel. + +The assurance story should be explicit about what is and is not provable. You can obtain strong guarantees for **boot integrity** and **artifact integrity** with **secure boot**, **dm-verity/fs-verity**, **IMA appraisal/measurement**, and optional **TPM-based attestation**. You can obtain strong guarantees for **structured-output validity** with CSD-style completion engines. You **cannot** obtain a blanket semantic-safety guarantee for unconstrained free-form text from decoding-time steering alone, and current LLM judges are themselves brittle under prompt sensitivity and adversarial attack. For air-gapped deployments, this means the system should treat model-backed safety as **one layer** in a broader secure systems design, not as the entire safety case. + +## Definitions and the approach landscape + +### What SecDecoding is + +**SecDecoding** is a **decoding-time defense** that uses two small contrastive models—a small base model and a safety-tuned expert—to estimate a **token-level safety signal** from the difference in their logits. That signal is scaled by a dynamic factor that depends on divergence and token position, then added to the target model’s logits before sampling. In the paper’s formulation, the approach is modular, requires only the auxiliary contrastive models rather than re-training the target model, and is compatible with speculative decoding; the paper reports near-zero attack success rates across advanced jailbreak attacks while preserving helpfulness, and states that only an auxiliary **1B-parameter** model is required in the demonstrated setup. + +The important implementation consequence is that SecDecoding is **sidecar-style safety steering**. It is attractive when you want to leave the main model untouched and drive its logits with a safety signal produced by smaller models. On resource-limited edge devices, though, the need to run auxiliary models during generation means SecDecoding is easiest to justify on **GPU/NPU-equipped devices** or when you can **gate the auxiliary path** to only the riskiest phases of decoding. That last point is an engineering inference from the paper’s architecture and from edge memory/latency constraints, rather than a claim made verbatim by the authors. + +### What SafeDecoding is and how it differs + +**SafeDecoding** is also a decoding-time defense, but its mechanism is simpler and more directly edge-friendly. It fine-tunes an **expert model** from the original model, constructs a **sample space** as the intersection of top tokens from the original and expert models, and then defines a new probability distribution by interpolating the original probabilities with the expert-vs-original difference. The paper explicitly recommends applying the strategy only to the **first few decoding steps**, motivated by the observation that jailbreaks often succeed by forcing the model into an affirmative initial trajectory; this is one reason SafeDecoding is attractive on constrained hardware. + +Relative to SecDecoding, SafeDecoding is usually the better starting point for a **0.5B offline edge system** if you want the main benefit of model-backed token steering without the full overhead of continuously running separate contrastive small models at every step. It is less expressive than SecDecoding’s dynamic divergence-based safety signal, but it has a cleaner **same-family deployment story**: one main model, one PEFT-tuned safety expert, selective use on early tokens, and straightforward fallback to standard decoding outside high-risk windows. + +### What CSD is and how it differs + +In **Synchromesh**, **Constrained Semantic Decoding** is a general framework for constraining generation to a set of valid programs by querying a **completion engine** over partial outputs. The completion engine can embody both **context-free** constraints, such as syntax, and **context-sensitive** constraints, such as schema validity or type/scoping rules. The paper’s main claim is strong: when the completion engine is exhaustive and correctly implemented, CSD guarantees that all sampled outputs satisfy the implemented constraints **by construction**. The authors also report an average **8% sampling overhead** in their GPT-3 experiments. + +That makes CSD conceptually different from both SafeDecoding and SecDecoding. CSD is not primarily a “safety signal” method; it is a **valid-continuation filtering** method. It is therefore strongest when “safety” can be re-expressed as membership in a formally checkable language or prefix-closed valid set: safe JSON schemas, command allow-lists, typed tool signatures, safe SQL fragments, or restricted code subsets. It is weaker for nuanced free-form harms that do not reduce cleanly to syntax or semantics. + +### Related model-backed safety families + +The last five years of literature show several adjacent families that matter for edge implementation design: + +| Technique | Core mechanism | Strength | Main weakness | Edge relevance | +|---|---|---|---|---| +| **DExperts** | Product-of-experts combining a base LM with expert and/or anti-expert LMs at decoding time | Clean general formulation for attribute control | Extra model calls and tuning burden | Good conceptual basis for safety sidecars | +| **SafeDecoding** | Intersection of top-token sets from original and safety expert, with reweighted probabilities | Simple, effective, early-token application helps edge fit | Limited expressivity; no hard guarantees | Very good default for 0.5B edge | +| **InferAligner** | Cross-model guidance using safety steering vectors to modify activations | Decouples harmlessness from helpfulness; no full re-training of target | Requires activation intervention support | Strong where runtime hooks are available | +| **SafeInfer** | Safety amplification vector in hidden states plus safety-guided decoding using harmful-model debiasing | Combines latent steering with decoding control | More moving parts than SafeDecoding | Good for NPU/GPU phones; harder on tiny CPUs | +| **SecDecoding** | Small base + safety expert produce divergence-based safety signal added to target logits | Modular, scalable, speculative-decoding compatible | Continuous auxiliary model use | Better on stronger edge hardware | +| **DeAL** | Decoding framed as search with custom reward/alignment objectives | Flexible, customizable objectives | Slower search-heavy inference | Usually too costly for 0.5B edge unless heavily simplified | +| **CoSA** | Inference-time adaptation to natural-language safety configurations | Useful for pluralistic or changing safety requirements | Needs training for controllability | Useful when policy configs must vary offline | +| **InferenceGuard** | Constrained MDP in latent space with augmented safety state | Formal “almost sure” guarantee relative to the defined safety cost model | Guarantee is cost-model-dependent; implementation is not yet mature as edge software | Important for assurance design, not yet the easiest production choice | +| **CSD** | Completion-engine-constrained valid-token decoding | Hard validity guarantees for structured outputs | Nontrivial constraint authoring; limited free-form semantic coverage | Best assurance for structured channels | + +The most important analytical takeaway is that **“model-backed safety” is not one thing**. Some methods steer **logits**; some steer **activations**; some change **search**; some attach **formal completion engines**; some add **rollback**; and some aim for **formal cost-model guarantees**. On edge devices, this matters because different families stress different resources: logits steering stresses repeated forward passes, activation methods stress runtime hooks, completion engines stress validator complexity, and rollback stresses snapshotting and loop control. + +## Runtime backtracking mechanisms + +Runtime backtracking is the missing piece that turns safety steering into a **recoverable control loop** instead of a one-shot gate. The recent safety and code-generation literature converges on the same systems lesson: once an autoregressive model has started generating along a bad trajectory, post-hoc repair is expensive and often too late; better systems detect problems during generation and **rollback to an earlier safe prefix**. That is explicit in the 2025 ICLR backtracking paper using a special **[RESET]** token, in **RoCode**’s incremental error detection with rollback and constrained regeneration, and in **Hydra**’s checkpoint-and-rollback support with asynchronous checking. + +```mermaid +flowchart TD + A[Input prompt] --> B[Prompt risk triage] + B -->|low risk| C[Normal decode] + B -->|elevated risk| D[Safety-steered decode] + C --> E[Checkpoint every K tokens or at boundaries] + D --> E + E --> F[Chunk guard scores recent output] + F -->|safe| G[Continue decoding] + F -->|unsafe| H[Rollback to last safe checkpoint] + H --> I[Strengthen safety controls] + I --> J[Lower temperature / raise alpha / inject refusal bias] + J --> D + G --> K{Done?} + K -->|no| E + K -->|yes| L[Emit output] +``` + +### Backtracking algorithms that are actually useful on edge + +There are four algorithms worth separating: + +**Token-buffer rollback** keeps only the last *B* emitted tokens and replays from the most recent safe prefix. This is the lightest implementation, but if you do not also checkpoint the **KV cache**, replay cost will dominate latency on weaker hardware. That makes token-buffer-only rollback acceptable mainly for very short outputs or CPU-only prototypes. The practical runtimes that matter for LLM deployment already revolve around explicit **generation loops** and **KV-cache management**, which is why real implementations should treat KV state as the rollback object, not just token text. + +**Checkpointed rollback** snapshots generation state every *K* tokens and at semantic boundaries such as newline, sentence break, closing brace, or tool-call delimiter. The state to snapshot is normally: emitted token IDs, sampler RNG state, safety-control state, and **KV-cache metadata or handles**. With block-based caches, such as TensorRT-LLM’s KV block pool, storing cache **block references** is much cheaper than copying tensors; with ONNX Runtime GenAI, the generation loop and KV management are explicit enough to support the same pattern. + +**Model-self-backtracking** trains the model to emit a reset action. In the ICLR 2025 backtracking work, the model produces a **[RESET]** token when it recognizes that a partial generation is unsafe; the serving stack then discards the preceding unsafe text and allows a fresh safe continuation. This is elegant because rollback becomes model-native, and the paper reports a large safety improvement without helpfulness regression in their evaluations, but the method still requires **specialized tuning** and does not remove the need for an external guard if you care about adversarial robustness on an air-gapped device. + +**Asynchronous validator rollback** is the most systems-efficient idea in the recent literature. **Hydra** shows that checking can proceed **asynchronously** with generation, so the checker adds almost no cost when the output remains valid; when a violation does appear, the system rolls back to a more appropriate earlier checkpoint instead of repairing only the immediate token. For constrained outputs this is often superior to pure per-token CSD, and the same principle transfers well to safety checks on edge devices: score every 4–16 tokens, not every token, unless the prompt is already high-risk. + +### Reversible computation versus checkpointing + +There is a long systems literature on rollback recovery and partially reversible execution, and recent work on checkpoint-based rollback recovery in concurrent programming explicitly combines **checkpoints** with **partially reversible semantics**. That theory is useful for thinking about minimal rollback state and consistency, but today’s practical LLM runtimes do **not** implement exact reversible transformer inference. They implement **forward-only kernels** plus explicit cached state. For edge deployment, the right primitive is therefore **checkpoint-and-restore**, not reversible compute in the strict sense. + +### Causal tracing and intervention points + +“Causal tracing,” “activation patching,” and related mechanistic interpretability tools are best understood as **offline design tools**, not as default runtime components. They are useful for identifying **where** to intervene—specific layers, heads, or representations associated with unsafe drift—so that activation-steering methods such as SafeInfer or InferAligner can place their interventions more efficiently. SafeInfer itself explicitly uses activation patching to identify influential heads for its safety amplification vector. However, best-practice work on activation patching warns that patching results are sensitive to methodological choices, so they should guide engineering decisions, not be mistaken for formal proofs. + +For a **0.5B edge** implementation, the best intervention points are usually: + +1. **Prompt ingress**, for risk triage and deterministic refusal rules. +2. **Early decode**, because several defenses and analyses show that initial affirmative tokens set the unsafe trajectory. +3. **Chunk boundaries**, every few tokens, for guard scoring and rollback. +4. **Structured-output emission points**, where CSD-style validators can enforce a hard language of allowed outputs. + +## Edge device constraints and secure deployment + +### What the hardware budget really means for ~0.5B models + +A 0.5B-class model is finally in the regime where **single-device offline deployment** is realistic across several edge classes, but the memory budget is still the dominant system driver. Official LiteRT-LM community artifacts provide especially useful empirical anchors here. For **Qwen2.5-0.5B-Instruct**, a dynamic-INT8 LiteRT artifact reports on a Samsung S24 Ultra CPU path roughly **29.97 decode tok/s**, **2.31 s time-to-first-token**, **1,363 MB RSS**, and **521 MB model size**. For **Qwen3-0.6B**, a mixed-INT4 LiteRT artifact reports on-device decode rates up to **69.38 tok/s** on a GPU OpenCL path with a **585 MB** peak private footprint, while the CPU path reports about **12.90 tok/s** with around **2.9 GB** footprint; an NPU-targeted artifact is also published. These results are device-specific and not cross-platform apples-to-apples, but they show that **0.5–0.6B quantized models are operationally plausible on edge hardware**. + +The practical implication is simple: for **low-latency chat**, CPU-only SBCs are usually acceptable only for minimal wrappers and limited contexts; **GPU/NPU-equipped** devices give enough headroom to add safety sidecars, chunk guards, and rollback without making the user experience unusable. The published Qwen community artifacts also show a large difference between CPU and accelerator paths even at 0.5–0.6B scale, which is why runtime safety design and hardware choice cannot be separated. + +### Representative hardware options + +| Hardware option | Official specs | Best-fit software stack | 0.5B safety-wrapper fit | +|---|---|---|---| +| **Raspberry Pi 5** | Quad-core Cortex-A76 at 2.4 GHz; RAM variants up to 16 GB | **llama.cpp**, **LiteRT-LM**, lightweight **ONNX Runtime Mobile** paths | Feasible for **single-model** CPU inference and light first-token steering; continuous dual-model steering plus frequent rollback is likely only comfortable for short outputs and small contexts. This is an engineering inference from the published model footprints and the Pi’s CPU-only profile. | +| **Jetson Orin Nano Super Developer Kit** | 67 INT8 TOPS; Ampere GPU with 1024 CUDA cores; 8 GB LPDDR5; 7–25 W | **TensorRT-LLM**, **llama.cpp** with CUDA/Vulkan, **ONNX Runtime** CUDA/TensorRT paths | Strong default for Linux ARM edge. Main 0.5B model plus a small safety expert or risk model is realistic; full SecDecoding-style always-on sidecars become plausible here. | +| **Intel Core Ultra Series 2 edge dev kit / NUC class systems** | Up to 48 NPU TOPS, up to 67 GPU TOPS, 120 platform TOPS on the referenced dev-kit page | **OpenVINO GenAI**, **ONNX Runtime GenAI**, **ExecuTorch** where PyTorch-native deployment matters | Excellent for x86 air-gapped deployments that need stronger assurance tooling, richer logging, and a more comfortable latency budget. | +| **AMD Ryzen AI 300 / PRO 300 systems** | 50–55 NPU TOPS on listed SKUs; 15–54 W range | **ONNX Runtime**, **LiteRT-LM**, vendor-specific ROCm/DirectML where available | Good for industrial mini-PC form factors; especially attractive if you want more RAM and storage than SBCs while staying offline and power-limited. | + +### Quantization, pruning, and loading strategy + +For small edge models, **post-training quantization** is usually the first lever and **pruning** the second. **AWQ** is explicitly framed as an on-device, hardware-friendly low-bit weight-only quantization approach; **GPTQ** remains the foundational one-shot 3–4 bit quantization method; **KVQuant** is relevant when context windows matter because KV cache footprint can become the dominant working set. Official framework support aligns with that pattern: TensorRT-LLM documents multiple quantization recipes including FP4, FP8, and INT8-style variants; ONNX Runtime documents dynamic and static INT8 quantization; LiteRT-LM community artifacts show practical INT8 and mixed-INT4 deployments for 0.5–0.6B Qwen models. + +Pruning is more nuanced. Research such as **SparseGPT** and **Wanda** shows that large language models can be pruned aggressively with modest quality loss, and newer work such as **ShortGPT** and **HAPE** argues for more hardware-aware or layer-level approaches. But on **today’s edge runtimes**, sparse kernels and sparse-memory layouts are often less mature than 4-bit dense paths, so **dense low-bit quantization usually wins first** for a 0.5B deployment unless you have a very specific sparse accelerator path. That is a systems recommendation derived from the pruned-model literature plus current deployment-tool maturity, not a theorem. + +For file loading, **memory mapping** is one of the most important practical tricks on air-gapped systems. **GGUF** is explicitly designed for fast loading, and `llama.cpp` enables **memory-mapped model loading by default** with `--mmap`; its docs also expose `--mlock` to keep models resident in RAM. That combination is highly attractive on Linux edge devices because it reduces startup cost, simplifies re-use across processes, and pairs naturally with **read-only verified filesystems** such as fs-verity. + +### Inference framework choices + +For **Linux-first embedded deployments**, **llama.cpp** remains the most pragmatic starting point because it is portable, supports multiple hardware backends, uses GGUF, and makes memory mapping a first-class part of the deployment story. **ONNX Runtime GenAI** is a strong choice when you want a clean on-device generation loop with built-in **KV cache management**, logits processing, and the option to custom-build smaller runtime packages. **LiteRT-LM** is now one of the most interesting cross-platform options for phones, WebGPU-capable desktops, and some NPU paths; its public artifacts provide real on-device data for Qwen-class models. **OpenVINO GenAI** is the best fit for Intel CPU/GPU/NPU systems, while **TensorRT-LLM** is the obvious choice for NVIDIA edge GPUs. **ExecuTorch** is attractive when you want a PyTorch-native path across mobile and embedded backends. + +## Verification, assurance, and threat modeling + +### What can be guaranteed + +There are three very different assurance layers in this problem, and they should not be conflated. + +The first is **system integrity**. On Linux and Jetson-class devices, **secure boot** creates a chain of trust that prevents execution of unauthorized boot code; **dm-verity** provides cryptographic integrity checking of read-only block devices; **fs-verity** provides Merkle-tree-based authenticity and integrity protection for individual read-only files; and **IMA appraisal/measurement** can verify and measure files accessed through execution and memory mapping. This is the layer that protects your model binaries, tokenizers, safety adapters, policies, validators, and runtime itself from silent corruption. + +The second is **output-language validity**. For structured channels, a CSD-style completion engine can provide a genuine construction-time guarantee that outputs stay inside the permitted language or semantic manifold represented by the completion engine. This is the strongest available guarantee for edge systems that emit machine-consumable outputs, and it is why CSD-style wrapping is so attractive for commands, JSON, database queries, safety case files, or code destined for later execution. + +The third is **semantic harmlessness of free-form text**. Here, guarantees are much weaker. The most ambitious recent work, **InferenceGuard**, argues for “almost sure” safety relative to a defined **safety cost model** by formulating inference-time alignment as a constrained MDP in latent space. That is important research, but the guarantee is explicitly **with respect to the modeled cost structure**, not an unrestricted guarantee against all real-world harms or all adversarial prompts. It should therefore be understood as a promising formal wrapper model, not a complete replacement for other system defenses. + +### Why single-judge safety is not enough + +A central caution from the recent evaluation literature is that **LLM-based judges are themselves attack surfaces**. “Know Thy Judge” shows that small changes in output style can change false-negative rates materially and that adversarial attacks on generations can fool some judges into misclassifying harmful generations as safe. Separate systematic benchmarking of jailbreak defenses also finds that simple baselines can remain competitive and that performance varies strongly by jailbreak style, which is another warning against relying on a single clever guard model. For edge deployment, that means the monitor should be **heterogeneous**: one model-backed signal plus deterministic rules plus optional structured constraints, not one judge to rule everything. + +### Threat model for air-gapped edge deployment + +The most relevant threats are: + +**Adversarial prompt inputs.** These include jailbreak suffixes, multilingual/encoded prompts, instruction nesting, and style attacks. The JailbreakBench project is valuable here because it standardizes artifacts, system prompts, threat models, and scoring functions rather than relying on ad hoc evaluation. + +**Supply-chain and artifact corruption.** On an air-gapped device, the danger is not cloud compromise but **tampered weights**, **modified tokenizers**, **malicious LoRA adapters**, **swapped safety policies**, or altered runtime binaries introduced during provisioning or field updates. OWASP’s LLM/GenAI top-risk material highlights supply-chain vulnerabilities, training data poisoning, prompt injection, denial of service, and model theft as recurring categories. + +**Model denial of service.** Long prompts, repeated rollback cycles, and degenerate validator loops can starve limited hardware or drain power. Air-gapped systems often lack orchestration elasticity, so defensive limits must be local and deterministic. OWASP explicitly identifies model DoS as a major risk category. + +**Physical access and local tampering.** Because devices are offline, operators often assume they are safe; in reality, physical access raises the premium on **measured boot**, sealed storage, signed update bundles, and immutable logging. Secure boot, IMA, and verity layers are therefore not optional hardening extras but core parts of the model-safety stack. + +### Recommended mitigations + +The most robust mitigation stack is layered: + +1. **Cryptographic trust of the platform and artifacts**: secure boot, verified/immutable root filesystem, fs-verity on model and policy files, IMA appraisal for executables and mapped files. +2. **Deterministic ingress controls**: prompt length caps, encoding/cipher detection if relevant to mission, allow/deny policies for capability classes, and output channel separation. +3. **Model-backed steering**: SafeDecoding/SecDecoding-style token steering, ideally gated by risk. +4. **Checkpointed rollback**: last-safe-prefix restore, bounded retry count, stronger refusal bias after rollback. +5. **Structured-output enforcement**: CSD or grammar-constrained decoding for any machine-consumable output. +6. **Auditability**: sealed logs of inputs, risk scores, rollback events, hashes, and policy versions for offline forensic review. This is a design recommendation consistent with NIST AI RMF and integrity-measurement practice. + +## Recommended architectures, code patterns, and implementation plan + +### Recommended reference architecture + +```mermaid +flowchart LR + U[User / upstream system] --> P[Deterministic ingress policy] + P --> R[Prompt risk scorer] + R --> M[Main 0.5B model] + R --> S[Safety expert or safety adapter] + M --> G[Safety-aware sampler] + S --> G + G --> C[Checkpoint manager] + C --> Q[Chunk guard] + Q -->|safe| O[Output formatter] + Q -->|unsafe| B[Rollback controller] + B --> C + B --> G + O --> X[Free-form response] + O --> Y[Structured channel via CSD / grammar validator] + + subgraph Platform trust + SB[Secure boot] + FV[fs-verity / dm-verity] + IMA[IMA appraisal / measurement] + end + + SB --> P + FV --> M + FV --> S + FV --> O + IMA --> M + IMA --> S +``` + +This architecture deliberately separates **text-generation safety**, **structured-output safety**, and **platform integrity**. That separation is what lets it degrade gracefully across hardware tiers. On a CPU-only box, you can keep the safety expert tiny and invoke it selectively. On a Jetson or Core Ultra system, you can afford denser chunk guards, more frequent checkpoints, or even a closer SecDecoding-style continuous sidecar. + +### Which architecture to choose + +For most **0.5B offline edge** deployments, I recommend the following order: + +**Default recommendation: selective SafeDecoding-style steering plus rollback.** Use a main 0.5–0.6B instruct model and a same-family safety expert or LoRA safety adapter. Run the safety path on **prompt ingress** and the **first 8–32 output tokens**, then switch to normal decode unless the chunk guard detects elevated risk. This gives the best “safety per watt” profile for constrained hardware. + +**If you have stronger accelerator headroom: SecDecoding-style sidecar steering.** This becomes attractive on Jetson, Core Ultra, or similar hardware where the extra forward passes are affordable and you want a more expressive divergence-based safety signal. Pair it with chunked rollback so that safety interventions can recover, not just steer. + +**If the output is structured or security-critical: add CSD.** Any output that could trigger local automation, feed another parser, or reach an operational process should be constrained by grammar/semantic validation. In those cases, free-form natural language should be treated as advisory and the structured channel as the bounded actuation interface. + +### Code-level patterns that map well to real runtimes + +The runtime loop below is the core pattern I would actually implement in C++ or Rust on Linux edge, or in Python only for prototyping. It synthesizes the mechanisms documented in SafeDecoding, SecDecoding, Hydra, and the modern on-device runtimes with explicit KV-cache management. + +```python +class DecodeState: + def __init__(self, token_ids, kv_handle, rng_state, alpha, temperature): + self.token_ids = list(token_ids) + self.kv_handle = kv_handle # block refs or cache handle, not full copy if possible + self.rng_state = rng_state + self.alpha = alpha + self.temperature = temperature + +def generate_with_safety(prompt_ids, max_new_tokens, cfg): + output_ids = [] + checkpoints = [] + rollback_count = 0 + guard_mode = risk_triage(prompt_ids) >= cfg.risk_threshold + + state = capture_state(prompt_ids, alpha=cfg.alpha_base, temperature=cfg.temperature) + + for step in range(max_new_tokens): + main_logits = main_model.forward_next(state.token_ids, state.kv_handle) + + if guard_mode: + expert_logits = safety_expert.forward_next(state.token_ids, safety_cache_handle()) + steered_logits = safety_aware_mix(main_logits, expert_logits, alpha=state.alpha) + else: + steered_logits = main_logits + + next_id = sample_token( + steered_logits, + temperature=state.temperature, + top_p=cfg.top_p + ) + + append_token(state, next_id) + output_ids.append(next_id) + + if should_checkpoint(step, next_id, cfg): + checkpoints.append(capture_state( + state.token_ids, + alpha=state.alpha, + temperature=state.temperature + )) + + if should_guard(step, next_id, cfg): + risk = chunk_guard_score(prompt_ids, output_ids, state) + if risk >= cfg.hard_threshold: + if not checkpoints or rollback_count >= cfg.max_rollbacks: + return hard_refusal(prompt_ids) + state = restore_last_safe(checkpoints) + rollback_count += 1 + guard_mode = True + state.alpha = min(state.alpha * cfg.alpha_upscale, cfg.alpha_max) + state.temperature = min(state.temperature, cfg.safe_temperature) + inject_refusal_bias(state) + continue + elif risk >= cfg.soft_threshold: + guard_mode = True + + if is_eos(next_id): + break + + return detokenize(output_ids) +``` + +The key engineering rule is that **rollback should restore cache state, not just text state**. Otherwise every rollback becomes a full replay. On runtimes with explicit generation loops and cache handles—such as **ONNX Runtime GenAI** and **TensorRT-LLM**—this is natural. If you use `llama.cpp`, the same idea still applies: keep checkpoints sparse, checkpoint on semantic boundaries, and never let rollback loops run unbounded. + +### Practical implementation plan + +**Phase one: choose the main model and baseline runtime.** +Start with one of the published small open-weight instruct models that already have real on-device evidence: **Qwen2.5-0.5B-Instruct** or **Qwen3-0.6B** are good defaults. Use **mixed INT4** or **Q4** on accelerator paths and **INT8** on simpler CPU/NPU paths. Establish baseline **TTFT**, **decode tok/s**, **peak RSS**, and **energy per request** before adding any safety wrapper. + +**Phase two: create the safety expert or safety adapter.** +The cheapest path is SafeDecoding-style: fine-tune a same-family expert or LoRA adapter on harmful-instruction → refusal data, preserving tokenizer compatibility. If the hardware budget is tighter, a small paired refusal classifier plus prompt triage can decide when to invoke the expert path. If the hardware budget is larger, a SecDecoding-style contrastive pair is more expressive. + +**Phase three: implement selective steering and checkpointed rollback.** +Begin with steering on the **first 8–32 tokens** only. Add checkpoints every **8–16 tokens** and at sentence/JSON boundaries. Add a chunk guard every **4–8 tokens**. Fail closed after a small number of rollbacks, and log the rollback event with hashes and policy version. This is the right initial operating point for constrained devices because it captures most of the benefit of early-trajectory control while bounding sidecar cost. + +**Phase four: add structured-output enforcement where applicable.** +If the device emits JSON, commands, code, or configuration, place a grammar or completion engine around that channel. For highly sensitive flows, it is better for the model to produce a natural-language explanation and a **separately constrained machine output** than to rely on free text alone. + +**Phase five: harden the platform.** +Enable secure boot; put the model, tokenizer, adapters, and policy files on read-only, signed media; use **fs-verity** or **dm-verity**; and enable **IMA appraisal/measurement** for executables and mapped model files. If you can support periodic offline audit or a custodian station, add TPM evidence collection and Keylime-compatible measurement logs for later verification. + +### Testing methodology and metrics + +Use three test suites, not one. First, a **malicious-input suite** such as JailbreakBench-style behaviors and attack artifacts; second, a **benign-utility suite** covering the mission tasks actually expected of the device; third, a **systems-resilience suite** for tampering, rollback loops, long-prompt DoS, power-cycling, and corrupted artifacts. The benchmarking literature now strongly suggests that narrow in-distribution safety tests overstate real robustness. + +The core metrics should include: **attack success rate**, **harm score or refusal-quality score**, **benign task success**, **over-refusal rate**, **TTFT**, **prefill tok/s**, **decode tok/s**, **p95 latency**, **peak RSS/VRAM**, **J/token or Wh/request**, **rollback rate**, **mean rollback depth**, **mean discarded tokens**, **tamper-detection rate**, and **secure-boot / verity / IMA pass rate**. If you use a judge model in evaluation, report **judge sensitivity analyses** and at least one non-judge or multi-judge cross-check, because safety judges alone can be fooled. + +## Paper and project landscape + +### Comparative paper table + +| Title | Authors | Year | Venue | Key contributions | +|---|---|---:|---|---| +| **DExperts: Decoding-Time Controlled Text Generation with Experts and Anti-Experts** | Liu, Sap, Lu, Swayamdipta, Bhagavatula, Smith, Choi | 2021 | ACL | Foundational product-of-experts framework for decoding-time control using smaller expert/anti-expert models. | +| **Reliable code generation from pre-trained language models** / **Synchromesh** with **Constrained Semantic Decoding** | Poesia, Dong, Zettlemoyer, Ericson, Scales, Polozov | 2022 | ICLR | Introduces CSD and completion engines; guarantees constraint satisfaction by construction when the completion engine is exhaustive and correct. | +| **SafeDecoding: Defending against Jailbreak Attacks via Safety-Aware Decoding** | Xu, Jiang, Niu, Jia, Lin, Poovendran | 2024 | ACL | Introduces safety-aware decoding using original/expert token intersections and reweighting; strong empirical jailbreak defense with small inference overhead. | +| **InferAligner: Inference-Time Alignment for Harmlessness through Cross-Model Guidance** | Wang, Zhang, Li, Tan, Wang, Zhang, Ren, Jiang, Qiu | 2024 | EMNLP | Uses cross-model safety steering vectors to modify activations at inference time for harmlessness alignment. | +| **AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration** | Lin et al. | 2024 | MLSys | Hardware-friendly low-bit weight quantization oriented toward on-device deployment; central for small edge models. | +| **JailbreakBench: An Open Robustness Benchmark for Jailbreaking Large Language Models** | Chao et al. | 2024 | NeurIPS Datasets and Benchmarks | Standardizes jailbreak artifacts, threat models, prompts, and scoring, making defense comparisons more credible. | +| **SafeInfer: Context Adaptive Decoding Time Safety Alignment for Large Language Models** | Banerjee, Layek, Tripathy, Kumar, Mukherjee, Hazra | 2025 | AAAI | Couples activation-level safety amplification with safety-guided decoding; introduces HarmEval. | +| **ROCODE: Integrating Backtracking Mechanism and Program Analysis in Large Language Models for Code Generation** | Jiang, Dong, Tao, Liu, Jin, Jiao, Li | 2025 | ICSE | Incremental error detection with rollback and constrained regeneration; strong evidence that rollback lowers error accumulation and token cost. | +| **Controllable Safety Alignment: Inference-Time Adaptation to Diverse Safety Requirements** | Zhang, Elgohary, Magooda, Khashabi, Van Durme | 2025 | ICLR | Introduces natural-language “safety configs” and inference-time adaptation to varying safety requirements. | +| **SecDecoding: Steerable Decoding for Safer LLM Generation** | Wang, Liu, Hu, Wu, He | 2025 | Findings of EMNLP | Uses small contrastive models to produce a divergence-based token-level safety signal; compatible with speculative decoding. | +| **Know Thy Judge: On the Robustness Meta-Evaluation of LLM Safety Judges** | Eiras, Zemour, Lin, Mugunthan | 2025 | ICLR workshop / PMLR | Shows that safety judges are brittle to style changes and adversarial outputs; essential caution for evaluation design. | +| **Hydra: Efficient, Correct Code Generation via Checkpoint-and-Rollback Support** | Du, Ou, Zhuo, Lentz | 2026 | arXiv preprint | Asynchronous checking plus checkpoint-and-rollback; reports large latency and token-consumption reductions relative to post-hoc repair. | +| **On Almost Surely Safe Alignment of Large Language Models at Inference-Time** | Ji, Ramesh, Zimmer, Bogunovic, Wang, Bou Ammar | 2025 | arXiv / TMLR under review | Formalizes inference-time safety as a constrained MDP with a cost-model-relative guarantee; introduces InferenceGuard. | + +### Open-source projects and implementations worth tracking + +| Project | What it provides | Why it matters here | +|---|---|---| +| **uw-nsl/SafeDecoding** | Official implementation of the ACL 2024 SafeDecoding paper | Best starting codebase for lightweight model-backed decoding-time safety. | +| **kanishkg/synchromesh** | CSD / completion-engine implementation | Reference design for hard-validity constrained decoding. | +| **microsoft/controllable-safety-alignment** | Codebase for CoSA | Useful if offline deployments need configurable safety policies. | +| **jiangxxxue/ROCODE** | Runtime backtracking with program analysis | Good concrete reference for rollback orchestration patterns. | +| **llama.cpp** | Portable local inference runtime with GGUF, mmap, many backends | Most pragmatic Linux/SBC starting point. | +| **onnxruntime-genai** | On-device generation loop with KV cache management and logits processing | Strong cross-platform runtime with extensible generation internals. | +| **LiteRT-LM** | Production-ready cross-platform edge LLM runtime | Especially attractive for Android/iOS/WebGPU deployments. | +| **openvino.genai** | Pre-built GenAI pipelines on OpenVINO Runtime | Best fit for Intel edge systems with CPU/GPU/NPU options. | +| **TensorRT-LLM** | NVIDIA-optimized LLM runtime with quantization and KV-cache facilities | Best fit for Jetson and NVIDIA edge GPUs. | +| **ExecuTorch** | PyTorch-native on-device inference across mobile and embedded targets | Best if the surrounding organization is already PyTorch-first. | +| **Keylime** | TPM-based attestation and runtime integrity monitoring | The right integrity/attestation companion for high-assurance deployments. | + +### Open questions and limitations + +The research base is strong on **decoding-time steering** and increasingly strong on **rollback**, but there are still real gaps. + +First, there is no mature, widely adopted open-source implementation that combines **SecDecoding-style steering**, **CSD-style structured constraints**, and **Hydra/RoCode-style rollback** in one edge-optimized stack. Today, a production system usually has to compose these pieces manually. + +Second, the strongest formal result in this space—**InferenceGuard**—is tied to a specified safety cost model, which is valuable, but still far from a blanket guarantee against arbitrary semantic harms. That is not a criticism of the paper; it is just the present state of the field. + +Third, **LLM judges remain fragile**, which means many published “safety gains” should be read as gains under specific evaluators rather than final proof of robustness. This is especially important for air-gapped devices, where a false sense of security is dangerous because operators often cannot patch or re-tune quickly in the field. + +Fourth, for **hard real-time** control systems, even 0.5B models remain poor candidates for direct control loops. They are much better used as **advisory**, **translation**, **summarization**, or **configuration-authoring** systems whose machine-facing outputs are then filtered through deterministic validators or approval steps. That recommendation follows from the variable latency and stochasticity documented by current on-device deployment stacks and benchmarks rather than from any single paper. + +The practical bottom line is therefore clear: for an offline, air-gapped edge device running a **~0.5B open-weight model**, the best current design is a **quantized small model + selective model-backed safety steering + checkpointed rollback + structured-output constraints + cryptographically verified platform integrity**. That stack is feasible today, gives a credible defense-in-depth story, and maps well to the strongest evidence currently available.