//! Request runtime: continuous-batching scheduler + radix prefix cache + the answer policy (no-think readout, tau gate, thought, post-thought readout),
//! on top of a `ModelRunner`.
//!
//! One engine thread owns the model. `Engine::run` (any thread) preprocesses images, expands the prompts, enqueues the request and blocks on the reply.
//! Per iteration: (1) admit waiting requests (earliest deadline first, page / branch / row budgets) and run ONE packed prefill forward for the new work of
//! all of them (shared prefix rows once per request + per-branch tails; cache hits skip rows except the last two of every branch), read the readout row
//! of every branch from it; (2) one decode step for ALL thinking branches of ALL requests (`ThinkDecoder`); (3) one readout forward for the branches that
//! finished their thought (end-of-thought row = post-thought letter distribution); (4) one decode cycle for every answering number question
//! (`ThinkDecoder::answer_step`).
//! Scheduling ideas follow SGLang `srt/managers/schedule_policy.py` (cache-aware ordering, in-batch prefix de-duplication by deferral) and
//! `srt/managers/scheduler.py` (iteration-level batching); the prefix tree is `radix.rs`.
use crate::model::*;
use crate::radix::{page_keys, NodeId, PKey, RadixCache, ROOT};
use crate::think::{GreedyDecoder, ThinkDecoder, ThinkReq};
use anyhow::{anyhow, bail, Result};
use dllm_core::policy::{GateMode, Policy, ThinkShape};
use dllm_core::serving::{Item, ItemResult};
use dllm_core::tokenizer::{Special, Tok};
use dllm_model::gpu_model::{RowBatch, PAGE};
use dllm_model::spec::greedy_logprob;
use std::collections::{HashMap, VecDeque};
use std::sync::mpsc::{channel, Receiver, Sender};
use std::sync::{Arc, Mutex};
use std::time::Instant;

pub const MK: i32 = 101;
pub const EOS: i32 = 1;
pub const TURN_END: i32 = 106;
/// longest generated answer of a number question (tokens, without the closing end-of-turn)
pub const ANSWER_MAX: usize = 32;
/// a request whose remaining wall-clock budget (minus the readout reserve) is below this does not start a thought (ms)
const MIN_THINK_MS: f64 = 60.0;

#[derive(Clone)]
pub struct EngineConfig {
    /// model context limit (thoughts are clamped to the room left in it)
    pub max_len: usize,
    /// max concurrent branches (prefilled or thinking)
    pub max_branches: usize,
    /// soft limit of new rows per prefill iteration (the model's max_rows is the hard limit of one forward; larger work is chunked)
    pub max_prefill_rows: usize,
    pub prefix_cache: bool,
    /// image preprocessing (None => requests with images fail)
    pub prep: Option<Arc<dyn ImagePrep>>,
    /// encoded images kept on the device (by content hash)
    pub soft_cache_images: usize,
    /// forward row cap while any branch is thinking (0 = model max_rows): keeps decode steps / readouts responsive; the model max_rows applies otherwise
    pub think_max_rows: usize,
}
impl Default for EngineConfig {
    fn default() -> Self {
        EngineConfig { max_len: 32768, max_branches: 256, max_prefill_rows: 2048, prefix_cache: true, prep: None, soft_cache_images: 8, think_max_rows: 1024 }
    }
}

/// counters of the scheduler (tests and diagnostics)
#[derive(Debug, Clone, Default)]
pub struct Stats {
    pub requests: u64,
    pub items: u64,
    pub think_items: u64,
    pub errors: u64,
    pub prefill_forwards: u64,
    pub prefill_rows: u64,
    pub decode_forwards: u64,
    pub probe_forwards: u64,
    pub cache_query_pages: u64,
    pub cache_hit_pages: u64,
    pub cache_evicted_pages: u64,
    pub deferred_in_batch: u64,
    pub vision_encodes: u64,
    pub soft_cache_hits: u64,
    pub peak_branches: u64,
    pub free_pages: usize,
    pub tree_pages: usize,
    pub num_pages: usize,
    /// sliding-window pool (0 pages = model without sliding layers): pool size, free pages, pages kept for cached prefixes
    pub sliding_pages: usize,
    pub sliding_free: usize,
    pub sliding_tree: usize,
    pub forced_by_budget: u64,
    /// thinking skipped because the budget left no useful time (or the deterministic cap was 0)
    pub skipped_think: u64,
    /// thinking cut because the sliding-window pool had no page for the next step
    pub swa_forced: u64,
    pub thought_tokens: u64,
}
impl Stats {
    pub fn hit_rate(&self) -> f64 {
        if self.cache_query_pages == 0 { 0.0 } else { self.cache_hit_pages as f64 / self.cache_query_pages as f64 }
    }
}

#[derive(Debug, Clone, Copy)]
struct Blk {
    s: usize,
    e: usize,
    img: usize,
}

struct ReqIn {
    items: Vec<Item>,
    ids: Vec<Vec<i32>>,
    keys: Vec<Vec<u64>>,
    blocks: Vec<Vec<Blk>>,
    images: Vec<PreparedImage>,
    pol: Policy,
    reply: Sender<Result<Vec<ItemResult>>>,
    /// request arrival (HTTP layer if it called `set_arrival`, else the call to `run`): the think budget counts from here
    arrive: Instant,
    /// request shape for the deterministic `--think-ms` cap (`thinking` is filled in after the gate)
    shape: ThinkShape,
}

enum Msg {
    Req(ReqIn),
    Shutdown,
}

pub struct Engine {
    tx: Mutex<Sender<Msg>>,
    stats: Arc<Mutex<Stats>>,
    sp: Special,
    prep: Option<Arc<dyn ImagePrep>>,
    has_vision: bool,
    join: Mutex<Option<std::thread::JoinHandle<()>>>,
    max_len: usize,
}

fn fnv_mix(a: u64, b: u64) -> u64 {
    let mut z = (a ^ b.wrapping_mul(0x9E3779B97F4A7C15)).wrapping_add(0x9E3779B97F4A7C15);
    z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
    z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
    z ^ (z >> 31)
}

impl Engine {
    pub fn start(cfg: EngineConfig, model: Box<dyn ModelRunner>, vision: Option<Box<dyn VisionEncoder>>, tok: Arc<Tok>) -> Result<Engine> {
        Self::start_with(cfg, model, vision, tok, Box::new(GreedyDecoder))
    }

    pub fn start_with(cfg: EngineConfig, mut model: Box<dyn ModelRunner>, vision: Option<Box<dyn VisionEncoder>>, tok: Arc<Tok>, think: Box<dyn ThinkDecoder>) -> Result<Engine> {
        if model.num_pages() < 4 {
            bail!("model has too few KV pages");
        }
        let t = Instant::now();
        match model.enable_graphs() {
            Ok(true) => eprintln!("decode CUDA graphs captured in {:.2}s", t.elapsed().as_secs_f32()),
            Ok(false) => {}
            Err(e) => eprintln!("decode CUDA graphs disabled: {e:#}"),
        }
        let (tx, rx) = channel();
        let stats = Arc::new(Mutex::new(Stats::default()));
        let has_vision = vision.is_some();
        let prep = cfg.prep.clone();
        let max_len = cfg.max_len;
        let mut sch = Sched::new(cfg, model, vision, think, stats.clone());
        let join = std::thread::Builder::new().name("dllm-engine".into()).spawn(move || sch.run_loop(rx))?;
        Ok(Engine { tx: Mutex::new(tx), stats, sp: tok.sp, prep, has_vision, join: Mutex::new(Some(join)), max_len })
    }

    pub fn stats(&self) -> Stats {
        self.stats.lock().unwrap().clone()
    }

    /// Blocking; callable concurrently from many threads. All questions of one call share the state prefix and the images.
    pub fn run(&self, items: &[Item], images: &[Vec<u8>], pol: &Policy) -> Result<Vec<ItemResult>> {
        if items.is_empty() {
            return Ok(vec![]);
        }
        let arrive = dllm_core::serving::take_arrival().unwrap_or_else(Instant::now);
        let mut prepared = vec![];
        if !images.is_empty() {
            if !self.has_vision {
                bail!("images are not supported by this engine (no vision encoder)");
            }
            let prep = self.prep.as_ref().ok_or_else(|| anyhow!("no image preprocessor configured"))?;
            // images of one request are prepared (decode + resize + tower) concurrently; the tower serialises internally
            let one = |i: usize, b: &Vec<u8>| -> Result<PreparedImage> { prep.prepare(b).map_err(|e| anyhow!("cannot decode image {i}: {e:#}")) };
            if images.len() == 1 {
                prepared.push(one(0, &images[0])?);
            } else {
                let rs: Vec<Result<PreparedImage>> = std::thread::scope(|sc| {
                    let hs: Vec<_> = images.iter().enumerate().map(|(i, b)| sc.spawn(move || one(i, b))).collect();
                    hs.into_iter().map(|h| h.join().unwrap_or_else(|_| Err(anyhow!("image preparation panicked")))).collect()
                });
                for r in rs {
                    prepared.push(r?);
                }
            }
        }
        let soft: Vec<usize> = prepared.iter().map(|p| p.n_soft).collect();
        let tok_sp = self.sp;
        let mut ids = vec![];
        let mut keys = vec![];
        let mut blocks = vec![];
        for it in items {
            let n_ph = it.ids.iter().filter(|&&t| t == tok_sp.image).count();
            if n_ph != images.len() {
                bail!("item {}: {} image placeholders for {} images", it.key, n_ph, images.len());
            }
            if it.ids.is_empty() {
                bail!("item {}: empty prompt", it.key);
            }
            let mut ex: Vec<i32> = Vec::with_capacity(it.ids.len() + soft.iter().sum::<usize>() + 2 * soft.len());
            let mut k = 0;
            for &t in &it.ids {
                if t == tok_sp.image {
                    ex.push(tok_sp.boi as i32);
                    ex.extend(std::iter::repeat(tok_sp.image as i32).take(soft[k]));
                    ex.push(tok_sp.eoi as i32);
                    k += 1;
                } else {
                    ex.push(t as i32);
                }
            }
            // keys + blocks (image soft tokens carry the image content hash and their offset: traps (a)/(e))
            let mut kv = Vec::with_capacity(ex.len());
            let mut bl: Vec<Blk> = vec![];
            let mut i = 0;
            while i < ex.len() {
                if ex[i] == tok_sp.image as i32 {
                    let img = bl.len();
                    let s = i;
                    while i < ex.len() && ex[i] == tok_sp.image as i32 {
                        let j = (i - s) as u64;
                        kv.push((1u64 << 63) | (fnv_mix(prepared[img].hash ^ ((prepared[img].n_soft as u64) << 40), j + 1) >> 1));
                        i += 1;
                    }
                    bl.push(Blk { s, e: i, img });
                } else {
                    kv.push(ex[i] as u32 as u64);
                    i += 1;
                }
            }
            ids.push(ex);
            keys.push(kv);
            blocks.push(bl);
        }
        let (rtx, rrx) = channel();
        let shape = ThinkShape { prefill_tokens: prefill_rows(&ids), images: images.len(), image_bytes: images.iter().map(|b| b.len()).sum(), questions: items.len(), thinking: 0 };
        let req = ReqIn { items: items.to_vec(), ids, keys, blocks, images: prepared, pol: *pol, reply: rtx, arrive, shape };
        self.tx.lock().unwrap().send(Msg::Req(req)).map_err(|_| anyhow!("engine thread has stopped"))?;
        rrx.recv().map_err(|_| anyhow!("engine thread died"))?
    }

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

impl Drop for Engine {
    fn drop(&mut self) {
        let _ = self.tx.lock().unwrap().send(Msg::Shutdown);
        if let Some(j) = self.join.lock().unwrap().take() {
            let _ = j.join();
        }
    }
}

// ================================================================================================================ scheduler

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum Phase {
    Prefill,
    /// `--think-ms`: the gate fired; waits for the rest of the request's gate decisions, then gets the request's deterministic cap and thinks
    Gated,
    Think,
    /// feed `feed` (starting at `probe_pos`) and read the letter logits of the last row (number questions: the first answer token)
    Probe,
    /// number question: greedy answer tokens after the end-of-thought token (`answer_step`)
    Answer,
}

struct Branch {
    req: u64,
    item: usize,
    /// number question (generated answer instead of a letter readout)
    number: bool,
    /// number answer: tokens so far, sum of their log-probabilities, position of `ans[0]`
    ans: Vec<i32>,
    ans_lp: f64,
    ans_pos: usize,
    l: usize,
    pol: Policy,
    pages: Vec<i32>,
    node: NodeId,
    phase: Phase,
    p0: Vec<f32>,
    gi: Vec<i32>,
    fed: usize,
    maxgen: usize,
    /// position of gi[0]
    base: usize,
    feed: Vec<i32>,
    probe_pos: usize,
    nt: usize,
    out_tok: usize,
    deadline: Option<Instant>,
    /// sliding-window pool: sliding ref held on page index j (`sl[j]`), trim cursor, shared-prefix pages of the request
    sl: Vec<bool>,
    sl_cur: usize,
    sp_pages: usize,
}

struct Req {
    inp: ReqIn,
    results: Vec<Option<ItemResult>>,
    remaining: usize,
    failed: bool,
    arrive: Instant,
    /// sliding pages reserved for this request's prefill (recycled within the request), and the reserve's size cap
    sl_res: Vec<i32>,
    sl_cap: usize,
}

struct Seg {
    tokens: Vec<i32>,
    start: usize,
    pages: Vec<i32>,
    blocks: Vec<Blk>,
    req: u64,
    branch: Option<usize>,
}

struct Pool {
    free: Vec<i32>,
    rc: Vec<u32>,
}
impl Pool {
    fn decref(&mut self, p: i32) {
        let r = &mut self.rc[p as usize];
        assert!(*r > 0, "decref of free page {p}");
        *r -= 1;
        if *r == 0 {
            self.free.push(p);
        }
    }
    fn incref(&mut self, p: i32) {
        assert!(self.rc[p as usize] > 0, "incref of free page {p}");
        self.rc[p as usize] += 1;
    }
}

struct Sched {
    cfg: EngineConfig,
    model: Box<dyn ModelRunner>,
    vision: Option<Box<dyn VisionEncoder>>,
    think: Box<dyn ThinkDecoder>,
    pool: Pool,
    tree: RadixCache,
    reqs: HashMap<u64, Req>,
    waiting: VecDeque<u64>,
    branches: Vec<Option<Branch>>,
    free_slots: Vec<usize>,
    live_branches: usize,
    next_id: u64,
    soft_cache: Vec<(u64, Arc<SoftTokens>)>,
    st: Stats,
    shared: Arc<Mutex<Stats>>,
    job: Option<PrefillJob>,
    ema_probe_ms: f64,
    /// EMA of one prefill forward (ms), used to cut thinking branches BEFORE a forward that would overrun their deadline
    ema_pf_fwd_ms: f64,
    ema_dec_ms: f64,
    /// vision-tower encode time per image (the next forward's estimate includes the encodes it will trigger)
    ema_vis_ms: f64,
    ema_pf_ms_row: f64,
    /// decayed least-squares sums of (rows, ms) of prefill forwards: ms ~ a + b * rows
    pf_fit: [f64; 5],
    /// sliding-window KV pool bookkeeping (None = the model has no separate sliding-window pool)
    swa: Option<Swa>,
}

struct PrefillJob {
    segs: Vec<Seg>,
    chunks: Vec<Vec<(usize, usize, usize)>>,
    next: usize,
    rids: Vec<u64>,
}

fn is_stop(t: i32) -> bool {
    t == MK || t == EOS || t == TURN_END
}

fn softmax64(z: &[f32]) -> Vec<f64> {
    let m = z.iter().cloned().fold(f32::NEG_INFINITY, f32::max) as f64;
    let e: Vec<f64> = z.iter().map(|&v| (v as f64 - m).exp()).collect();
    let s: f64 = e.iter().sum();
    e.iter().map(|v| v / s).collect()
}

impl Sched {
    fn new(cfg: EngineConfig, model: Box<dyn ModelRunner>, vision: Option<Box<dyn VisionEncoder>>, think: Box<dyn ThinkDecoder>, shared: Arc<Mutex<Stats>>) -> Self {
        let n = model.num_pages();
        let swa = match (model.sliding_pages(), model.sliding_window()) {
            (s, w) if s > 0 && w > 0 => Some(Swa::new(n, s, w)),
            _ => None,
        };
        let st = Stats { num_pages: n, sliding_pages: model.sliding_pages(), ..Default::default() };
        Sched {
            cfg,
            model,
            vision,
            think,
            pool: Pool { free: (0..n as i32).rev().collect(), rc: vec![0; n] },
            tree: RadixCache::new(),
            reqs: HashMap::new(),
            waiting: VecDeque::new(),
            branches: vec![],
            free_slots: vec![],
            live_branches: 0,
            next_id: 1,
            soft_cache: vec![],
            st,
            shared,
            job: None,
            ema_probe_ms: 3.0,
            ema_pf_fwd_ms: 10.0,
            ema_dec_ms: 3.0,
            ema_vis_ms: 10.0,
            ema_pf_ms_row: 0.02,
            pf_fit: [0.0; 5],
            swa,
        }
    }

    fn run_loop(&mut self, rx: Receiver<Msg>) {
        let mut shutdown = false;
        loop {
            let idle = self.waiting.is_empty() && self.live_branches == 0;
            if idle {
                if cfg!(debug_assertions) {
                    self.check_idle();
                }
                if shutdown {
                    return;
                }
                match rx.recv() {
                    Ok(Msg::Req(r)) => self.enqueue(r),
                    Ok(Msg::Shutdown) | Err(_) => return,
                }
            }
            while let Ok(m) = rx.try_recv() {
                match m {
                    Msg::Req(r) => self.enqueue(r),
                    Msg::Shutdown => shutdown = true,
                }
            }
            // load-aware budget accounting: if the next prefill forward would run past (deadline - reserve) of a thinking branch, cut and read that branch out FIRST
            let est = self.next_prefill_est_ms();
            if est > 0.0 && self.cut_due(est) {
                let tp0 = Instant::now();
                if self.probe_step() {
                    self.ema_probe_ms = 0.9 * self.ema_probe_ms + 0.1 * tp0.elapsed().as_secs_f64() * 1e3;
                }
            }
            self.prefill_step();
            self.decode_step();
            let tp = Instant::now();
            let probed = self.probe_step();
            if probed {
                self.ema_probe_ms = 0.9 * self.ema_probe_ms + 0.1 * tp.elapsed().as_secs_f64() * 1e3;
            }
            self.answer_step();
            self.publish();
            if shutdown && self.waiting.is_empty() && self.live_branches == 0 {
                return;
            }
        }
    }

    /// measured time needed to finish a cut request: one readout (probe) pass + one decode step of slack
    /// prefill cost model ms = a + b*rows (fixed expert-streaming floor + per-row slope), least squares over recent forwards
    fn pf_model(&self) -> (f64, f64) {
        let [n, sx, sy, sxx, sxy] = self.pf_fit;
        let var = n * sxx - sx * sx;
        if n < 6.0 || var < 1e-6 * n * n {
            return (8.0, self.ema_pf_ms_row.min(0.1));
        }
        let b = ((n * sxy - sx * sy) / var).clamp(0.005, 0.3);
        let a = ((sy - b * sx) / n).clamp(2.0, 40.0);
        (a, b)
    }

    /// expected duration of the prefill forward the scheduler is about to run (0 = none pending)
    fn next_prefill_est_ms(&self) -> f64 {
        let (a, b) = self.pf_model();
        let cap = if self.thinking_live() { self.think_cap() } else { usize::MAX };
        if let Some(j) = &self.job {
            if j.next < j.chunks.len() {
                let rows: usize = j.chunks[j.next].iter().map(|c| c.2 - c.1).sum();
                return (a + b * rows.min(cap) as f64 + self.uncached_images(&j.segs, &j.chunks[j.next]) as f64 * self.ema_vis_ms).min(400.0);
            }
        }
        if !self.waiting.is_empty() {
            return self.ema_pf_fwd_ms.min(a + b * cap.min(self.model.max_rows()) as f64).min(400.0);
        }
        0.0
    }

    /// images of a chunk whose soft tokens are not in the device cache (the chunk will run the vision tower for them)
    fn uncached_images(&self, segs: &[Seg], ch: &[(usize, usize, usize)]) -> usize {
        let mut seen: Vec<u64> = vec![];
        for &(si, lo, hi) in ch {
            let sg = &segs[si];
            let Some(r) = self.reqs.get(&sg.req) else { continue };
            for bl in sg.blocks.iter().filter(|bl| bl.s < sg.start + hi && bl.e > sg.start + lo) {
                let p = &r.inp.images[bl.img];
                let h = p.hash ^ ((p.n_soft as u64) << 48);
                if !seen.contains(&h) && !self.soft_cache.iter().any(|(x, _)| *x == h) {
                    seen.push(h);
                }
            }
        }
        seen.len()
    }

    /// forward row cap while thinking is live (`think_max_rows`, 0 = the model's max_rows)
    fn think_cap(&self) -> usize {
        if self.cfg.think_max_rows == 0 { self.model.max_rows() } else { self.cfg.think_max_rows }
    }

    /// force the readout of thinking branches whose deadline is within reserve + `extra_ms`; returns whether any was cut
    fn cut_due(&mut self, extra_ms: f64) -> bool {
        let now = Instant::now() + std::time::Duration::from_secs_f64((self.reserve_ms() + extra_ms) / 1e3);
        let late: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.phase == Phase::Think && b.deadline.map_or(false, |d| now >= d))).collect();
        for &bi in &late {
            self.st.forced_by_budget += 1;
            self.end_think(bi);
        }
        !late.is_empty()
    }

    fn reserve_ms(&self) -> f64 {
        (1.5 * self.ema_probe_ms + self.ema_dec_ms).max(2.0)
    }

    fn publish(&mut self) {
        self.st.free_pages = self.pool.free.len();
        if let Some(sw) = &self.swa {
            self.st.sliding_free = sw.free.len() + self.reqs.values().map(|r| r.sl_res.len()).sum::<usize>();
            self.st.sliding_tree = sw.tree_n;
        }
        self.st.tree_pages = self.tree.total_pages();
        *self.shared.lock().unwrap() = self.st.clone();
    }

    fn enqueue(&mut self, inp: ReqIn) {
        let id = self.next_id;
        self.next_id += 1;
        let n = inp.items.len();
        self.st.requests += 1;
        self.st.items += n as u64;
        let arrive = inp.arrive;
        self.reqs.insert(id, Req { inp, results: vec![None; n], remaining: n, failed: false, arrive, sl_res: vec![], sl_cap: 0 });
        self.waiting.push_back(id);
    }

    // -------------------------------------------------------------------------------------------- pages
    fn alloc(&mut self, n: usize) -> Result<Vec<i32>> {
        if self.pool.free.len() < n {
            let need = n - self.pool.free.len();
            self.evict_tree(need);
        }
        if self.pool.free.len() < n {
            bail!("KV pool exhausted ({} pages total, {} needed, {} free)", self.pool.rc.len(), n, self.pool.free.len());
        }
        let mut v = Vec::with_capacity(n);
        for _ in 0..n {
            let p = self.pool.free.pop().unwrap();
            self.pool.rc[p as usize] = 1;
            v.push(p);
        }
        Ok(v)
    }

    fn release_branch(&mut self, bi: usize) {
        if self.swa_on() {
            let held: Vec<i32> = {
                let b = self.branches[bi].as_ref().unwrap();
                b.sl.iter().enumerate().filter(|(_, &h)| h).map(|(j, _)| b.pages[j]).collect()
            };
            for f in held {
                self.sl_unref(f);
            }
            self.branches[bi].as_mut().unwrap().sl.clear();
        }
        if let Some(b) = self.branches[bi].take() {
            for &p in &b.pages {
                self.pool.decref(p);
            }
            if b.node != ROOT {
                self.tree.dec_lock(b.node);
            }
            self.free_slots.push(bi);
            self.live_branches -= 1;
        }
    }

    fn new_branch(&mut self, b: Branch) -> usize {
        self.live_branches += 1;
        self.st.peak_branches = self.st.peak_branches.max(self.live_branches as u64);
        if let Some(i) = self.free_slots.pop() {
            self.branches[i] = Some(b);
            i
        } else {
            self.branches.push(Some(b));
            self.branches.len() - 1
        }
    }

    // -------------------------------------------------------------------------------------------- completion
    fn finish_item(&mut self, bi: usize, res: ItemResult) {
        let (rid, item) = {
            let b = self.branches[bi].as_ref().unwrap();
            (b.req, b.item)
        };
        self.release_branch(bi);
        if let Some(r) = self.reqs.get_mut(&rid) {
            if r.failed {
                return;
            }
            r.results[item] = Some(res);
            r.remaining -= 1;
            if r.remaining == 0 {
                let mut r = self.reqs.remove(&rid).unwrap();
                self.sl_drop_reserve(std::mem::take(&mut r.sl_res));
                self.publish();
                let out: Vec<ItemResult> = r.results.into_iter().map(|x| x.unwrap()).collect();
                let _ = r.inp.reply.send(Ok(out));
            }
        }
    }

    fn fail_req(&mut self, rid: u64, msg: String) {
        self.waiting.retain(|&x| x != rid);
        let bis: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.req == rid)).collect();
        for bi in bis {
            self.release_branch(bi);
        }
        if let Some(mut r) = self.reqs.remove(&rid) {
            self.sl_drop_reserve(std::mem::take(&mut r.sl_res));
            self.st.errors += 1;
            let _ = r.inp.reply.send(Err(anyhow!(msg)));
        }
    }

    fn fail_branches(&mut self, bis: &[usize], msg: &str) {
        let mut rids: Vec<u64> = bis.iter().filter_map(|&i| self.branches[i].as_ref().map(|b| b.req)).collect();
        rids.sort();
        rids.dedup();
        for r in rids {
            self.fail_req(r, msg.to_string());
        }
    }

    // -------------------------------------------------------------------------------------------- admission + prefill
    /// One scheduler iteration of prefill: admit (EDF) while the pending job is below `max_prefill_rows`, then run ONE forward of the job (every forward
    /// streams all experts, so big batches are cheap per row); the rest of the job continues next iteration, after a decode step of the thinking branches,
    /// so a long prefill never stalls them.
    fn prefill_step(&mut self) {
        if !self.waiting.is_empty() {
            let pending: usize = self.job.as_ref().map_or(0, |j| j.chunks[j.next..].iter().map(|c| c.iter().map(|p| p.2 - p.1).sum::<usize>()).sum());
            if pending < self.cfg.max_prefill_rows {
                self.admit(pending);
            }
        }
        let Some(mut job) = self.job.take() else { return };
        let mut spent = 0.0f64;
        let mut failed: Option<String> = None;
        while job.next < job.chunks.len() {
            let ch: Vec<(usize, usize, usize)> = job.chunks[job.next].iter().copied().filter(|c| self.reqs.contains_key(&job.segs[c.0].req)).collect();
            if ch.is_empty() {
                job.next += 1;
                continue;
            }
            let mut ch = ch;
            let mut rows: usize = ch.iter().map(|c| c.2 - c.1).sum();
            // while branches think, forwards stay at `think_max_rows` (a big no-think pack would delay their decode steps / readouts)
            let tcap = self.think_cap();
            if tcap > 0 && rows > tcap && self.thinking_live() {
                let parts = self.split_chunk(&job.segs, &ch, tcap);
                if parts.len() > 1 {
                    ch = parts[0].clone();
                    rows = ch.iter().map(|c| c.2 - c.1).sum();
                    job.chunks.splice(job.next + 1..job.next + 1, parts[1..].to_vec());
                    job.chunks[job.next] = ch.clone();
                }
            }
            // exactly one forward per iteration
            if spent > 0.0 {
                break;
            }
            let t0 = Instant::now();
            if let Err(e) = self.run_chunk(&job.segs, &ch) {
                failed = Some(format!("{e:#}"));
                break;
            }
            let el = t0.elapsed().as_secs_f64() * 1e3;
            self.ema_pf_fwd_ms = 0.9 * self.ema_pf_fwd_ms + 0.1 * el;
            spent += el;
            job.next += 1;
        }
        if let Some(e) = failed {
            for r in job.rids.clone() {
                self.fail_req(r, format!("prefill failed: {e}"));
            }
        } else if job.next < job.chunks.len() {
            self.job = Some(job);
        }
    }

    /// chunk an ordered list of pieces (seg, lo, hi) into forwards of at most `cap` rows (same cut rules as `build_chunks`)
    fn chunk_pieces(&self, segs: &[Seg], ch: &[(usize, usize, usize)], cap: usize) -> Result<Vec<Vec<(usize, usize, usize)>>> {
        let tmp: Vec<Seg> = ch
            .iter()
            .map(|&(si, lo, hi)| {
                let s = &segs[si];
                Seg { tokens: s.tokens[lo..hi].to_vec(), start: s.start + lo, pages: vec![], blocks: s.blocks.clone(), req: s.req, branch: if hi == s.tokens.len() { s.branch } else { None } }
            })
            .collect();
        Ok(self.build_chunks(&tmp, cap)?.into_iter().map(|p| p.into_iter().map(|(ti, lo2, hi2)| (ch[ti].0, ch[ti].1 + lo2, ch[ti].1 + hi2)).collect()).collect())
    }

    /// any branch thinking / waiting for its readout (prefill forwards are then kept at `think_max_rows` so decode steps are not starved)
    fn thinking_live(&self) -> bool {
        self.branches.iter().flatten().any(|b| matches!(b.phase, Phase::Think | Phase::Probe | Phase::Answer))
    }

    /// re-chunk the pieces of one chunk with a smaller row cap (same cut rules as `build_chunks`)
    fn split_chunk(&self, segs: &[Seg], ch: &[(usize, usize, usize)], cap: usize) -> Vec<Vec<(usize, usize, usize)>> {
        let tmp: Vec<Seg> = ch
            .iter()
            .map(|&(si, lo, hi)| {
                let s = &segs[si];
                Seg { tokens: s.tokens[lo..hi].to_vec(), start: s.start + lo, pages: vec![], blocks: s.blocks.clone(), req: s.req, branch: if hi == s.tokens.len() { s.branch } else { None } }
            })
            .collect();
        match self.build_chunks(&tmp, cap) {
            Ok(parts) => parts.into_iter().map(|p| p.into_iter().map(|(ti, lo2, hi2)| (ch[ti].0, ch[ti].1 + lo2, ch[ti].1 + hi2)).collect()).collect(),
            Err(_) => vec![ch.to_vec()],
        }
    }

    fn admit(&mut self, pending: usize) {
        // earliest deadline first (arrival + budget); a cached page is worth 0.5 ms of urgency (cache hits go slightly earlier, deadlines dominate).
        // (A shortest-job-first variant lowered nothing measurable at c=16 and raised p99: not used.)
        let mut order: Vec<(f64, u64)> = vec![];
        for &rid in &self.waiting {
            let r = &self.reqs[&rid];
            let b = r.inp.pol.budget_ms as u64;
            let dl = (if b == 0 { 60_000 } else { b }) as f64 - r.arrive.elapsed().as_secs_f64() * 1e3;
            let k = &r.inp.keys[0];
            let c = (k.len() - 1) / PAGE;
            let m = if self.cfg.prefix_cache { self.tree.match_len(&page_keys(&k[..c * PAGE])) } else { 0 };
            order.push((dl - 0.5 * m as f64, rid));
        }
        order.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let mut segs: Vec<Seg> = vec![];
        let mut inflight: Vec<Vec<PKey>> = vec![];
        let mut rows = 0usize;
        let mut admitted: Vec<u64> = vec![];
        for (_, rid) in order {
            if pending + rows >= self.cfg.max_prefill_rows && !segs.is_empty() {
                break;
            }
            let n = self.reqs[&rid].inp.items.len();
            if self.live_branches > 0 && self.live_branches + n > self.cfg.max_branches {
                continue;
            }
            match self.plan_request(rid, &mut segs, &mut inflight) {
                Ok(Some(r)) => {
                    rows += r;
                    admitted.push(rid);
                    self.waiting.retain(|&x| x != rid);
                }
                Ok(None) => {}
                Err(e) => self.fail_req(rid, format!("{e:#}")),
            }
        }
        if segs.is_empty() {
            return;
        }
        let cap = self.model.max_rows();
        // re-pack: the not-yet-run pieces of the pending job and the new segments are chunked TOGETHER (piece order kept, so a request's shared
        // prefix still precedes its branch tails), so every forward is as full as possible instead of the old job's tail chunk running alone
        if let Some(j) = self.job.take() {
            let mut pieces: Vec<(usize, usize, usize)> = j.chunks[j.next..].concat();
            let mut all = j.segs;
            let base = all.len();
            for (i, s) in segs.iter().enumerate() {
                pieces.push((base + i, 0, s.tokens.len()));
            }
            let n_new = segs.len();
            all.extend(segs);
            let mut rids = j.rids;
            rids.extend(admitted.iter().copied());
            match self.chunk_pieces(&all, &pieces, cap) {
                Ok(chunks) => self.job = Some(PrefillJob { segs: all, chunks, next: 0, rids }),
                Err(_) => {
                    // fall back: old pieces as they were, new segments chunked on their own behind them
                    let old: Vec<(usize, usize, usize)> = pieces[..pieces.len() - n_new].to_vec();
                    let new_segs: Vec<Seg> = all.drain(base..).collect();
                    let mut chunks = if old.is_empty() { vec![] } else { self.chunk_pieces(&all, &old, cap).unwrap_or_else(|_| vec![old.clone()]) };
                    match self.build_chunks(&new_segs, cap) {
                        Ok(c) => {
                            all.extend(new_segs);
                            for c in c {
                                chunks.push(c.into_iter().map(|(si, lo, hi)| (si + base, lo, hi)).collect());
                            }
                            self.job = Some(PrefillJob { segs: all, chunks, next: 0, rids });
                        }
                        Err(e) => {
                            if !chunks.is_empty() {
                                self.job = Some(PrefillJob { segs: all, chunks, next: 0, rids });
                            }
                            for r in admitted {
                                self.fail_req(r, format!("prefill failed: {e:#}"));
                            }
                        }
                    }
                }
            }
            return;
        }
        match self.build_chunks(&segs, cap) {
            Ok(chunks) => self.job = Some(PrefillJob { segs, chunks, next: 0, rids: admitted }),
            Err(e) => {
                for r in admitted {
                    self.fail_req(r, format!("prefill failed: {e:#}"));
                }
            }
        }
    }

    /// Plan one request: cache lookup per branch, shared region, page allocation, segments. Ok(None) = does not fit now (deferred).
    fn plan_request(&mut self, rid: u64, segs: &mut Vec<Seg>, inflight: &mut Vec<Vec<PKey>>) -> Result<Option<usize>> {
        let (n, lens): (usize, Vec<usize>) = {
            let r = &self.reqs[&rid];
            (r.inp.items.len(), r.inp.ids.iter().map(|x| x.len()).collect())
        };
        // 1. match + lock
        let mut matched: Vec<Vec<i32>> = vec![];
        let mut nodes: Vec<NodeId> = vec![];
        let mut cmax: Vec<usize> = vec![];
        let mut qpages = 0usize;
        for b in 0..n {
            let l = lens[b];
            let c = (l - 1) / PAGE; // trap (b): the last two prompt positions are always recomputed
            cmax.push(c);
            qpages += c;
            let (mut pages, node) = if self.cfg.prefix_cache && c > 0 {
                let k = page_keys(&self.reqs[&rid].inp.keys[b][..c * PAGE]);
                self.tree.match_prefix(&k)
            } else {
                (vec![], ROOT)
            };
            if let Some(sw) = &self.swa {
                // SWA partial miss: the hit is only usable up to the last point whose preceding window is still resident (the node lock stays deeper: harmless)
                pages.truncate(sw.valid_len(&pages));
            }
            self.tree.inc_lock(node);
            for &p in &pages {
                self.pool.incref(p);
            }
            matched.push(pages);
            nodes.push(node);
        }
        let rollback = |s: &mut Sched, matched: &Vec<Vec<i32>>, nodes: &Vec<NodeId>| {
            for (m, &nd) in matched.iter().zip(nodes) {
                for &p in m {
                    s.pool.decref(p);
                }
                s.tree.dec_lock(nd);
            }
        };
        // in-batch prefix de-duplication (SGLang schedule_policy in-batch prefix caching): wait for the request admitted earlier in this iteration
        if self.cfg.prefix_cache && !inflight.is_empty() {
            let k0 = page_keys(&self.reqs[&rid].inp.keys[0][..cmax[0] * PAGE]);
            let best = inflight.iter().map(|o| o.iter().zip(&k0).take_while(|(a, b)| a == b).count()).max().unwrap_or(0);
            if best > matched[0].len() {
                rollback(self, &matched, &nodes);
                self.st.deferred_in_batch += 1;
                return Ok(None);
            }
        }
        // 2. shared region
        let r = &self.reqs[&rid];
        let keys = &r.inp.keys;
        let mut sp_pages = 0usize;
        if n >= 2 {
            let mut lcp = keys[0].len();
            for k in &keys[1..] {
                lcp = lcp.min(keys[0].iter().zip(k).take_while(|(a, b)| a == b).count());
            }
            sp_pages = (lcp / PAGE).min(*cmax.iter().min().unwrap());
            // trap (e): an image block must not straddle the shared/tail boundary
            loop {
                let bnd = sp_pages * PAGE;
                match r.inp.blocks[0].iter().find(|bl| bl.s < bnd && bnd < bl.e) {
                    Some(bl) => sp_pages = bl.s / PAGE,
                    None => break,
                }
            }
        }
        let mcl: Vec<usize> = matched.iter().map(|m| m.len()).collect();
        let ms = (0..n).map(|b| mcl[b].min(sp_pages)).min().unwrap();
        let need_shared = sp_pages - ms;
        let total_pages: Vec<usize> = lens.iter().map(|&l| (l + 1 + PAGE - 1) / PAGE).collect();
        let start_pages: Vec<usize> = (0..n).map(|b| sp_pages.max(matched[b].len())).collect();
        let need: usize = need_shared + (0..n).map(|b| total_pages[b] - start_pages[b]).sum::<usize>();
        let _ = r;
        // a request that needs more pages than the WHOLE pool can never be admitted (not a transient shortage): fail it as a capacity limit
        if need > self.pool.rc.len() {
            rollback(self, &matched, &nodes);
            bail!("KV pool exhausted: KV cache capacity exceeded ({} pages needed, {} pages total)", need, self.pool.rc.len());
        }
        // 3. allocate
        let alloc = match self.alloc(need) {
            Ok(a) => a,
            Err(e) => {
                rollback(self, &matched, &nodes);
                if self.live_branches == 0 && segs.is_empty() {
                    return Err(e);
                }
                return Ok(None);
            }
        };
        // sliding-window reserve: the most sliding pages this request's prefill can hold at once (exact for short prompts; window + one forward for long ones)
        if self.swa_on() {
            let sw = self.swa.as_ref().unwrap();
            let bound = sw.wpages() + 2 * SWA_SLACK_PAGES + self.model.max_rows() / PAGE + 4;
            let p_req = need_shared.min(bound) + (0..n).map(|b| (total_pages[b] - start_pages[b]).min(bound)).sum::<usize>();
            if p_req > sw.n {
                for &p in &alloc {
                    self.pool.decref(p);
                }
                rollback(self, &matched, &nodes);
                let total = self.swa.as_ref().unwrap().n;
                bail!("sliding-window KV pool exhausted: KV cache capacity exceeded ({} pages needed, {} pages total)", p_req, total);
            }
            if !self.sl_make_free(p_req) {
                for &p in &alloc {
                    self.pool.decref(p);
                }
                rollback(self, &matched, &nodes);
                if self.live_branches == 0 && segs.is_empty() {
                    let sw = self.swa.as_ref().unwrap();
                    bail!("sliding-window KV pool exhausted ({} pages total, {} needed, {} free)", sw.n, p_req, sw.free.len());
                }
                return Ok(None);
            }
            let sw = self.swa.as_mut().unwrap();
            let at = sw.free.len() - p_req;
            let res: Vec<i32> = sw.free.drain(at..).collect();
            let r = self.reqs.get_mut(&rid).unwrap();
            r.sl_res = res;
            r.sl_cap = p_req;
        }
        let mut it = alloc.into_iter();
        let new_shared: Vec<i32> = (0..need_shared).map(|_| it.next().unwrap()).collect();
        let mut shared: Vec<i32> = matched[0][..ms].to_vec();
        shared.extend_from_slice(&new_shared);
        for _ in 1..n {
            for &p in &new_shared {
                self.pool.incref(p);
            }
        }
        // 4. branches + segments
        let (pol, ids0_blocks) = {
            let r = &self.reqs[&rid];
            (r.inp.pol, r.inp.blocks[0].clone())
        };
        let mut rows = 0;
        if need_shared > 0 {
            let r = &self.reqs[&rid];
            segs.push(Seg { tokens: r.inp.ids[0][ms * PAGE..sp_pages * PAGE].to_vec(), start: ms * PAGE, pages: shared.clone(), blocks: ids0_blocks, req: rid, branch: None });
            rows += need_shared * PAGE;
        }
        for b in 0..n {
            let mut pages = shared.clone();
            if matched[b].len() > sp_pages {
                pages.extend_from_slice(&matched[b][sp_pages..]);
            }
            for _ in start_pages[b]..total_pages[b] {
                pages.push(it.next().unwrap());
            }
            // matched pages beyond the usable cap are dropped (their refs were taken: release the surplus)
            let r = &self.reqs[&rid];
            let l = lens[b];
            let start = start_pages[b] * PAGE;
            let mut toks: Vec<i32> = r.inp.ids[b][start..l].to_vec();
            toks.push(MK);
            rows += toks.len();
            let blocks = r.inp.blocks[b].clone();
            let item = &r.inp.items[b];
            let n_unexp = item.ids.len();
            let n_img = r.inp.images.len();
            let room = self.cfg.max_len as i64 - n_unexp as i64 - ((dllm_core::tokenizer::img_soft_max() + 2) * n_img) as i64 - 4;
            let mt = (pol.cap as i64 + 3).min(room).max(4);
            let maxgen = (pol.cap as i64 + 1).min(mt) as usize;
            let br = Branch {
                req: rid,
                item: b,
                number: item.number,
                ans: vec![],
                ans_lp: 0.0,
                ans_pos: 0,
                l,
                pol,
                pages: pages.clone(),
                node: nodes[b],
                phase: Phase::Prefill,
                p0: vec![],
                gi: vec![],
                fed: 0,
                maxgen,
                base: l,
                feed: vec![],
                probe_pos: 0,
                nt: 0,
                out_tok: 1,
                deadline: if pol.budget_ms > 0 { Some(self.reqs[&rid].arrive + std::time::Duration::from_millis(pol.budget_ms as u64)) } else { None },
                sl: vec![],
                sl_cur: 0,
                sp_pages,
            };
            let bi = self.new_branch(br);
            if self.swa_on() {
                // refs on the resident (cached / shared) pages inside the window of this branch's first query (its tail, or the shared segment it waits for)
                let q0 = if need_shared > 0 { ms * PAGE } else { start };
                let lo = self.swa.as_ref().unwrap().keep_lo(q0);
                let np = self.branches[bi].as_ref().unwrap().pages.len();
                self.branches[bi].as_mut().unwrap().sl = vec![false; np];
                self.branches[bi].as_mut().unwrap().sl_cur = lo;
                for j in lo..start_pages[b].max(sp_pages) {
                    let f = self.branches[bi].as_ref().unwrap().pages[j];
                    if self.swa.as_ref().unwrap().smap[f as usize] != 0 {
                        self.sl_ref(bi, j);
                    }
                }
            }
            segs.push(Seg { tokens: toks, start, pages, blocks, req: rid, branch: Some(bi) });
        }
        self.st.cache_query_pages += qpages as u64;
        self.st.cache_hit_pages += (0..n).map(|b| matched[b].len().min(cmax[b])).sum::<usize>() as u64;
        if self.cfg.prefix_cache {
            inflight.push(page_keys(&self.reqs[&rid].inp.keys[0][..cmax[0] * PAGE]));
        }
        Ok(Some(rows))
    }

    fn soft_for(&mut self, rid: u64, img: usize) -> Result<Arc<SoftTokens>> {
        let (hash, gh, gw, patches) = {
            let p = &self.reqs[&rid].inp.images[img];
            (p.hash ^ ((p.n_soft as u64) << 48), p.gh, p.gw, p.patches.clone())
        };
        if let Some(pos) = self.soft_cache.iter().position(|(h, _)| *h == hash) {
            let e = self.soft_cache.remove(pos);
            self.soft_cache.push(e.clone());
            self.st.soft_cache_hits += 1;
            return Ok(e.1);
        }
        let v = self.vision.as_mut().ok_or_else(|| anyhow!("no vision encoder"))?;
        let s = Arc::new(v.encode(&patches, gh, gw)?);
        self.st.vision_encodes += 1;
        self.soft_cache.push((hash, s.clone()));
        while self.soft_cache.len() > self.cfg.soft_cache_images.max(1) {
            self.soft_cache.remove(0);
        }
        Ok(s)
    }

    fn build_chunks(&self, segs: &[Seg], max_rows: usize) -> Result<Vec<Vec<(usize, usize, usize)>>> {
        // chunking: pieces (seg, lo, hi)
        let mut chunks: Vec<Vec<(usize, usize, usize)>> = vec![vec![]];
        let mut cur = 0usize;
        for (si, s) in segs.iter().enumerate() {
            let mut lo = 0;
            while lo < s.tokens.len() {
                let rem = s.tokens.len() - lo;
                let room = max_rows - cur;
                if rem <= room {
                    chunks.last_mut().unwrap().push((si, lo, s.tokens.len()));
                    cur += rem;
                    break;
                }
                // split: cut <= room, not inside an image block, tail segs keep their last two rows together
                let mut cut = room.min(rem);
                if s.branch.is_some() {
                    cut = cut.min(rem.saturating_sub(2));
                }
                loop {
                    let a = s.start + lo + cut;
                    match s.blocks.iter().find(|b| b.s < a && a < b.e) {
                        Some(b) => cut = if b.s > s.start + lo { b.s - (s.start + lo) } else { 0 },
                        None => break,
                    }
                    if cut == 0 {
                        break;
                    }
                }
                if cut == 0 {
                    if chunks.last().unwrap().is_empty() {
                        bail!("cannot chunk prefill: an unsplittable image block / row group exceeds max_rows {max_rows}");
                    }
                    chunks.push(vec![]);
                    cur = 0;
                    continue;
                }
                chunks.last_mut().unwrap().push((si, lo, lo + cut));
                lo += cut;
                chunks.push(vec![]);
                cur = 0;
            }
        }
        chunks.retain(|c| !c.is_empty());
        Ok(chunks)
    }

    fn run_chunk(&mut self, segs: &[Seg], ch: &[(usize, usize, usize)]) -> Result<()> {
        let mut b = RowBatch::default();
        let mut offs: HashMap<usize, i32> = HashMap::new();
        let mut done: Vec<(usize, usize, usize)> = vec![]; // (branch, gate row, readout row)
        let mut soft_rows: Vec<(u64, usize, i32, i32)> = vec![]; // (req, img, row, idx)
        for &(si, lo, hi) in ch {
            let s = &segs[si];
            let off = *offs.entry(si).or_insert_with(|| {
                let o = b.block_table.len() as i32;
                b.block_table.extend_from_slice(&s.pages);
                o
            });
            for i in lo..hi {
                let pos = s.start + i;
                let row = b.tokens.len() as i32;
                let blk = s.blocks.iter().find(|bl| bl.s <= pos && pos < bl.e);
                b.tokens.push(s.tokens[i]);
                b.pos.push(pos as i32);
                b.slot.push(s.pages[pos / PAGE] * PAGE as i32 + (pos % PAGE) as i32);
                b.bt_off.push(off);
                b.kv_hi.push(pos as i32);
                b.kv_hi_sliding.push(match blk {
                    Some(bl) => (bl.e - 1) as i32,
                    None => pos as i32,
                });
                if let Some(bl) = blk {
                    soft_rows.push((s.req, bl.img, row, (pos - bl.s) as i32));
                }
            }
            if let (Some(bi), true) = (s.branch, hi == s.tokens.len()) {
                let r = b.tokens.len() - 1;
                done.push((bi, r - 1, r));
            }
        }
        // soft tokens (tower encodes of uncached images run here, asynchronously, and finish inside the forward's sync below)
        let (tv, enc0) = (Instant::now(), self.st.vision_encodes);
        let mut holders: Vec<Arc<SoftTokens>> = vec![];
        let mut map: HashMap<(u64, usize), usize> = HashMap::new();
        let mut rows: Vec<SoftRow> = vec![];
        for (rid, img, row, idx) in soft_rows {
            let k = *match map.get(&(rid, img)) {
                Some(k) => k,
                None => {
                    let t = self.soft_for(rid, img)?;
                    holders.push(t);
                    map.insert((rid, img), holders.len() - 1);
                    map.get(&(rid, img)).unwrap()
                }
            };
            rows.push(SoftRow { img: k, row, idx });
        }
        if self.swa_on() {
            self.sl_chunk(segs, ch, false)?;
            self.swa_flush()?;
        }
        let t0 = Instant::now();
        if rows.is_empty() {
            self.model.forward(&b, None)?;
        } else {
            let sel = SoftSel { imgs: holders.iter().map(|h| h.as_ref()).collect(), rows };
            self.model.forward(&b, Some(&sel))?;
        }
        self.model.sync()?;
        let el = t0.elapsed().as_secs_f64() * 1e3;
        let n_enc = self.st.vision_encodes - enc0;
        if n_enc > 0 {
            // time beyond the forward model = the tower encodes; the chunk stays out of the forward fit
            let (a, bb) = self.pf_model();
            let vis = tv.elapsed().as_secs_f64() * 1e3 - (a + bb * b.rows() as f64);
            self.ema_vis_ms = (0.8 * self.ema_vis_ms + 0.2 * (vis / n_enc as f64)).clamp(1.0, 100.0);
        }
        if b.rows() >= 32 && n_enc == 0 {
            self.ema_pf_ms_row = 0.8 * self.ema_pf_ms_row + 0.2 * (el / b.rows() as f64);
            let (x, y) = (b.rows() as f64, el);
            let f = &mut self.pf_fit;
            for v in f.iter_mut() {
                *v *= 0.98;
            }
            f[0] += 1.0;
            f[1] += x;
            f[2] += y;
            f[3] += x * x;
            f[4] += x * y;
        }
        self.st.prefill_forwards += 1;
        self.st.prefill_rows += b.rows() as u64;
        if self.swa_on() {
            self.sl_chunk(segs, ch, true)?;
        }
        if !done.is_empty() {
            self.finish_prefill(done)?;
        }

        Ok(())
    }

    /// end of a branch's prefill: letters at the readout row, gate decision, first thought token, cache insertion
    fn finish_prefill(&mut self, done: Vec<(usize, usize, usize)>) -> Result<()> {
        // letter logits grouped by label set
        let mut p0: HashMap<usize, Vec<f32>> = HashMap::new();
        let mut groups: Vec<(Vec<i32>, Vec<(usize, usize)>)> = vec![];
        for &(bi, _g, r) in &done {
            let b = self.branches[bi].as_ref().unwrap();
            if b.number {
                continue;
            }
            let lab: Vec<i32> = self.reqs[&b.req].inp.items[b.item].label_ids.iter().map(|&x| x as i32).collect();
            match groups.iter_mut().find(|g| g.0 == lab) {
                Some(g) => g.1.push((bi, r)),
                None => groups.push((lab, vec![(bi, r)])),
            }
        }
        for (lab, members) in &groups {
            let rows: Vec<usize> = members.iter().map(|m| m.1).collect();
            let lg = self.model.letter_logits(&rows, lab)?;
            for ((bi, _), z) in members.iter().zip(lg) {
                p0.insert(*bi, z);
            }
        }
        // gate decisions
        let mut think: Vec<(usize, usize)> = vec![]; // (branch, gate row)
        let rsv_g = self.reserve_ms();
        let mut skipped = 0u64;
        let mut answer_now: Vec<(usize, usize)> = vec![]; // number questions answered without a thought: (branch, end-of-thought row)
        for &(bi, g, r) in &done {
            let b = self.branches[bi].as_mut().unwrap();
            let t = if b.number {
                // the trained gate reads option letters: number questions think only when the request asks for it
                b.pol.think_numbers && b.pol.gate != GateMode::Never
            } else {
                let z = p0.remove(&bi).unwrap();
                let pm = softmax64(&z).into_iter().fold(0.0f64, f64::max);
                b.p0 = z;
                match b.pol.gate {
                    GateMode::Tau => pm < b.pol.tau as f64,
                    GateMode::Never => false,
                }
            };
            // a request whose budget is (almost) gone cannot think usefully; skip thinking (answer from p0) instead of burning decode steps + a readout on it
            let hopeless = t && b.deadline.map_or(false, |d| d.saturating_duration_since(Instant::now()).as_secs_f64() * 1e3 - rsv_g < MIN_THINK_MS);
            if hopeless {
                skipped += 1;
            } else if t {
                think.push((bi, g));
            }
            if b.number && !(t && !hopeless) {
                answer_now.push((bi, r));
            }
        }
        self.st.skipped_think += skipped;
        // first thought token d = argmax excluding MK / EOS / <turn|>
        if !think.is_empty() {
            let rows: Vec<usize> = think.iter().map(|x| x.1).collect();
            let d = self.model.argmax_rows(&rows, &[MK, EOS, TURN_END])?;
            for ((bi, _), t) in think.iter().zip(d) {
                let b = self.branches[*bi].as_mut().unwrap();
                b.gi = vec![t];
                b.fed = 0;
                if b.pol.think_ms > 0 {
                    b.phase = Phase::Gated;
                } else {
                    b.phase = Phase::Think;
                    self.st.think_items += 1;
                }
            }
        }
        // number questions without a thought: the end-of-thought row gives the first answer token
        if !answer_now.is_empty() {
            let rows: Vec<usize> = answer_now.iter().map(|x| x.1).collect();
            let lg = self.model.logits_rows(&rows)?;
            for (&(bi, row), z) in answer_now.iter().zip(lg) {
                let pos = self.branches[bi].as_ref().unwrap().l + 1;
                self.branches[bi].as_mut().unwrap().phase = Phase::Answer;
                self.start_answer(bi, &z, pos, row);
            }
        }
        let mut rids: Vec<u64> = done.iter().filter_map(|d| self.branches[d.0].as_ref().map(|b| b.req)).collect();
        rids.sort_unstable();
        rids.dedup();
        // sliding pool: a request whose prefill is complete returns the rest of its reserve (thinking growth draws from the free list)
        if self.swa_on() {
            for &rid in &rids {
                let pending = (0..self.branches.len()).any(|i| !done.iter().any(|d| d.0 == i) && self.branches[i].as_ref().map_or(false, |b| b.req == rid && b.phase == Phase::Prefill));
                if !pending {
                    if let Some(r) = self.reqs.get_mut(&rid) {
                        let res = std::mem::take(&mut r.sl_res);
                        r.sl_cap = 0;
                        self.sl_drop_reserve(res);
                    }
                }
            }
        }
        // cache insertion + completion of non-thinking branches
        let bis: Vec<usize> = done.iter().map(|d| d.0).filter(|&bi| self.branches[bi].is_some()).collect();
        for bi in bis {
            if self.cfg.prefix_cache {
                self.insert_branch(bi);
            }
            let b = self.branches[bi].as_ref().unwrap();
            if self.swa_on() && matches!(b.phase, Phase::Think | Phase::Answer) {
                let q = if b.phase == Phase::Answer { b.ans_pos + b.ans.len() - 1 } else { b.l };
                self.sl_trim(bi, q);
            }
            let b = self.branches[bi].as_ref().unwrap();
            if b.phase == Phase::Prefill {
                let res = ItemResult { p0: b.p0.clone(), think: None, trace_tokens: 0, out_tokens: 1, ..Default::default() };
                self.finish_item(bi, res);
            }
        }
        // --think-ms: once every question of a request has its gate decision, its gated branches get the request's deterministic cap
        for rid in rids {
            self.release_gated(rid);
        }
        Ok(())
    }

    /// `--think-ms`: when no question of request `rid` is still in prefill, give its gated branches the per-branch cap of the fixed cost model
    /// (a pure function of the request: prefill rows, images, questions and the number of thinking questions) and start them; a cap of 0
    /// answers them from p0 without thinking
    fn release_gated(&mut self, rid: u64) {
        let mut gated = vec![];
        for (i, b) in self.branches.iter().enumerate() {
            if let Some(b) = b.as_ref().filter(|b| b.req == rid) {
                match b.phase {
                    Phase::Prefill => return,
                    Phase::Gated => gated.push(i),
                    _ => {}
                }
            }
        }
        if gated.is_empty() {
            return;
        }
        let shape = ThinkShape { thinking: gated.len(), ..self.reqs[&rid].inp.shape };
        let pol = self.branches[gated[0]].as_ref().unwrap().pol;
        let cap = pol.think_cost.cap(pol.think_ms as f64, &shape, pol.cap);
        for bi in gated {
            if cap == 0 && self.branches[bi].as_ref().unwrap().number {
                // no thought: re-read the end-of-thought row (position l, already in the KV cache) for the first answer token
                self.st.skipped_think += 1;
                let b = self.branches[bi].as_mut().unwrap();
                b.gi.clear();
                b.feed = vec![MK];
                b.probe_pos = b.l;
                b.phase = Phase::Probe;
            } else if cap == 0 {
                self.st.skipped_think += 1;
                let b = self.branches[bi].as_ref().unwrap();
                let res = ItemResult { p0: b.p0.clone(), think: None, trace_tokens: 0, out_tokens: 1, ..Default::default() };
                self.finish_item(bi, res);
            } else {
                let b = self.branches[bi].as_mut().unwrap();
                b.pol.cap = cap;
                b.maxgen = b.maxgen.min(cap + 1);
                b.phase = Phase::Think;
                self.st.think_items += 1;
            }
        }
    }

    /// hand the full prompt pages of a prefilled branch to the tree (SGLang cache_unfinished_req)
    fn insert_branch(&mut self, bi: usize) {
        let (rid, item, l) = {
            let b = self.branches[bi].as_ref().unwrap();
            (b.req, b.item, b.l)
        };
        let n_full = l / PAGE;
        if n_full == 0 {
            return;
        }
        let key = page_keys(&self.reqs[&rid].inp.keys[item][..n_full * PAGE]);
        let pages: Vec<i32> = self.branches[bi].as_ref().unwrap().pages[..n_full].to_vec();
        let (mm, _) = self.tree.insert(&key, &pages);
        for &p in &pages[mm..n_full] {
            self.pool.incref(p); // adopted by the tree
        }
        let (tp, node) = self.tree.match_prefix(&key);
        debug_assert_eq!(tp.len(), n_full);
        self.tree.inc_lock(node);
        let old = self.branches[bi].as_ref().unwrap().node;
        if old != ROOT {
            self.tree.dec_lock(old);
        }
        let b = self.branches[bi].as_mut().unwrap();
        b.node = node;
        for i in 0..n_full {
            let b = self.branches[bi].as_mut().unwrap();
            if b.pages[i] != tp[i] {
                // identical content already cached under another page: switch to it, drop ours
                // (sliding pool: only if the cached page still has the sliding KV this branch's window needs)
                let held = b.sl.get(i).copied().unwrap_or(false);
                if held && self.swa.as_ref().unwrap().smap[tp[i] as usize] == 0 {
                    continue;
                }
                let ours = b.pages[i];
                b.pages[i] = tp[i];
                self.pool.incref(tp[i]);
                if held {
                    self.branches[bi].as_mut().unwrap().sl[i] = false;
                    self.sl_ref(bi, i);
                    self.sl_unref(ours);
                }
                self.pool.decref(ours);
            }
        }
        // the tree keeps the sliding KV of the prompt's window tail (+1 window) and of the window before the shared-prefix boundary
        if self.swa_on() {
            let sp = self.branches[bi].as_ref().unwrap().sp_pages;
            for r in self.sl_claim_ranges(l, sp) {
                for j in r.start..r.end.min(n_full) {
                    let f = tp[j] as usize;
                    let sw = self.swa.as_mut().unwrap();
                    if sw.smap[f] != 0 && !sw.tree_sl[f] {
                        sw.tree_sl[f] = true;
                        sw.tree_n += 1;
                        sw.src[f] += 1;
                    }
                }
            }
        }
    }

    // -------------------------------------------------------------------------------------------- decode
    fn ensure_pages(&mut self, bi: usize, last_pos: usize) -> Result<()> {
        let have = self.branches[bi].as_ref().unwrap().pages.len();
        let need = last_pos / PAGE + 1;
        if need > have {
            let np = self.alloc(need - have)?;
            self.branches[bi].as_mut().unwrap().pages.extend(np);
        }
        Ok(())
    }

    fn decode_step(&mut self) {
        let la = self.think.lookahead();
        let now = Instant::now() + std::time::Duration::from_secs_f64(self.reserve_ms() / 1e3);
        let late: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.phase == Phase::Think && b.deadline.map_or(false, |d| now >= d))).collect();
        for bi in late {
            self.st.forced_by_budget += 1;
            self.end_think(bi);
        }
        let mut act: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.phase == Phase::Think)).collect();
        if act.is_empty() {
            return;
        }
        // pages for the positions this step may write; a branch that cannot grow fails its request
        let mut failed: Vec<(usize, String)> = vec![];
        let mut cut: Vec<usize> = vec![];
        for &bi in &act {
            let pos = {
                let b = self.branches[bi].as_ref().unwrap();
                b.base + b.gi.len() - 1
            };
            if let Err(e) = self.ensure_pages(bi, pos + la) {
                failed.push((bi, format!("{e:#}")));
            } else if self.swa_on() {
                if self.sl_write(bi, pos, pos + la).is_err() {
                    // sliding pool exhausted: the branch stops thinking and is read out (no new pages needed)
                    cut.push(bi);
                }
            }
        }
        for &bi in &cut {
            self.st.swa_forced += 1;
            self.end_think(bi);
        }
        act.retain(|i| !cut.contains(i));
        for (bi, m) in &failed {
            self.fail_branches(&[*bi], m);
        }
        act.retain(|&i| self.branches[i].is_some());
        if act.is_empty() {
            return;
        }
        if let Err(e) = self.swa_flush() {
            self.fail_branches(&act, &format!("decode failed: {e:#}"));
            return;
        }
        let out = {
            let reqs: Vec<ThinkReq> = act
                .iter()
                .map(|&bi| {
                    let b = self.branches[bi].as_ref().unwrap();
                    ThinkReq { pages: &b.pages, pos: b.base + b.gi.len() - 1, tok: *b.gi.last().unwrap(), uid: bi }
                })
                .collect();
            let t0 = Instant::now();
            let r = self.think.step(self.model.as_mut(), &reqs);
            let el = t0.elapsed().as_secs_f64() * 1e3;
            self.ema_dec_ms = 0.9 * self.ema_dec_ms + 0.1 * el;
            r
        };
        let out = match out {
            Ok(o) => o,
            Err(e) => {
                self.fail_branches(&act, &format!("decode failed: {e:#}"));
                return;
            }
        };
        self.st.decode_forwards += 1;
        for (&bi, toks) in act.iter().zip(out) {
            let b = self.branches[bi].as_mut().unwrap();
            b.fed = b.gi.len(); // the last token has been fed by this step
            let mut fin = false;
            for t in toks {
                b.gi.push(t);
                if is_stop(t) || b.gi.len() - 1 >= b.maxgen {
                    fin = true;
                    break;
                }
            }
            {
                // every returned token except the last was a row of the decoder's forward (KV written); only the last one still has to be fed
                // (equals the old `fed = len before the step` for one-token decoders; also when the step ended thinking: a speculative step may have accepted several tokens)
                let b = self.branches[bi].as_mut().unwrap();
                b.fed = b.gi.len() - 1;
            }
            if fin {
                self.end_think(bi);
            }
        }
        if self.swa_on() {
            for &bi in &act {
                let q = match self.branches[bi].as_ref() {
                    Some(b) if b.phase == Phase::Probe => b.probe_pos,
                    Some(b) if b.phase == Phase::Think => b.base + b.gi.len() - 1,
                    _ => continue,
                };
                self.sl_trim(bi, q);
            }
        }
    }

    fn end_think(&mut self, bi: usize) {
        let b = self.branches[bi].as_mut().unwrap();
        // thought = gi[:cap] cut at the first end-of-thought / stop token
        let cap = b.pol.cap;
        let mut th = b.gi[..b.gi.len().min(cap)].to_vec();
        if let Some(j) = th.iter().position(|&t| is_stop(t)) {
            th.truncate(j);
        }
        b.nt = th.len();
        b.out_tok = b.gi.len() + 1;
        let fed = b.fed.min(th.len());
        b.probe_pos = b.l + fed;
        b.feed = th[fed..].to_vec();
        b.feed.push(MK);
        b.phase = Phase::Probe;
    }

    fn probe_step(&mut self) -> bool {
        let mut act: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.phase == Phase::Probe)).collect();
        if act.is_empty() {
            return false;
        }
        let mut failed = vec![];
        for &bi in &act {
            let (last, ) = {
                let b = self.branches[bi].as_ref().unwrap();
                (b.probe_pos + b.feed.len() - 1,)
            };
            if let Err(e) = self.ensure_pages(bi, last) {
                failed.push((bi, format!("{e:#}")));
            } else if self.swa_on() {
                let p0 = self.branches[bi].as_ref().unwrap().probe_pos;
                if let Err(e) = self.sl_write(bi, p0, last) {
                    failed.push((bi, format!("{e:#}")));
                }
            }
        }
        for (bi, m) in &failed {
            self.fail_branches(&[*bi], m);
        }
        act.retain(|&i| self.branches[i].is_some());
        let max_rows = self.model.max_rows();
        // chunk by rows
        let mut start = 0;
        while start < act.len() {
            let mut end = start;
            let mut rows = 0;
            while end < act.len() {
                let n = self.branches[act[end]].as_ref().unwrap().feed.len();
                if rows + n > max_rows && end > start {
                    break;
                }
                rows += n;
                end += 1;
            }
            let chunk: Vec<usize> = act[start..end].to_vec();
            start = end;
            if let Err(e) = self.probe_chunk(&chunk) {
                self.fail_branches(&chunk, &format!("probe failed: {e:#}"));
            }
        }
        true
    }

    fn probe_chunk(&mut self, chunk: &[usize]) -> Result<()> {
        let mut toks = vec![];
        let mut seqs_meta: Vec<(usize, usize)> = vec![]; // (branch, pos)
        let mut last_rows = vec![];
        for &bi in chunk {
            let b = self.branches[bi].as_ref().unwrap();
            for (i, &t) in b.feed.iter().enumerate() {
                toks.push(t);
                seqs_meta.push((bi, b.probe_pos + i));
            }
            last_rows.push(toks.len() - 1);
        }
        let batch = {
            let seqs: Vec<(&[i32], usize)> = seqs_meta.iter().map(|&(bi, pos)| (self.branches[bi].as_ref().unwrap().pages.as_slice(), pos)).collect();
            RowBatch::decode(&toks, &seqs)
        };
        self.swa_flush()?;
        self.model.forward(&batch, None)?;
        self.st.probe_forwards += 1;
        let mut groups: Vec<(Vec<i32>, Vec<(usize, usize)>)> = vec![];
        let mut answers: Vec<(usize, usize)> = vec![];
        for (k, &bi) in chunk.iter().enumerate() {
            let b = self.branches[bi].as_ref().unwrap();
            if b.number {
                answers.push((bi, last_rows[k]));
                continue;
            }
            let lab: Vec<i32> = self.reqs[&b.req].inp.items[b.item].label_ids.iter().map(|&x| x as i32).collect();
            match groups.iter_mut().find(|g| g.0 == lab) {
                Some(g) => g.1.push((bi, last_rows[k])),
                None => groups.push((lab, vec![(bi, last_rows[k])])),
            }
        }
        for (lab, members) in groups {
            let rows: Vec<usize> = members.iter().map(|m| m.1).collect();
            let lg = self.model.letter_logits(&rows, &lab)?;
            for ((bi, _), z) in members.into_iter().zip(lg) {
                self.probe_done(bi, z);
            }
        }
        if !answers.is_empty() {
            let rows: Vec<usize> = answers.iter().map(|m| m.1).collect();
            let lg = self.model.logits_rows(&rows)?;
            for ((bi, row), z) in answers.into_iter().zip(lg) {
                let b = self.branches[bi].as_mut().unwrap();
                let pos = b.probe_pos + b.feed.len();
                if b.nt > 0 || !b.gi.is_empty() {
                    self.st.thought_tokens += b.nt as u64;
                }
                b.phase = Phase::Answer;
                self.start_answer(bi, &z, pos, row);
            }
        }
        Ok(())
    }

    /// number question: first answer token from the logits of the end-of-thought row (`row` of the last forward); the answer continues in `answer_step`
    fn start_answer(&mut self, bi: usize, z: &[f32], pos: usize, row: usize) {
        let (t, lp) = greedy_logprob(z);
        let b = self.branches[bi].as_mut().unwrap();
        b.ans_pos = pos;
        if is_stop(t) {
            self.finish_number(bi);
        } else {
            b.ans = vec![t];
            b.ans_lp = lp;
            if let Err(e) = self.think.prime(self.model.as_mut(), bi, row, pos, t) {
                self.fail_branches(&[bi], &format!("answer step failed: {e:#}"));
            }
        }
    }

    /// one decode cycle for every answering number question (all of them in one step of the decoder: speculative with the MTP draft, or one greedy row)
    fn answer_step(&mut self) {
        let act: Vec<usize> = (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.phase == Phase::Answer)).collect();
        if act.is_empty() {
            return;
        }
        let la = self.think.lookahead();
        let mut ok = vec![];
        for &bi in &act {
            let pos = {
                let b = self.branches[bi].as_ref().unwrap();
                b.ans_pos + b.ans.len() - 1
            };
            let room = self.ensure_pages(bi, pos + la).is_ok() && (!self.swa_on() || self.sl_write(bi, pos, pos + la).is_ok());
            if room {
                ok.push(bi);
            } else {
                self.finish_number(bi); // pool exhausted: answer with the tokens so far
            }
        }
        if ok.is_empty() {
            return;
        }
        if let Err(e) = self.swa_flush() {
            self.fail_branches(&ok, &format!("answer step failed: {e:#}"));
            return;
        }
        let out = {
            let reqs: Vec<ThinkReq> = ok
                .iter()
                .map(|&bi| {
                    let b = self.branches[bi].as_ref().unwrap();
                    ThinkReq { pages: &b.pages, pos: b.ans_pos + b.ans.len() - 1, tok: *b.ans.last().unwrap(), uid: bi }
                })
                .collect();
            self.think.answer_step(self.model.as_mut(), &reqs)
        };
        let out = match out {
            Ok(o) => o,
            Err(e) => {
                self.fail_branches(&ok, &format!("answer step failed: {e:#}"));
                return;
            }
        };
        self.st.decode_forwards += 1;
        for (&bi, toks) in ok.iter().zip(out) {
            let b = self.branches[bi].as_mut().unwrap();
            let mut fin = false;
            for (t, lp) in toks {
                if is_stop(t) || b.ans.len() >= ANSWER_MAX {
                    fin = true;
                    break;
                }
                b.ans.push(t);
                b.ans_lp += lp;
            }
            let q = b.ans_pos + b.ans.len() - 1;
            if fin {
                self.finish_number(bi);
            } else if self.swa_on() {
                self.sl_trim(bi, q);
            }
        }
    }

    fn finish_number(&mut self, bi: usize) {
        let b = self.branches[bi].as_ref().unwrap();
        let thought = !b.gi.is_empty();
        let res = ItemResult {
            p0: vec![],
            think: thought.then(Vec::new),
            trace_tokens: if thought { b.nt } else { 0 },
            out_tokens: if thought { b.out_tok } else { 1 } + b.ans.len() + 1,
            trace_ids: if thought { trace_ids(&b.gi, b.nt, &b.pol) } else { vec![] },
            answer_ids: b.ans.iter().map(|&t| t as u32).collect(),
            answer_logprob: b.ans_lp,
        };
        self.finish_item(bi, res);
    }

    fn probe_done(&mut self, bi: usize, z: Vec<f32>) {
        let b = self.branches[bi].as_mut().unwrap();
        self.st.thought_tokens += b.nt as u64;
        let res = ItemResult { p0: b.p0.clone(), think: Some(z), trace_tokens: b.nt, out_tokens: b.out_tok, trace_ids: trace_ids(&b.gi, b.nt, &b.pol), ..Default::default() };
        self.finish_item(bi, res);
    }
}

/// the first `n` generated thought tokens (the trace the readout was conditioned on), only when the request asked for it
/// prefill rows of a request: the longest prefix shared by all its (image-expanded) prompts once, plus every prompt's own tail
pub fn prefill_rows(ids: &[Vec<i32>]) -> usize {
    let Some(first) = ids.first() else { return 0 };
    let lcp = ids[1..].iter().fold(first.len(), |l, x| x.iter().zip(&first[..l]).take_while(|(a, b)| a == b).count());
    lcp + ids.iter().map(|x| x.len() - lcp).sum::<usize>()
}

fn trace_ids(gi: &[i32], n: usize, pol: &Policy) -> Vec<u32> {
    if pol.return_trace { gi[..n.min(gi.len())].iter().map(|&t| t as u32).collect() } else { vec![] }
}

// ================================================================================================================ sliding-window KV pool
//
// The 25 sliding-window layers only ever read the last `window` (1024) positions, so their KV lives in a SEPARATE, smaller page pool (vLLM hybrid KV-cache
// manager: per-group block tables with null blocks outside the window; SGLang SWAKVPool / SWARadixCache). Full-attention layers keep the full paged pool.
// Identity of a page = its FULL page id (block tables, slots, radix tree, refcounts are unchanged); the model translates it for the sliding layers through a
// device page map smap[full page] -> sliding page (0 = none, the never-read null page).
//  * attach on first write: before a forward every page a row writes gets a sliding page (from the owning request's reserve, else the global free list);
//  * refcounts: a branch holds a ref on page j while j >= lo(q) where q = the lowest position the branch may still query
//    (lo(q) = (q - window + 1) / 16 - SWA_SLACK_PAGES); the radix tree holds one ref per cached page whose sliding KV it keeps (`tree_sl`);
//  * a page whose refs drop to 0 is detached and its sliding page recycled (into the owner request's reserve while it lives, else the free list);
//  * prefix cache: a hit at page m is only usable if the pages covering the window before position 16m are resident (`sl_valid_len`); otherwise
//    the match is cut back (partial miss: only the cut part is recomputed). Inserting a branch lets the tree keep the window tail of the prompt
//    (plus one extra window and the window before the request's shared-prefix boundary, the usual divergence point of later requests);
//  * admission: a request reserves the sliding pages its prefill can hold at peak (all new pages for short prompts; window + one forward of rows for long ones),
//    so an admitted prefill never runs out mid-way (same guarantee as the up-front full-page allocation).
pub(crate) const SWA_SLACK_PAGES: usize = 2;

pub(crate) struct Swa {
    pub window: usize,
    pub smap: Vec<i32>,
    pub src: Vec<u32>,
    pub tree_sl: Vec<bool>,
    /// per sliding page: request that attached it
    pub owner: Vec<u64>,
    pub free: Vec<i32>,
    pub n: usize,
    /// pages with tree_sl set
    pub tree_n: usize,
    dirty: Option<(usize, usize)>,
}

impl Swa {
    fn new(full_pages: usize, n: usize, window: usize) -> Swa {
        Swa { window, smap: vec![0; full_pages], src: vec![0; full_pages], tree_sl: vec![false; full_pages], owner: vec![0; n + 1], free: (1..=n as i32).rev().collect(), n, tree_n: 0, dirty: None }
    }
    fn mark(&mut self, f: usize) {
        self.dirty = Some(match self.dirty {
            None => (f, f),
            Some((a, b)) => (a.min(f), b.max(f)),
        });
    }
    /// first page index a query at position `q` may read
    pub fn need_lo(&self, q: usize) -> usize {
        q.saturating_sub(self.window - 1) / PAGE
    }
    /// first page index a branch whose lowest future query is `q` keeps (window + slack)
    pub fn keep_lo(&self, q: usize) -> usize {
        self.need_lo(q).saturating_sub(SWA_SLACK_PAGES)
    }
    /// pages of one window
    pub fn wpages(&self) -> usize {
        (self.window + PAGE - 1) / PAGE + 1
    }
    /// longest usable prefix of a radix match `pages`: every page covering the window before the first recomputed position must be resident
    pub fn valid_len(&self, pages: &[i32]) -> usize {
        let mut m = pages.len();
        loop {
            let lo = self.need_lo(m * PAGE);
            match (lo..m).rev().find(|&j| self.smap[pages[j] as usize] == 0) {
                None => return m,
                Some(t) => m = t,
            }
        }
    }
}

impl Sched {
    fn swa_on(&self) -> bool {
        self.swa.is_some()
    }

    /// evict LRU radix leaves (whole pages) until `need` pages were handed back; releases the tree's sliding refs of those pages
    fn evict_tree(&mut self, need: usize) -> usize {
        let mut ev: Vec<i32> = vec![];
        let got = self.tree.evict(need, &mut |ps| ev.extend_from_slice(ps));
        for p in ev {
            if let Some(sw) = self.swa.as_mut() {
                if sw.tree_sl[p as usize] {
                    sw.tree_sl[p as usize] = false;
                    sw.tree_n -= 1;
                    self.sl_unref(p);
                }
            }
            self.pool.decref(p);
        }
        self.st.cache_evicted_pages += got as u64;
        got
    }

    /// make `n` sliding pages free in the global list (evicting cached prefixes); false if impossible
    fn sl_make_free(&mut self, n: usize) -> bool {
        loop {
            let free = self.swa.as_ref().unwrap().free.len();
            if free >= n {
                return true;
            }
            if self.evict_tree((n - free).max(4)) == 0 {
                return false;
            }
        }
    }

    /// drop one sliding ref of full page `f`; at zero the sliding page is detached and recycled
    fn sl_unref(&mut self, f: i32) {
        let sw = self.swa.as_mut().unwrap();
        let r = &mut sw.src[f as usize];
        assert!(*r > 0, "sliding unref of page {f} without refs");
        *r -= 1;
        if *r > 0 {
            return;
        }
        let s = sw.smap[f as usize];
        sw.smap[f as usize] = 0;
        sw.mark(f as usize);
        let owner = sw.owner[s as usize];
        sw.owner[s as usize] = 0;
        match self.reqs.get_mut(&owner) {
            Some(r) if r.sl_res.len() < r.sl_cap => r.sl_res.push(s),
            _ => self.swa.as_mut().unwrap().free.push(s),
        }
    }

    /// give full page `f` a sliding page (owner request `rid`) if it has none
    fn sl_attach(&mut self, f: i32, rid: u64) -> Result<()> {
        if self.swa.as_ref().unwrap().smap[f as usize] != 0 {
            return Ok(());
        }
        let s = match self.reqs.get_mut(&rid).and_then(|r| r.sl_res.pop()) {
            Some(s) => s,
            None => {
                if !self.sl_make_free(1) {
                    let sw = self.swa.as_ref().unwrap();
                    bail!("sliding-window KV pool exhausted ({} pages, {} free)", sw.n, sw.free.len());
                }
                self.swa.as_mut().unwrap().free.pop().unwrap()
            }
        };
        let sw = self.swa.as_mut().unwrap();
        sw.smap[f as usize] = s;
        sw.owner[s as usize] = rid;
        sw.mark(f as usize);
        Ok(())
    }

    /// branch `bi` takes its sliding ref on page index `j` (the page must be attached)
    fn sl_ref(&mut self, bi: usize, j: usize) {
        let b = self.branches[bi].as_mut().unwrap();
        if b.sl.len() <= j {
            b.sl.resize(j + 1, false);
        }
        if b.sl[j] {
            return;
        }
        let f = b.pages[j] as usize;
        let sw = self.swa.as_mut().unwrap();
        assert!(sw.smap[f] != 0, "sliding ref on unattached page {f}");
        b.sl[j] = true;
        sw.src[f] += 1;
    }

    /// attach + ref the pages of branch `bi` that hold positions `p0..=p1` (rows about to be written)
    fn sl_write(&mut self, bi: usize, p0: usize, p1: usize) -> Result<()> {
        let rid = self.branches[bi].as_ref().unwrap().req;
        for j in p0 / PAGE..=p1 / PAGE {
            let f = self.branches[bi].as_ref().unwrap().pages[j];
            self.sl_attach(f, rid)?;
            self.sl_ref(bi, j);
        }
        Ok(())
    }

    /// the lowest position branch `bi` may still query is `q`: release its refs below the window (+ slack)
    fn sl_trim(&mut self, bi: usize, q: usize) {
        let lo = self.swa.as_ref().unwrap().keep_lo(q);
        let (from, to) = {
            let b = self.branches[bi].as_ref().unwrap();
            (b.sl_cur, lo.min(b.sl.len()))
        };
        for j in from..to {
            let b = self.branches[bi].as_mut().unwrap();
            if b.sl[j] {
                b.sl[j] = false;
                let f = b.pages[j];
                self.sl_unref(f);
            }
        }
        let b = self.branches[bi].as_mut().unwrap();
        b.sl_cur = b.sl_cur.max(to);
    }

    /// push pending page-map changes to the model (before every forward)
    fn swa_flush(&mut self) -> Result<()> {
        let d = match self.swa.as_mut() {
            Some(sw) => sw.dirty.take(),
            None => return Ok(()),
        };
        if let Some((a, b)) = d {
            self.model.set_sliding_map(a, &self.swa.as_ref().unwrap().smap[a..=b])?;
        }
        Ok(())
    }

    /// request `rid` leaves: its unused sliding reserve goes back to the free list
    fn sl_drop_reserve(&mut self, res: Vec<i32>) {
        if let Some(sw) = self.swa.as_mut() {
            sw.free.extend(res);
        }
    }

    /// pages the tree keeps sliding KV for after inserting a branch whose prompt is `l` tokens and whose request shares `sp` prefix pages:
    /// the window tail of the prompt plus one more window, and the window before the shared-prefix boundary
    fn sl_claim_ranges(&self, l: usize, sp: usize) -> Vec<std::ops::Range<usize>> {
        let sw = self.swa.as_ref().unwrap();
        let mut v = vec![sw.keep_lo(l).saturating_sub(sw.wpages())..l / PAGE];
        if sp > 0 {
            v.push(sw.keep_lo(sp * PAGE)..sp);
        }
        v
    }
}

impl Sched {
    /// sliding pool around one prefill forward: before (`after` = false) attach + ref the pages its rows write; after it, release what no future query reads
    fn sl_chunk(&mut self, segs: &[Seg], ch: &[(usize, usize, usize)], after: bool) -> Result<()> {
        for &(si, lo, hi) in ch {
            let s = &segs[si];
            if !self.reqs.contains_key(&s.req) || hi <= lo {
                continue;
            }
            let (p0, p1) = (s.start + lo, s.start + hi - 1);
            let owners: Vec<usize> = match s.branch {
                Some(bi) => vec![bi],
                None => (0..self.branches.len()).filter(|&i| self.branches[i].as_ref().map_or(false, |b| b.req == s.req && b.phase == Phase::Prefill)).collect(),
            };
            if after {
                if s.branch.is_none() || hi < s.tokens.len() {
                    for &bi in &owners {
                        self.sl_trim(bi, p1 + 1);
                    }
                }
            } else {
                for j in p0 / PAGE..=p1 / PAGE {
                    self.sl_attach(s.pages[j], s.req)?;
                    for &bi in &owners {
                        debug_assert_eq!(self.branches[bi].as_ref().unwrap().pages[j], s.pages[j]);
                        self.sl_ref(bi, j);
                    }
                }
            }
        }
        Ok(())
    }
}

impl Sched {
    /// invariants when nothing is live: every full page is free or in the tree; every sliding page is free, reserved, or kept by the tree
    fn check_idle(&self) {
        let n = self.pool.rc.len();
        assert_eq!(self.pool.free.len() + self.tree.total_pages(), n, "idle: full pages leaked");
        if let Some(sw) = &self.swa {
            let mut tree_sl = 0;
            for f in 0..n {
                let t = sw.tree_sl[f] as u32;
                tree_sl += t as usize;
                assert_eq!(sw.src[f], t, "idle: page {f} has {} sliding refs (tree {t})", sw.src[f]);
                assert_eq!(sw.smap[f] != 0, t == 1, "idle: page {f} sliding attachment without a tree ref");
            }
            let res: usize = self.reqs.values().map(|r| r.sl_res.len()).sum();
            assert_eq!(tree_sl, sw.tree_n);
            assert_eq!(sw.free.len() + res + tree_sl, sw.n, "idle: sliding pages leaked");
        }
    }
}