inferencelayer 0.2.8

//! Multi-user serving core: paged KV-block allocator, radix prefix cache, and the
//! continuous-batching scheduler that drives [`crate::forward::BatchPlan`].
//!
//! Design (grounded in `docs/ai-sota/serving-engines-vllm-sglang-p2p.md`):
//! - **Blocks of 16 tokens** (vLLM's measured sweet spot). A block becomes IMMUTABLE once full —
//!   full blocks are the only sharing granularity, so prefix reuse needs **no copy-on-write**:
//!   the writing (partial) block is always private to one sequence.
//! - **Radix prefix cache** as a hash-chained block index (vLLM-V1-style chain
//!   `h_i = hash(h_{i-1}, chunk_i)`, SGLang-style LRU eviction of unreferenced entries). Token
//!   chunks are stored in each entry and compared exactly on hit, so a 64-bit collision cannot
//!   alias two different prefixes.
//! - **Iteration-level scheduling** (Orca/vLLM-V1): no prefill/decode phases — each step fills up
//!   to `k` columns with running sequences' next tokens plus one admitted sequence's prefill chunk.
//! - **Determinism by construction**: every column replays the SAME compiled pipelines as a solo
//!   run (see `make_batch_plan`), so a sequence's output is bitwise-identical no matter what else
//!   shares its batches — the property vLLM lacks and SGLang sells at a measured ~34% tax.
use std::collections::HashMap;
use std::collections::hash_map::DefaultHasher;
use std::hash::{Hash, Hasher};

use anyhow::{Result, anyhow};

use crate::GpuCtx;
use crate::forward::{
    BatchCol, BatchPlan, CHAIN_MAX, DraftPair, DraftReq, Lfm2Gpu, MAX_SLOTS, MtpEngine, PAGE_BLK,
};
use crate::sampling::{self, TokenLogprobs};
use crate::sampling::prompt_lookup_drafts;

/// Maximum columns a single `batch_step` launch can carry — bounded by the engine's u64
/// per-step `need_logit` mask. Caps the split-fuse wide-prefill budget.
const MAX_BATCH_COLS: usize = 64;

/// Paged KV-block allocator over the engine's pool. Block ids are physical pool slots; the last
/// block is reserved as the trash block (idle batch columns write there) and never allocated.
pub struct KvPool {
    free: Vec<u32>,
    /// Reference counts, indexed by block id (0 = free or unreferenced-but-cached).
    refs: Vec<u32>,
}

impl KvPool {
    pub fn new(pool_blocks: usize) -> Self {
        assert!(
            pool_blocks >= 2,
            "pool needs at least one usable block + trash"
        );
        Self {
            // Last block = trash; hand out low ids first (matches the identity row for debug).
            free: (0..pool_blocks as u32 - 1).rev().collect(),
            refs: vec![0; pool_blocks],
        }
    }

    pub fn available(&self) -> usize {
        self.free.len()
    }

    /// Total allocatable blocks (pool size minus the reserved trash block) — the most any single
    /// sequence could ever hold, used to reject never-fits requests up front.
    pub fn capacity(&self) -> usize {
        self.refs.len().saturating_sub(1)
    }

    /// Allocate one block (refcount 1).
    pub fn alloc(&mut self) -> Option<u32> {
        let b = self.free.pop()?;
        self.refs[b as usize] = 1;
        Some(b)
    }

    /// Add a reference to a cached block (radix hit).
    pub fn retain(&mut self, b: u32) {
        self.refs[b as usize] += 1;
    }

    /// Drop a reference; at zero the block is NOT freed — it may live on in the radix index
    /// (call [`Self::release`] for uncached blocks).
    pub fn unref(&mut self, b: u32) -> u32 {
        let r = &mut self.refs[b as usize];
        *r = r.saturating_sub(1);
        *r
    }

    /// Return an unreferenced block to the free list (evicted from cache or never cached).
    pub fn release(&mut self, b: u32) {
        debug_assert_eq!(self.refs[b as usize], 0, "releasing a referenced block");
        self.free.push(b);
    }
}

/// One cached full block: the chain hash identifies the prefix ending at this block.
struct CacheEntry {
    block: u32,
    /// Exact token chunk (collision armor — verified on every hit).
    chunk: Vec<u32>,
    parent: u64,
    last_used: u64,
}

/// Hash-chained radix prefix index at block granularity (attention-only architectures: recurrent
/// state is not prefix-addressable — the SGLang MambaRadixCache caveat — so conv/DeltaNet models
/// serve without prefix reuse).
pub struct RadixIndex {
    map: HashMap<u64, CacheEntry>,
    clock: u64,
}

fn chain_hash(parent: u64, chunk: &[u32]) -> u64 {
    let mut h = DefaultHasher::new();
    parent.hash(&mut h);
    chunk.hash(&mut h);
    h.finish()
}

impl RadixIndex {
    pub fn new() -> Self {
        Self {
            map: HashMap::new(),
            clock: 0,
        }
    }

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

    pub fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Longest cached prefix of `tokens` in FULL blocks: returns (matched blocks, matched tokens).
    /// Bumps LRU stamps and refcounts (caller owns the references).
    pub fn match_prefix(&mut self, pool: &mut KvPool, tokens: &[u32]) -> (Vec<u32>, usize) {
        let mut blocks = Vec::new();
        let mut parent = 0u64;
        self.clock += 1;
        for chunk in tokens.chunks_exact(PAGE_BLK) {
            let h = chain_hash(parent, chunk);
            match self.map.get_mut(&h) {
                Some(e) if e.chunk == chunk => {
                    e.last_used = self.clock;
                    pool.retain(e.block);
                    blocks.push(e.block);
                    parent = h;
                }
                _ => break,
            }
        }
        let n = blocks.len() * PAGE_BLK;
        (blocks, n)
    }

    /// Insert a sequence's full blocks (called as blocks fill). `blocks[i]` holds
    /// `tokens[i·B..(i+1)·B]`. Already-cached prefixes are left as-is (their block stays the
    /// canonical one). Returns how many entries were newly inserted.
    pub fn insert(&mut self, tokens: &[u32], blocks: &[u32]) -> usize {
        let mut parent = 0u64;
        let mut fresh = 0;
        self.clock += 1;
        for (i, chunk) in tokens.chunks_exact(PAGE_BLK).enumerate() {
            if i >= blocks.len() {
                break;
            }
            let h = chain_hash(parent, chunk);
            match self.map.get_mut(&h) {
                Some(e) if e.chunk == chunk => {
                    e.last_used = self.clock;
                }
                Some(_) => break, // 64-bit collision with a different chunk: stop caching here
                None => {
                    self.map.insert(
                        h,
                        CacheEntry {
                            block: blocks[i],
                            chunk: chunk.to_vec(),
                            parent,
                            last_used: self.clock,
                        },
                    );
                    fresh += 1;
                }
            }
            parent = h;
        }
        fresh
    }

    /// The set of block ids the index actually owns. A sequence that PUBLISHED a chunk whose
    /// hash was already present kept its own duplicate block — `published` alone cannot tell the
    /// two apart, and treating the duplicate as radix-owned on release leaked it forever (the
    /// 16-way identical-page bench drained the whole pool into a scheduler livelock).
    pub fn owned_blocks(&self) -> std::collections::HashSet<u32> {
        self.map.values().map(|e| e.block).collect()
    }

    /// Evict least-recently-used unreferenced entries until `want` blocks are free (or nothing
    /// evictable remains). Children keep working: a child entry whose parent is evicted can no
    /// longer be REACHED by match_prefix (the walk stops at the gap), so evict leaf-first by
    /// preferring entries no other entry names as parent.
    pub fn evict(&mut self, pool: &mut KvPool, want: usize) {
        while pool.available() < want {
            let parents: std::collections::HashSet<u64> =
                self.map.values().map(|e| e.parent).collect();
            let victim = self
                .map
                .iter()
                .filter(|(h, e)| pool.refs[e.block as usize] == 0 && !parents.contains(h))
                .min_by_key(|(_, e)| e.last_used)
                .map(|(h, _)| *h);
            let Some(h) = victim else { break };
            let e = self.map.remove(&h).expect("victim vanished");
            pool.release(e.block);
        }
    }
}

impl Default for RadixIndex {
    fn default() -> Self {
        Self::new()
    }
}

/// Per-request sampling: `temperature == 0` ⇒ greedy (the deterministic default). With a seed,
/// sampling is REPRODUCIBLE UNDER ANY LOAD: logits are batch-invariant (the bitwise gates) and
/// the per-request RNG stream depends only on the seed — the same request replays identically no
/// matter what traffic it shared batches with. Neither vLLM nor SGLang can promise that.
#[derive(Clone, Copy, Debug)]
pub struct SamplingParams {
    pub temperature: f32,
    pub top_p: f32,
    pub seed: u64,
}

impl Default for SamplingParams {
    fn default() -> Self {
        Self {
            temperature: 0.0,
            top_p: 1.0,
            seed: 0,
        }
    }
}

/// Per-request parameters beyond the `Copy` [`SamplingParams`] core: the Vec-bearing options
/// (stop token ids, penalties, logit bias) and logprobs config. Kept as a separate owned struct so
/// [`SamplingParams`] stays `Copy` and the hot chained-decode gate can inspect the scalar sampling
/// knobs without cloning.
#[derive(Clone, Debug)]
pub struct RequestParams {
    pub sampling: SamplingParams,
    /// Token ids that terminate generation like EOS (vLLM's `stop_token_ids`); the matched id is
    /// reported as the finish reason.
    pub stop_token_ids: Vec<u32>,
    /// OpenAI presence penalty (subtracted once per emitted token id).
    pub presence_penalty: f32,
    /// OpenAI frequency penalty (subtracted ∝ emitted count).
    pub frequency_penalty: f32,
    /// HF/vLLM repetition penalty (multiplicative over prompt ∪ emitted); `1.0` = off.
    pub repetition_penalty: f32,
    /// Per-token `(id, bias)` additions, clamped to ±100 at application.
    pub logit_bias: Vec<(u32, f32)>,
    /// Restrict sampling to the `k` highest-probability tokens.
    pub top_k: Option<usize>,
    /// Drop tokens below `min_p × max_prob` before sampling.
    pub min_p: Option<f32>,
    /// Number of top logprobs to attach per emitted token (`Some(0)` = the chosen token only).
    pub logprobs: Option<usize>,
}

impl Default for RequestParams {
    fn default() -> Self {
        Self {
            sampling: SamplingParams::default(),
            stop_token_ids: Vec::new(),
            presence_penalty: 0.0,
            frequency_penalty: 0.0,
            // The multiplicative penalty's identity is 1.0, NOT the derived 0.0.
            repetition_penalty: 1.0,
            logit_bias: Vec::new(),
            top_k: None,
            min_p: None,
            logprobs: None,
        }
    }
}

impl RequestParams {
    /// Greedy with no per-column CPU post-processing ⇒ eligible for the on-GPU chained decode.
    /// Stop token ids and `n` stay chain-safe (checked post-hoc on the read-back tokens); penalties,
    /// logit bias and logprobs force the per-column read-back path. `top_k`/`min_p` are no-ops under
    /// greedy, so they do not disqualify chaining.
    pub fn chain_safe(&self) -> bool {
        self.sampling.temperature == 0.0
            && self.presence_penalty == 0.0
            && self.frequency_penalty == 0.0
            && self.repetition_penalty == 1.0
            && self.logit_bias.is_empty()
            && self.logprobs.is_none()
    }
}

/// Why a sequence stopped generating, decided in the scheduler. The serving layer maps this to the
/// OpenAI `finish_reason` string (`Eos`/`StopTokenId` ⇒ `"stop"`, `Length` ⇒ `"length"`) and to
/// the vendor `stop_reason` detail (the specific token id for `StopTokenId`).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum FinishReason {
    /// Natural end-of-sequence (an `eos` token id).
    Eos,
    /// Hit the request's `max_tokens` budget.
    Length,
    /// Emitted one of the request's `stop_token_ids`.
    StopTokenId(u32),
}

/// Lifecycle of one request inside the scheduler.
enum SeqState {
    /// Prefilling: next chunk starts at `done` (tokens 0..done already ran or were cache hits).
    Prefill {
        done: usize,
    },
    /// Decoding: last emitted token is `tokens.last()`.
    Decode,
    Finished,
}

/// One admitted sequence occupying a slot (block-table row `slot + 1`).
/// The vision half of a multimodal request: what the tower produced, and how it maps onto the
/// prompt's tokens. Built by [`crate::vision::prepare_prompt`].
#[derive(Clone, Debug)]
pub struct VisionPrompt {
    /// `[rows, hidden]` — the tower's merged image tokens, all images concatenated in prompt order.
    pub embeds: Vec<f32>,
    /// Per PROMPT token: its 3-D M-RoPE position.
    pub mpos: Vec<[u32; 3]>,
    /// Per PROMPT token: which row of `embeds` replaces its token embedding, or `-1` for text.
    pub embed_index: Vec<i32>,
    /// The position the first GENERATED token takes. NOT `prompt_len`: an image advances the
    /// position counter by `max(llm_h, llm_w)`, not by its token count.
    pub next_pos: u32,
}

/// A sequence's vision state once admitted: the host rows are uploaded to the engine's arena and
/// dropped, leaving the mapping the column builder needs.
struct SeqVision {
    mpos: Vec<[u32; 3]>,
    embed_index: Vec<i32>,
    next_pos: u32,
    /// Base row this sequence owns in the engine's image arena (assigned at admission).
    row0: usize,
    rows: usize,
    /// Host-side rows, uploaded to the GPU at (every) admission and KEPT for the sequence's
    /// lifetime: recompute-preemption frees the arena rows, and the replay re-admission must
    /// re-upload them (~6 MB per live multimodal sequence).
    pending: Option<Vec<f32>>,
}

struct Sequence {
    id: u64,
    slot: usize,
    tokens: Vec<u32>,
    prompt_len: usize,
    /// Multimodal requests only; `None` is an ordinary text sequence and every code path below
    /// behaves exactly as it did before vision existed.
    vision: Option<SeqVision>,
    /// Physical blocks backing tokens (last one may be partial → private, uncached).
    blocks: Vec<u32>,
    /// How many leading blocks came from the radix cache (referenced, not owned).
    cached_blocks: usize,
    /// Blocks already inserted into the radix index (avoid re-inserting each step).
    published: usize,
    /// How many table entries the GPU block-table row already holds (rewrite only on growth).
    table_synced: usize,
    state: SeqState,
    max_tokens: usize,
    emitted: Vec<u32>,
    params: RequestParams,
    /// Radix-cache prefix hit for this sequence's prompt (reported as `cached_tokens` in usage).
    cached_prompt_tokens: usize,
    rng: u64,
    /// Serving-MTP: drafts awaiting the next verify span (len = the depth at the round that
    /// drafted them; may differ from the current occupancy-gated depth across a threshold).
    pending_drafts: Vec<u32>,
    /// Parked for a spec verify span — the plain decode packer must skip it.
    drafts_ready: bool,
}

/// One sequence's verify span within a serving-MTP spec round: `base` = the span head's column
/// index, followed by the `drafts` columns.
struct SpanOwner {
    slot: usize,
    base: usize,
    drafts: Vec<u32>,
}

/// A surviving sequence's re-draft request, recorded while its emitting column is at hand:
/// seed the draft chain from column `seed` of the plan that ran, chain from `req`.
struct Rearm {
    req: DraftReq,
    seed: usize,
    slot: usize,
    /// Token budget left (`max_tokens - emitted`) — a span must fit inside it to be worth arming.
    remaining: usize,
}

/// Scheduler statistics (served over `/metrics`).
#[derive(Default, Debug, Clone)]
pub struct ServeStats {
    pub steps: u64,
    pub tokens_out: u64,
    pub prefill_tokens: u64,
    pub cache_hit_tokens: u64,
    pub admitted: u64,
    pub finished: u64,
    pub rejected_full: u64,
    pub cancelled: u64,
    /// Running sequences preempted (recompute-evicted under KV-pool pressure) and requeued.
    pub preempted: u64,
    /// Zero-decode steps that ran the wide (KC=16) prefill plan (split-fuse).
    pub wide_steps: u64,
    /// Serving-MTP verify rounds run.
    pub spec_rounds: u64,
    /// Draft tokens verified / accepted — `spec_accepted / spec_drafted` is the acceptance rate,
    /// the go/no-go number for speculation on a given workload.
    pub spec_drafted: u64,
    pub spec_accepted: u64,
    /// Rounds whose drafts came from the prompt lookup instead of the MTP head.
    pub pld_hits: u64,
    /// Cumulative µs inside `batch_step` for wide (prefill) and non-wide (decode band) steps —
    /// with wall time and the tower counter this attributes an under-load page without log
    /// scraping: slack = wall − (wide + decode) − readback/sampling.
    pub wide_us: u64,
    pub decode_us: u64,
    /// Cumulative µs for the WHOLE step() call — minus wide_us+decode_us this is the
    /// scheduler's CPU overhead (admit, plan build, sampling, readback handling).
    pub step_us: u64,
    /// Step-CPU split: before batch_step (admit + cols/plan build) and after it (emission
    /// bookkeeping, per-token block growth, readback handling) — the pair that names WHICH
    /// side of the GPU call the scheduler overhead lives on.
    pub build_us: u64,
    pub post_us: u64,
    /// Chained-decode fast path: cumulative µs inside `batch_step_chained` and how many of
    /// `steps` came from chains — the path early-returns before the plain-step timers, which
    /// hid ~10s/window of REAL decode GPU work behind an apparent "scheduler CPU" mystery.
    pub chain_us: u64,
    pub chained_steps: u64,
}

/// Continuous-batching scheduler over one engine + one batch plan.
///
/// Not `Sync`: owns GPU work; drive it from one thread (the server binary's engine thread) and
/// feed it via channels.
pub struct Scheduler {
    /// The FULL-width decode plan (`k` columns).
    plan: BatchPlan,
    /// Narrower decode plans, built lazily and keyed by width.
    ///
    /// A plan's dispatch grid is sized by its width, so a step run on the k-wide plan pays for `k`
    /// columns even when one sequence is live. `serve` defaults `--batch` to 8, so a lone request —
    /// which is every claim-extraction — was paying for 7 idle columns on every token it decoded:
    /// **6.5 tok/s at width 8 versus 18.8 at width 1**. The per-column results are bitwise-equal
    /// across widths (`batched_serving_is_bitwise_equal_to_solo_runs`), so picking the smallest
    /// plan that covers the live columns is free correctness-wise and 2.9× on latency.
    ///
    /// Widths are powers of two, so a fleet of long-lived sequences settles on a handful of plans
    /// rather than one per distinct column count.
    bands: std::collections::BTreeMap<usize, BatchPlan>,
    /// Wide (KC=16) prefill plan for zero-decode steps — `Some` only when `--max-batched-tokens`
    /// was set AND the adapter has subgroups. Its columns are bitwise-equal to `plan`'s, so a
    /// prompt prefilled wide continues into decode identically.
    wide: Option<BatchPlan>,
    /// Round-robin start slot for decode packing (fairness under the prefill reserve).
    scan_cursor: usize,
    pool: KvPool,
    radix: RadixIndex,
    /// Radix reuse only for attention-only stacks (recurrent state is not prefix-addressable).
    radix_enabled: bool,
    /// Live extents of the engine's image-embedding arena: `(row0, rows, seq_id)`.
    img_alloc: Vec<(usize, usize, u64)>,
    /// Rows the engine's arena holds (0 = text-only model).
    img_rows: usize,
    /// Model hidden width — needed to validate a vision payload at submit(), where there is no gpu.
    hidden: usize,
    slots: Vec<Option<Sequence>>,
    queue: std::collections::VecDeque<Sequence>,
    next_id: u64,
    eos: Vec<u32>,
    /// Per-sequence (prompt + generated) token cap (from the engine's [`crate::EngineOpts`]).
    max_seq_tokens: usize,
    /// `(id, cached_prompt_tokens)` for sequences admitted since the last [`Scheduler::drain_admissions`].
    admissions: Vec<(u64, usize)>,
    /// Serving-MTP draft engine (`OSFKB_SERVE_MTP` > 0 AND the checkpoint ships a draft head AND
    /// the stack is attention-only). `None` = plain decoding, bit-for-bit the pre-MTP scheduler.
    mtp: Option<MtpEngine>,
    /// Spec-verify plans by WIDTH (wide GEMV families + DECODE attention — see
    /// [`Lfm2Gpu::make_batch_plan_verify`]), built lazily like the decode bands. Width matters
    /// as much as it does for decode: a wide-family step costs ~its PLAN width regardless of
    /// live columns (measured 33 ms for a 5-column span on a 64-wide verify plan — 85% of the
    /// round). Below one KC16 tile (16 columns) the narrow decode bands serve instead, where
    /// the KC4 family wins anyway. Empty on adapters without subgroups.
    verify_bands: std::collections::BTreeMap<usize, BatchPlan>,
    /// Column budget of a verify round (env `OSFKB_SPEC_VERIFY_COLS`, default 32, capped by the
    /// wide width): more spans amortize a round, but cost grows with the band the columns land
    /// in — 32 measured as the knee for 8-way spans of 5.
    spec_verify_cols: usize,
    /// Base draft depth. The per-round depth is occupancy-gated — see [`Self::spec_depth_for`].
    spec_k: usize,
    /// Occupancy schedule (envs `OSFKB_SPEC_C_FULL` / `OSFKB_SPEC_C_HALF`): decode concurrency ≤
    /// `spec_c_full` runs the base depth, ≤ `spec_c_half` runs `min(base, 4)`, above it
    /// speculation drains (verify cost scales with batch × span).
    spec_c_full: usize,
    spec_c_half: usize,
    /// Prompt-lookup drafting (`OSFKB_PLD`, default on; `0` disables): replace a round's MTP
    /// drafts with the continuation of the longest matching suffix n-gram of the sequence's own
    /// prompt+history. Free acceptance on copy/echo spans; wrong drafts only cost speed.
    pld: bool,
    /// The previous step was a spec verify round — used to alternate with plain rounds while
    /// prefill is pending, so redrafting sequences cannot starve admissions.
    last_spec: bool,
    /// The previous step was a FORCED wide prefill launch — the alternation bit of split-fuse v2.
    ///
    /// v1 packed prompt chunks into the narrow decode step's leftover columns, so under
    /// continuous decode a 968-token vision prompt trickled in at 2-3 columns per ~7 ms step
    /// (~0.36 tok/ms) while the wide plan — 64 columns per ~27 ms step, 1.9 tok/ms — sat unused
    /// except on zero-decode steps. v2 alternates: while any slot is prefilling (and the wide
    /// plan exists), every other step is a dedicated wide prefill launch and decode steps keep
    /// their FULL width (no reserve). Prompts ingest ~5× faster; decoders run every other step
    /// during prefill phases and full-rate otherwise.
    last_wide_prefill: bool,
    pub stats: ServeStats,
}

/// A progressing or terminal emission returned by [`Scheduler::step`]. `finish` is `Some` on the
/// last token of a sequence (and only then), carrying the reason it stopped — this replaces the
/// old caller-side `ids.is_empty()` guess with the scheduler's ground truth.
pub struct Emission {
    pub id: u64,
    pub token: u32,
    pub finish: Option<FinishReason>,
    /// Per-token logprobs (the chosen token + top alternatives), present only when the request asked
    /// for logprobs — which forces the per-column read-back path, never the chained decode.
    pub logprobs: Option<TokenLogprobs>,
}

impl Emission {
    /// Whether this emission terminates its sequence.
    pub fn done(&self) -> bool {
        self.finish.is_some()
    }
}

impl Scheduler {
    pub fn new(gpu: &Lfm2Gpu, ctx: &GpuCtx, k: usize, eos: Vec<u32>) -> Self {
        Self::new_with_wide(gpu, ctx, k, 0, eos)
    }

    /// The narrowest decode plan that can carry `n` live columns, building it on first use.
    ///
    /// Rounded up to a power of two so that a workload whose live count wobbles (7, 8, 6, …) reuses
    /// one plan instead of thrashing new ones.
    fn decode_band(&self, n: usize) -> usize {
        n.max(1).next_power_of_two().min(self.plan.k)
    }

    fn ensure_band(&mut self, gpu: &Lfm2Gpu, ctx: &GpuCtx, band: usize) {
        if band < self.plan.k && !self.bands.contains_key(&band) {
            self.bands.insert(band, gpu.make_batch_plan(ctx, band));
        }
    }

    /// The plan for `band` — the full-width one when the band IS the full width.
    fn band_plan(&self, band: usize) -> &BatchPlan {
        if band >= self.plan.k {
            &self.plan
        } else {
            self.bands
                .get(&band)
                .expect("band plan built by ensure_band")
        }
    }

    /// [`Self::new`] with a split-fuse wide-prefill budget `max_batched`: zero-decode steps pack up
    /// to `max_batched` prompt columns through the KC=16 wide plan. `max_batched == 0` or an adapter
    /// without subgroups ⇒ no wide plan (status quo — the narrow `k` plan serves every step).
    pub fn new_with_wide(
        gpu: &Lfm2Gpu,
        ctx: &GpuCtx,
        k: usize,
        max_batched: usize,
        eos: Vec<u32>,
    ) -> Self {
        let plan = gpu.make_batch_plan(ctx, k);
        // The engine's per-step `need_logit` mask is a u64, so `batch_step` handles at most 64
        // columns per launch; the wide prefill budget is capped there.
        let wide_k = max_batched.min(MAX_BATCH_COLS);
        let wide = (wide_k > 0 && ctx.subgroups).then(|| gpu.make_batch_plan_wide(ctx, wide_k));
        let pool = KvPool::new(gpu.pool_blocks());
        // Recurrent (conv) state is not prefix-addressable — radix reuse is attention-only.
        let has_conv = gpu.w.cfg.layer_is_attn.iter().any(|a| !a);
        // Serving-MTP: draft depth from the env, engine only if the checkpoint shipped a head.
        // ATTENTION-ONLY: KV rolls back positionally for free on a rejected draft, but recurrent
        // state (DeltaNet/conv) advances irreversibly over the rejected columns — the shard path
        // snapshots and restores it per column; this scheduler does not, so a hybrid stack must
        // never speculate here.
        let mut spec_k = std::env::var("OSFKB_SERVE_MTP")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(0);
        if spec_k > 0 && has_conv {
            eprintln!(
                "serve-mtp: recurrent-state architecture — no per-column state rollback on the \
                 single-node scheduler; speculation disabled"
            );
            spec_k = 0;
        }
        // A verify span is 1 + depth columns and must fit one launch of the verify plan.
        let verify_budget = wide.as_ref().map(|w| w.k).unwrap_or(k);
        spec_k = spec_k.min(verify_budget.saturating_sub(1));
        let mtp = (spec_k > 0)
            .then(|| MtpEngine::new_matching(ctx, gpu))
            .flatten();
        match (&mtp, spec_k) {
            (Some(_), _) => eprintln!(
                "serve-mtp: on (depth {spec_k}, verify budget {verify_budget} cols{})",
                if wide.is_some() { ", wide plan" } else { "" }
            ),
            (None, k) if k > 0 => {
                eprintln!("serve-mtp: requested but the checkpoint ships no draft head — plain decode")
            }
            _ => {}
        }
        Self {
            plan,
            bands: std::collections::BTreeMap::new(),
            wide,
            scan_cursor: 0,
            pool,
            radix: RadixIndex::new(),
            radix_enabled: !has_conv,
            img_alloc: Vec::new(),
            img_rows: gpu.image_rows(),
            hidden: gpu.w.cfg.hidden,
            slots: (0..MAX_SLOTS).map(|_| None).collect(),
            queue: std::collections::VecDeque::new(),
            next_id: 1,
            eos,
            max_seq_tokens: gpu.max_seq_tokens(),
            admissions: Vec::new(),
            mtp,
            verify_bands: std::collections::BTreeMap::new(),
            spec_verify_cols: std::env::var("OSFKB_SPEC_VERIFY_COLS")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(32),
            spec_k,
            spec_c_full: std::env::var("OSFKB_SPEC_C_FULL")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(4),
            spec_c_half: std::env::var("OSFKB_SPEC_C_HALF")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(8),
            pld: std::env::var("OSFKB_PLD").ok().as_deref() != Some("0"),
            last_spec: false,
            last_wide_prefill: false,
            stats: ServeStats::default(),
        }
    }

    /// Live occupancy: ADMITTED sequences (prefilling or decoding) — the concurrency any new
    /// speculation will share the GPU with. Prefilling sequences count because they flip to
    /// Decode within a round or two.
    fn live_occupancy(&self) -> usize {
        self.slots.iter().filter(|s| s.is_some()).count()
    }

    /// The occupancy-gated draft depth: base at c ≤ full, `min(base, 4)` at c ≤ half, else 0.
    fn spec_depth_for(&self, c: usize) -> usize {
        if self.spec_k == 0 {
            0
        } else if c <= self.spec_c_full {
            self.spec_k
        } else if c <= self.spec_c_half {
            self.spec_k.min(4)
        } else {
            0
        }
    }

    /// Arm each re-arm record's sequence with depth-`k_next` drafts. PROMPT-LOOKUP FIRST: the
    /// lookup is a free CPU scan and on OCR/copy-heavy output it wins most spans (measured 85%
    /// of spans on a dense page), so the MTP chain — 4 sequential draft-model steps, each a GPU
    /// sync — runs only for the sequences the lookup missed. Associated fn over the disjoint
    /// scheduler fields so `plan` can borrow the plan fields while `slots`/`stats` mutate.
    #[allow(clippy::too_many_arguments)]
    fn arm_drafts(
        mtp: &MtpEngine,
        pld: bool,
        slots: &mut [Option<Sequence>],
        stats: &mut ServeStats,
        ctx: &GpuCtx,
        plan: &BatchPlan,
        mut rearms: Vec<Rearm>,
        k_next: usize,
    ) -> Result<()> {
        // A sequence too near its budget to profit from a span drains plainly: accepting k
        // drafts + the bonus token must fit `remaining`.
        rearms.retain(|r| r.remaining > k_next);
        if pld {
            rearms.retain(|r| {
                let seq = slots[r.slot].as_mut().expect("re-armed seq is live");
                let look = prompt_lookup_drafts(&seq.tokens, k_next);
                if look.is_empty() {
                    return true; // no lookup hit — the MTP chain drafts it below
                }
                seq.pending_drafts = look;
                seq.drafts_ready = true;
                stats.pld_hits += 1;
                false
            });
        }
        for chunk in rearms.chunks(MtpEngine::K_DRAFT_BATCH) {
            let seeds: Vec<usize> = chunk.iter().map(|r| r.seed).collect();
            let reqs: Vec<DraftReq> = chunk.iter().map(|r| r.req).collect();
            // Single-submission chain (one sync for the whole round) — the per-step-sync v1
            // chain measured as a net serving loss; see `draft_chain_round`.
            let drafts = mtp.draft_chain_round(ctx, plan, &seeds, &reqs, k_next)?;
            for (r, d) in chunk.iter().zip(drafts) {
                let seq = slots[r.slot].as_mut().expect("re-armed seq is live");
                seq.pending_drafts = d;
                seq.drafts_ready = true;
            }
        }
        Ok(())
    }

    /// Write a block-table row on the MAIN engine and MIRROR it into the draft engine (the draft
    /// model keeps per-slot KV of its own, addressed by the same block ids). Associated fn so it
    /// can be called while a sequence is mutably borrowed out of `self.slots`.
    fn write_btab_mirror(
        mtp: Option<&MtpEngine>,
        gpu: &Lfm2Gpu,
        ctx: &GpuCtx,
        row: u32,
        blocks: &[u32],
    ) {
        gpu.write_btab_row(ctx, row, blocks);
        if let Some(m) = mtp {
            m.gpu.write_btab_row(ctx, row, blocks);
        }
    }

    /// Enqueue a request; returns its id. Tokens are prompt token ids (caller tokenizes).
    pub fn submit(&mut self, prompt: Vec<u32>, max_tokens: usize) -> Result<u64> {
        self.submit_with(prompt, max_tokens, SamplingParams::default())
    }

    /// [`Self::submit`] with sampling parameters (greedy when `temperature == 0`).
    pub fn submit_with(
        &mut self,
        prompt: Vec<u32>,
        max_tokens: usize,
        sampling: SamplingParams,
    ) -> Result<u64> {
        self.submit_with_params(
            prompt,
            max_tokens,
            RequestParams {
                sampling,
                ..Default::default()
            },
        )
    }

    /// [`Self::submit`] with the full per-request parameter bundle (sampling + stop token ids +,
    /// once wired, penalties/bias/logprobs). Returns the request id.
    pub fn submit_with_params(
        &mut self,
        prompt: Vec<u32>,
        max_tokens: usize,
        params: RequestParams,
    ) -> Result<u64> {
        anyhow::ensure!(!prompt.is_empty(), "empty prompt");
        anyhow::ensure!(
            prompt.len() + max_tokens <= self.max_seq_tokens,
            "request of {} tokens exceeds the {}-token per-sequence cap",
            prompt.len() + max_tokens,
            self.max_seq_tokens
        );
        // Reject a request that can never fit the KV pool even running alone, up front — otherwise
        // it would loop through admit() forever (rejected_full) and stall the scheduler.
        let need_blocks = (prompt.len() + max_tokens).div_ceil(PAGE_BLK);
        if need_blocks > self.pool.capacity() {
            self.stats.rejected_full += 1;
            anyhow::bail!(
                "request needs {need_blocks} KV blocks but the pool holds only {} (raise --vram-gb / --kv-fraction, or lower max_tokens)",
                self.pool.capacity()
            );
        }
        let id = self.next_id;
        self.next_id += 1;
        let rng = params.sampling.seed ^ 0x9e37_79b9_7f4a_7c15;
        self.queue.push_back(Sequence {
            id,
            slot: usize::MAX,
            tokens: prompt.clone(),
            prompt_len: prompt.len(),
            blocks: Vec::new(),
            cached_blocks: 0,
            published: 0,
            table_synced: 0,
            state: SeqState::Prefill { done: 0 },
            vision: None,
            max_tokens,
            emitted: Vec::new(),
            params,
            cached_prompt_tokens: 0,
            rng,
            pending_drafts: Vec::new(),
            drafts_ready: false,
        });
        Ok(id)
    }

    /// [`Self::submit_with_params`] for a MULTIMODAL prompt: the placeholder tokens in `prompt` are
    /// replaced, before layer 0, by the rows the vision tower produced, and positions come from
    /// `vision.mpos` rather than the token index.
    pub fn submit_with_vision(
        &mut self,
        prompt: Vec<u32>,
        max_tokens: usize,
        params: RequestParams,
        vision: VisionPrompt,
    ) -> Result<u64> {
        let hidden = self.hidden;
        anyhow::ensure!(
            vision.mpos.len() == prompt.len() && vision.embed_index.len() == prompt.len(),
            "vision prompt describes {} / {} tokens but the prompt has {}",
            vision.mpos.len(),
            vision.embed_index.len(),
            prompt.len()
        );
        anyhow::ensure!(
            vision.embeds.len().is_multiple_of(hidden),
            "vision embeds are {} floats, not a multiple of hidden ({hidden})",
            vision.embeds.len()
        );
        let rows = vision.embeds.len() / hidden;
        let referenced = vision.embed_index.iter().filter(|&&i| i >= 0).count();
        anyhow::ensure!(
            referenced == rows,
            "{referenced} prompt tokens want an image row but the tower produced {rows}"
        );
        anyhow::ensure!(
            rows <= self.img_rows,
            "this request needs {rows} image rows; the arena holds {} (raise OSFKB_IMG_TOKENS)",
            self.img_rows
        );

        let id = self.submit_with_params(prompt, max_tokens, params)?;
        let seq = self
            .queue
            .back_mut()
            .expect("submit_with_params just queued it");
        seq.vision = Some(SeqVision {
            mpos: vision.mpos,
            embed_index: vision.embed_index,
            next_pos: vision.next_pos,
            row0: 0, // assigned at admission
            rows,
            pending: Some(vision.embeds),
        });
        Ok(id)
    }

    /// The 3-D rope position of `seq`'s token at LINEAR position `p`.
    ///
    /// Text (and every generated token) is uniform, so this is `[p; 3]` and the rope is bit-identical
    /// to the 1-D path. A prompt token of a multimodal request takes its position from the vision
    /// plan; a GENERATED token resumes at `next_pos`.
    ///
    /// The vision-plan boundary is `v.mpos.len()` — the ORIGINAL prompt length — never
    /// `seq.prompt_len`: recompute-preemption folds already-generated tokens into `prompt_len`
    /// for the replay, and those positions are text (indexing `v.mpos` with them panicked on the
    /// first preempted vision sequence).
    fn mpos_at(seq: &Sequence, p: usize) -> [u32; 3] {
        match &seq.vision {
            None => [p as u32; 3],
            Some(v) if p < v.mpos.len() => v.mpos[p],
            Some(v) => [v.next_pos + (p - v.mpos.len()) as u32; 3],
        }
    }

    /// The arena row that replaces `seq`'s token embedding at linear position `p` (`-1` = text).
    /// Same folded-replay boundary as [`Self::mpos_at`].
    fn embed_row_at(seq: &Sequence, p: usize) -> i32 {
        match &seq.vision {
            Some(v) if p < v.embed_index.len() && v.embed_index[p] >= 0 => {
                v.row0 as i32 + v.embed_index[p]
            }
            _ => -1,
        }
    }

    /// First-fit allocation in the image arena. Sequences are few and short-lived, so a linear scan
    /// over the live extents beats any structure worth maintaining.
    fn alloc_image_rows(&mut self, id: u64, rows: usize) -> Option<usize> {
        if rows == 0 {
            return Some(0);
        }
        let mut taken: Vec<(usize, usize)> =
            self.img_alloc.iter().map(|&(s, n, _)| (s, n)).collect();
        taken.sort_unstable();
        let mut at = 0usize;
        for (s, n) in taken {
            if at + rows <= s {
                break;
            }
            at = at.max(s + n);
        }
        (at + rows <= self.img_rows).then(|| {
            self.img_alloc.push((at, rows, id));
            at
        })
    }

    fn free_image_rows(&mut self, id: u64) {
        self.img_alloc.retain(|&(_, _, owner)| owner != id);
    }

    /// Cancel a request by id: dequeue it if still waiting, or release its running slot (freeing
    /// its KV blocks and any private radix entries immediately). Returns whether the id was found.
    /// Serves stop-sequence truncation, `stop_token_ids`, and client-disconnect cancellation.
    pub fn cancel(&mut self, id: u64) -> bool {
        if let Some(pos) = self.queue.iter().position(|s| s.id == id) {
            self.queue.remove(pos);
            self.stats.cancelled += 1;
            return true;
        }
        for slot in 0..self.slots.len() {
            if matches!(&self.slots[slot], Some(s) if s.id == id) {
                let mut seq = self.slots[slot].take().expect("slot occupied");
                self.release_sequence_blocks(&mut seq);
                self.stats.cancelled += 1;
                return true;
            }
        }
        false
    }

    /// Drain the `(id, cached_prompt_tokens)` records of sequences admitted since the last call.
    /// The serving layer uses the cached count to report `prompt_tokens_details.cached_tokens`.
    pub fn drain_admissions(&mut self) -> Vec<(u64, usize)> {
        std::mem::take(&mut self.admissions)
    }

    pub fn pending(&self) -> usize {
        self.queue.len() + self.slots.iter().flatten().count()
    }

    /// Ensure `seq.blocks` covers positions `0..need` (allocating fresh blocks), evicting cache
    /// entries under pressure. Returns false if the pool is exhausted even after eviction.
    fn ensure_blocks(
        pool: &mut KvPool,
        radix: &mut RadixIndex,
        seq: &mut Sequence,
        need: usize,
    ) -> bool {
        let want = need.div_ceil(PAGE_BLK);
        while seq.blocks.len() < want {
            if pool.available() == 0 {
                radix.evict(pool, 1);
            }
            match pool.alloc() {
                Some(b) => seq.blocks.push(b),
                None => return false,
            }
        }
        true
    }

    /// Admit queued sequences into free slots (radix match + first block allocation).
    fn admit(&mut self, gpu: &Lfm2Gpu, ctx: &GpuCtx) {
        while !self.queue.is_empty() {
            let Some(slot) = self.slots.iter().position(Option::is_none) else {
                return;
            };
            let mut seq = self.queue.pop_front().expect("queue non-empty");
            seq.slot = slot;
            // Vision payload: claim arena rows and push the tower's output to the GPU. Done HERE,
            // not at submit(), because this is where `gpu`/`ctx` exist.
            if let Some(v) = seq.vision.as_mut() {
                let Some(row0) = ({
                    // `alloc_image_rows` needs &mut self while `seq` is detached from the queue —
                    // it is, so this is fine, but keep the borrow tight.
                    let rows = v.rows;
                    let id = seq.id;
                    self.alloc_image_rows(id, rows)
                }) else {
                    // Arena full: the images of the sequences already running have it. Put this one
                    // back and stop admitting — the same discipline as a full KV pool.
                    self.stats.rejected_full += 1;
                    seq.slot = usize::MAX;
                    self.queue.push_front(seq);
                    return;
                };
                let v = seq.vision.as_mut().expect("just checked");
                v.row0 = row0;
                // Upload from a KEPT copy, not a take(): recompute-preemption frees the arena
                // rows and re-admits through this same path, which must be able to re-upload —
                // a drop here left a preempted vision sequence replaying over another request's
                // (or stale) arena rows. ~6 MB host RAM per live multimodal sequence.
                if let Some(rows) = v.pending.as_deref()
                    && let Err(e) = gpu.upload_image_embeds(ctx, row0, rows)
                {
                    // Nothing sane to do but drop the request: proceeding would run the model
                    // over whatever stale rows the arena happens to hold.
                    eprintln!("[vision] upload failed for seq {}: {e:#}", seq.id);
                    self.free_image_rows(seq.id);
                    self.stats.cancelled += 1;
                    continue;
                }
            }
            // Prefix reuse: full blocks only, and never the final prompt block even when the
            // prompt length is block-aligned — the LAST prompt position must be re-run so its
            // logits (the first generated token) exist. Cap the match at the last full block
            // strictly before position prompt_len - 1.
            //
            // NEVER for a multimodal prompt. The radix key is TOKEN IDS, and every image expands to
            // the same repeated placeholder id — so a photo of a cat and a scan of an invoice have
            // byte-identical prefixes and would share each other's KV blocks. (Today this is moot on
            // the shipped VLM: Qwen3.5 is hybrid, so `radix_enabled` is already false because
            // recurrent state is not prefix-addressable. It would NOT be moot on a pure-attention
            // VLM, and the failure would be silent and data-dependent — hence the explicit guard.)
            if self.radix_enabled && seq.vision.is_none() {
                let cap = (seq.prompt_len - 1) / PAGE_BLK * PAGE_BLK;
                let (blocks, mut matched) = self
                    .radix
                    .match_prefix(&mut self.pool, &seq.tokens[..cap.min(seq.tokens.len())]);
                matched = matched.min(cap);
                seq.blocks = blocks;
                seq.cached_blocks = seq.blocks.len();
                seq.published = seq.blocks.len();
                seq.state = SeqState::Prefill { done: matched };
                seq.cached_prompt_tokens = matched;
                self.stats.cache_hit_tokens += matched as u64;
            }
            // First working block(s).
            let need = (seq.prompt_len.min(matched_plus_one(&seq))).max(1);
            if !Self::ensure_blocks(&mut self.pool, &mut self.radix, &mut seq, need) {
                // Pool full: put it back and stop admitting.
                self.stats.rejected_full += 1;
                self.release_sequence_blocks(&mut seq);
                seq.slot = usize::MAX;
                self.queue.push_front(seq);
                return;
            }
            Self::write_btab_mirror(self.mtp.as_ref(), gpu, ctx, (slot + 1) as u32, &seq.blocks);
            seq.table_synced = seq.blocks.len();
            if !self.radix_enabled {
                // Fresh recurrent state for the slot (conv rings and/or DeltaNet S).
                gpu.zero_conv_slot(ctx, &self.plan, slot);
                gpu.zero_dn_slot(ctx, slot + 1);
            }
            self.stats.admitted += 1;
            self.admissions.push((seq.id, seq.cached_prompt_tokens));
            self.slots[slot] = Some(seq);
        }
    }

    fn release_sequence_blocks(&mut self, seq: &mut Sequence) {
        // Every path that gives a sequence's KV blocks back also gives back its image-arena rows —
        // one place, so a new retirement path cannot forget and slowly leak the arena until
        // multimodal requests start getting rejected for no visible reason.
        if seq.vision.is_some() {
            self.free_image_rows(seq.id);
        }
        // `i < seq.published` says the block was OFFERED to the radix index — not that the index
        // took it. A duplicate prefix (identical prompt in flight twice) keeps the FIRST block as
        // canonical, so the later sequence's offered blocks must come back to the free list here
        // or they leak with refcount 0 and no owner.
        let owned = self.radix.owned_blocks();
        for (i, &b) in seq.blocks.iter().enumerate() {
            let now = self.pool.unref(b);
            let cached = i < seq.published && owned.contains(&b);
            if now == 0 && !cached {
                self.pool.release(b);
            }
        }
        seq.blocks.clear();
    }

    /// Recompute-preemption: evict the latest-admitted running sequence (highest id) that is NOT
    /// `current` and NOT already `columned` this step, releasing its KV blocks and requeuing it at
    /// the FRONT to replay ALL its tokens (prompt + generated). An attention model resumes cheaply —
    /// its published prefix blocks stay in the radix index for a cache re-hit; a conv/DeltaNet model
    /// re-zeros state and fully replays through the ordinary [`Self::admit`] path. Returns whether a
    /// victim was found (false ⇒ nothing left to preempt).
    fn preempt_latest_admitted(&mut self, current: usize, columned: &[usize]) -> bool {
        let victim = (0..self.slots.len())
            .filter(|s| *s != current && !columned.contains(s))
            .filter_map(|s| self.slots[s].as_ref().map(|seq| (s, seq.id)))
            .max_by_key(|&(_, id)| id)
            .map(|(s, _)| s);
        let Some(vs) = victim else {
            return false;
        };
        let mut seq = self.slots[vs].take().expect("victim slot occupied");
        // Release blocks first (this reads the OLD `published` to keep cached prefixes in the radix
        // index), then reset the sequence to replay from scratch.
        self.release_sequence_blocks(&mut seq);
        seq.slot = usize::MAX;
        seq.prompt_len = seq.tokens.len();
        seq.state = SeqState::Prefill { done: 0 };
        seq.cached_blocks = 0;
        seq.published = 0;
        seq.table_synced = 0;
        seq.cached_prompt_tokens = 0;
        seq.pending_drafts.clear();
        seq.drafts_ready = false;
        self.queue.push_front(seq);
        self.stats.preempted += 1;
        true
    }

    /// One scheduler iteration: build a batch (decode columns + one prefill chunk), run it,
    /// distribute tokens. Returns the emissions of this step.
    ///
    /// Steady state (nothing queued, every active sequence decoding) runs a CHAINED submit of up
    /// to [`CHAIN_MAX`] steps with on-GPU token feedback — one CPU sync per chain instead of per
    /// token. EOS/max-token overshoot inside a chain is truncated after the fact; the overshot
    /// positions are never attended by later steps (their sequences are gone), so soundness holds.
    pub fn step(&mut self, gpu: &Lfm2Gpu, ctx: &GpuCtx) -> Result<Vec<Emission>> {
        let _t_whole = std::time::Instant::now();
        let r = self.step_inner(gpu, ctx);
        self.stats.step_us += _t_whole.elapsed().as_micros() as u64;
        r
    }

    fn step_inner(&mut self, gpu: &Lfm2Gpu, ctx: &GpuCtx) -> Result<Vec<Emission>> {
        let _t_build = std::time::Instant::now();
        self.admit(gpu, ctx);
        // Serving-MTP verify round: sequences whose drafts are ready verify a (1 + k)-column
        // span each and emit every accepted token in one step. Falls through to the plain path
        // when nothing is ready (or when alternating with pending prefill).
        if self.mtp.is_some()
            && let Some(out) = self.step_spec(gpu, ctx)?
        {
            return Ok(out);
        }
        // Whatever runs below is a plain round — the alternation gate re-arms.
        self.last_spec = false;
        if self.queue.is_empty() {
            let active: Vec<usize> = self
                .slots
                .iter()
                .enumerate()
                .filter_map(|(i, s)| s.as_ref().map(|_| i))
                .collect();
            // The chained fast path pre-allocates a block per sequence for the round; only take it
            // when the pool has that headroom, so a tight pool de-chains into the per-column path
            // below where preemption can free blocks (chaining never runs a sequence short of KV).
            // Under serving-MTP the spec verify rounds ARE the multi-token mechanism — the
            // chained submit would bypass the draft/verify bookkeeping entirely.
            let all_decode = self.mtp.is_none()
                && !active.is_empty()
                && active.len() <= self.plan.k
                && self.pool.available() >= active.len()
                && active.iter().all(|&i| {
                    self.slots[i].as_ref().is_some_and(|s| {
                        matches!(s.state, SeqState::Decode) && s.params.chain_safe()
                    })
                });
            if all_decode {
                // Band the chained plan too: dispatch grids are sized by PLAN width, so one
                // active stream on the full k-wide plan pays k weight sweeps per step (measured
                // 6.6 vs 32.9 tok/s solo on the 27B). Same pow2 banding the per-column path uses.
                let band = self.decode_band(active.len());
                self.ensure_band(gpu, ctx, band);
                return self.step_chained(gpu, ctx, &active, band);
            }
        }
        let k = self.plan.k;
        let mut cols: Vec<BatchCol> = Vec::with_capacity(k);
        // (slot, kind) per column; kind: decode = emits, prefill-final = emits, prefill-mid = no.
        let mut owners: Vec<(usize, bool)> = Vec::new();

        let n_slots = self.slots.len();
        let prefill_pending = (0..n_slots).any(|s| {
            matches!(
                self.slots[s].as_ref().map(|q| &q.state),
                Some(SeqState::Prefill { .. })
            )
        });
        // Split-fuse v2 (see `last_wide_prefill`): with a wide plan, prompts NEVER trickle
        // through narrow leftovers — every other step while prefill is pending is a dedicated
        // wide launch (decode packing skipped ⇒ `use_wide` routes it), and decode steps keep the
        // full width. Without a wide plan, v1's reserve packing stands.
        let force_wide_prefill = self.wide.is_some() && prefill_pending && !self.last_wide_prefill;
        self.last_wide_prefill = force_wide_prefill;
        let decode_budget = if force_wide_prefill {
            0
        } else if prefill_pending && self.wide.is_none() {
            // v1 reserve: decode keeps the majority (k − k/4); prefill gets at least max(k/4, 1).
            k - (k / 4).max(1)
        } else {
            k
        };
        // Decode round-robin from `scan_cursor` for fairness under the reserve. Index-based so a
        // starved column can PREEMPT another sequence to free KV (S-C3) — never one already columned
        // this step (its KV is about to be read).
        let mut columned: Vec<usize> = Vec::new();
        let mut cursor = self.scan_cursor;
        for i in 0..n_slots {
            if cols.len() == decode_budget {
                break;
            }
            let slot = (self.scan_cursor + i) % n_slots;
            if !matches!(
                self.slots[slot].as_ref().map(|s| &s.state),
                Some(SeqState::Decode)
            ) {
                continue;
            }
            // Parked for a spec verify span (drafts pending) — never re-decode it plainly.
            if self
                .slots[slot]
                .as_ref()
                .is_some_and(|s| s.drafts_ready)
            {
                continue;
            }
            let pos = self.slots[slot].as_ref().expect("decode slot").tokens.len() as u32 - 1;
            let need = pos as usize + 1;
            let mut fits = false;
            loop {
                let ok = {
                    let seq = self.slots[slot].as_mut().expect("decode slot");
                    Self::ensure_blocks(&mut self.pool, &mut self.radix, seq, need)
                };
                if ok {
                    fits = true;
                    break;
                }
                if !self.preempt_latest_admitted(slot, &columned) {
                    break;
                }
            }
            if !fits {
                continue;
            }
            let seq = self.slots[slot].as_mut().expect("decode slot");
            if seq.blocks.len() > seq.table_synced {
                Self::write_btab_mirror(self.mtp.as_ref(), gpu, ctx, (slot + 1) as u32, &seq.blocks);
                seq.table_synced = seq.blocks.len();
            }
            cols.push(BatchCol {
                token: *seq.tokens.last().expect("decode has tokens"),
                pos,
                btrow: (slot + 1) as u32,
                need_logit: true,
                // A generated token is text, but on a multimodal sequence its position resumes at
                // `next_pos`, NOT at the token index — the image consumed fewer positions than slots.
                mpos: Self::mpos_at(seq, pos as usize),
                embed_row: -1,
            });
            owners.push((slot, true));
            columned.push(slot);
            cursor = (slot + 1) % n_slots;
        }
        if !cols.is_empty() {
            self.scan_cursor = cursor;
        }
        // Then prefill chunks from as MANY prefilling sequences as fit (split-fuse). A step with NO
        // decode work routes through the WIDE (KC=16) plan at its wider column budget — 4× weight
        // amortization and 4× fewer syncs on the prompt (TTFT) phase.
        let use_wide = cols.is_empty() && self.wide.is_some();
        let budget = match &self.wide {
            Some(w) if use_wide => w.k,
            _ => k,
        };
        // Serving-MTP: draft-KV prefill pairs collected while packing — the draft layer's KV must
        // cover the prompt or the chain attends over zeros and acceptance collapses. Unlike the
        // shard path, the NEXT token is known even at a chunk boundary (the whole prompt is
        // here), so the draft KV has no per-chunk holes.
        let mtp_on = self.mtp.is_some();
        let mut pairs: Vec<DraftPair> = Vec::new();
        // Prefill columns pack ONLY into wide launches when a wide plan exists (split-fuse v2 —
        // the alternation above guarantees one every other step while prompts are pending);
        // narrow-leftover trickle is the no-wide-plan fallback.
        for slot in 0..n_slots {
            if !(use_wide || self.wide.is_none()) {
                break;
            }
            if cols.len() >= budget {
                break;
            }
            let Some(seq) = self.slots[slot].as_mut() else {
                continue;
            };
            if let SeqState::Prefill { done } = seq.state {
                let start = cols.len();
                let mut p = done;
                while cols.len() < budget && p < seq.prompt_len {
                    if !Self::ensure_blocks(&mut self.pool, &mut self.radix, seq, p + 1) {
                        break;
                    }
                    let last = p + 1 == seq.prompt_len;
                    cols.push(BatchCol {
                        token: seq.tokens[p],
                        pos: p as u32,
                        btrow: (slot + 1) as u32,
                        need_logit: last,
                        mpos: Self::mpos_at(seq, p),
                        embed_row: Self::embed_row_at(seq, p),
                    });
                    owners.push((slot, last));
                    if mtp_on && !last && seq.params.chain_safe() {
                        pairs.push(DraftPair {
                            btrow: (slot + 1) as u32,
                            cur_col: cols.len() - 1,
                            token: seq.tokens[p + 1],
                            pos: p,
                            mpos: Self::mpos_at(seq, p + 1),
                        });
                    }
                    p += 1;
                }
                if seq.blocks.len() > seq.table_synced {
                    Self::write_btab_mirror(
                        self.mtp.as_ref(),
                        gpu,
                        ctx,
                        (slot + 1) as u32,
                        &seq.blocks,
                    );
                    seq.table_synced = seq.blocks.len();
                }
                if p < seq.prompt_len && std::env::var("OSFKB_SCHED_TRACE").is_ok() {
                    eprintln!(
                        "[sched] prefill stalled: done {done} -> {p} of {} (blocks {}, budget {budget}, cols {})",
                        seq.prompt_len,
                        seq.blocks.len(),
                        cols.len()
                    );
                }
                seq.state = if p == seq.prompt_len {
                    SeqState::Decode // flips to Decode when its final column emits below
                } else {
                    SeqState::Prefill { done: p }
                };
                self.stats.prefill_tokens += (cols.len() - start) as u64;
                // NO break — pack the whole prefill fringe (that was the one-sequence-per-step tail).
            }
        }

        if cols.is_empty() {
            return Ok(Vec::new());
        }
        // Zero-decode prefill steps run the WIDE plan; every other step the narrow plan. Its columns
        // are bitwise-equal, so which plan ran never changes a token — but `read_batch_logits` below
        // MUST target the plan that actually ran (sampled/penalized prefill-final columns).
        let use_wide = use_wide && !cols.is_empty();
        // `OSFKB_PLAN_TRACE=1` prints which plan each step actually ran, and how many columns it
        // packed. Twice now a prefill optimization has measured as a no-op, and the first question
        // both times was "is the code I changed even executing?" — this answers it in one run
        // instead of an afternoon. (It is: 29 WIDE steps × 64 columns for a 1,801-token prompt.)
        // Decode steps run the NARROWEST plan that covers the live columns; prefill steps run the
        // wide plan at its own width.
        let band = self.decode_band(cols.len());
        if !use_wide {
            self.ensure_band(gpu, ctx, band);
        }
        if std::env::var("OSFKB_PLAN_TRACE").is_ok() {
            eprintln!(
                "[plan] {} cols={} band={} (full k={}, wide k={:?})",
                if use_wide { "WIDE" } else { "narrow" },
                cols.len(),
                band,
                self.plan.k,
                self.wide.as_ref().map(|w| w.k)
            );
        }
        self.stats.build_us += _t_build.elapsed().as_micros() as u64;
        let _t_step = std::time::Instant::now();
        let chosen = if use_wide {
            self.stats.wide_steps += 1;
            let r = gpu.batch_step(ctx, self.wide.as_ref().expect("wide plan"), &cols)?;
            self.stats.wide_us += _t_step.elapsed().as_micros() as u64;
            r
        } else {
            let r = gpu.batch_step(ctx, self.band_plan(band), &cols)?;
            self.stats.decode_us += _t_step.elapsed().as_micros() as u64;
            r
        };
        // `OSFKB_DECODE_TRACE=1` splits a decoded token into GPU (`batch_step`) vs CPU (readback,
        // sampling) — which is the split that decides where decode work belongs. It is how the
        // vocabulary-sized sampler sort was found: the readback everyone suspects costs ~0.4 ms on a
        // quiet box, and the sampler nobody suspects cost 5-6 ms.
        if std::env::var("OSFKB_DECODE_TRACE").is_ok() {
            eprintln!(
                "[step] batch_step {} us (cols={} band={})",
                _t_step.elapsed().as_micros(),
                cols.len(),
                band
            );
        }
        self.stats.steps += 1;
        let _t_post = std::time::Instant::now();
        let r = self.step_post(gpu, ctx, use_wide, band, &cols, &owners, &chosen, pairs);
        self.stats.post_us += _t_post.elapsed().as_micros() as u64;
        r
    }

    #[allow(clippy::too_many_arguments)]
    fn step_post(
        &mut self,
        gpu: &Lfm2Gpu,
        ctx: &GpuCtx,
        use_wide: bool,
        band: usize,
        cols: &[BatchCol],
        owners: &[(usize, bool)],
        chosen: &[u32],
        pairs: Vec<DraftPair>,
    ) -> Result<Vec<Emission>> {
        let mtp_on = self.mtp.is_some();
        let mut out = Vec::new();
        // Serving-MTP re-arm records: every emitting column of a surviving greedy sequence
        // seeds a fresh draft chain after the loop (the sequence's next step is a verify span).
        let mut rearms: Vec<Rearm> = Vec::new();
        for (ci, &(slot, emits)) in owners.iter().enumerate() {
            if !emits {
                continue;
            }
            let mut token = chosen[ci];
            let mut logprobs = None;
            let mut new_rng = None;
            // Non-chain-safe columns (sampling, penalties, logit bias, or logprobs) read THIS
            // column's logits and post-process on the CPU. Logits are batch-invariant and the RNG
            // stream is seed-only ⇒ the pick is reproducible under any load. Logprobs come from the
            // RAW column (the model's own probabilities), independent of the pick.
            {
                let seq = self.slots[slot].as_ref().expect("owner slot occupied");
                if !seq.params.chain_safe() {
                    let _t_rb = std::time::Instant::now();
                    // MUST read from the plan that actually ran — its logits buffer is the one
                    // that was written.
                    let raw = if use_wide {
                        gpu.read_batch_logits(ctx, self.wide.as_ref().expect("wide plan"), ci)?
                    } else {
                        gpu.read_batch_logits(ctx, self.band_plan(band), ci)?
                    };
                    let rb_us = _t_rb.elapsed().as_micros();
                    let _t_sm = std::time::Instant::now();
                    let mut work = raw.clone();
                    // Gemma-4 final-logit softcapping `s·tanh(z/s)` — HF applies it to the
                    // logits themselves, so every consumer here (penalties, sampling,
                    // logprobs) sees the capped values. Monotonic ⇒ the pure-greedy GPU-argmax
                    // fast path (which never reads logits back) is exempt by construction.
                    let cap = gpu.w.cfg.edge.final_logit_softcapping;
                    if cap > 0.0 {
                        for v in &mut work {
                            *v = cap * (*v / cap).tanh();
                        }
                    }
                    // Logprobs must be capped but UN-penalized: snapshot before penalties.
                    let capped_raw = (cap > 0.0 && seq.params.logprobs.is_some())
                        .then(|| work.clone());
                    let prompt = &seq.tokens[..seq.prompt_len];
                    sampling::apply_penalties(
                        &mut work,
                        prompt,
                        &seq.emitted,
                        seq.params.presence_penalty,
                        seq.params.frequency_penalty,
                        seq.params.repetition_penalty,
                    );
                    sampling::apply_logit_bias(&mut work, &seq.params.logit_bias);
                    if seq.params.sampling.temperature > 0.0 {
                        let mut rng = seq.rng;
                        token = sampling::sample(
                            &work,
                            seq.params.sampling.temperature,
                            seq.params.top_k,
                            seq.params.sampling.top_p,
                            seq.params.min_p,
                            &mut rng,
                        );
                        new_rng = Some(rng);
                    } else {
                        token = sampling::argmax(&work);
                    }
                    if let Some(n) = seq.params.logprobs {
                        logprobs = Some(sampling::log_softmax_at(
                            capped_raw.as_deref().unwrap_or(&raw),
                            token,
                            n,
                        ));
                    }
                    if std::env::var("OSFKB_DECODE_TRACE").is_ok() {
                        eprintln!(
                            "[decode] readback {} us | sample {} us",
                            rb_us,
                            _t_sm.elapsed().as_micros()
                        );
                    }
                }
            }
            if let Some(rng) = new_rng {
                self.slots[slot].as_mut().expect("slot").rng = rng;
            }
            let seq = self.slots[slot].as_mut().expect("owner slot occupied");
            seq.tokens.push(token);
            seq.emitted.push(token);
            self.stats.tokens_out += 1;
            let finish = decide_finish(&self.eos, seq, token);
            out.push(Emission {
                id: seq.id,
                token,
                finish,
                logprobs,
            });
            seq.state = if finish.is_some() {
                SeqState::Finished
            } else {
                SeqState::Decode
            };
            if mtp_on && finish.is_none() && seq.params.chain_safe() {
                rearms.push(Rearm {
                    req: DraftReq {
                        btrow: (slot + 1) as u32,
                        first_token: token,
                        pos0: cols[ci].pos as usize,
                        mpos0: cols[ci].mpos,
                    },
                    seed: ci,
                    slot,
                    remaining: seq.max_tokens.saturating_sub(seq.emitted.len()),
                });
            }
            // Publish freshly-FILLED full blocks to the radix index (immutable from now on).
            // Same guard as the admit-side match: NEVER publish a multimodal prompt — image
            // placeholder ids make different images byte-identical, so cached vision blocks
            // both poison the index and (being duplicates) leak on release.
            if self.radix_enabled && seq.vision.is_none() {
                let full = seq.tokens.len() / PAGE_BLK;
                if full > seq.published {
                    self.radix
                        .insert(&seq.tokens[..full * PAGE_BLK], &seq.blocks[..full]);
                    seq.published = full;
                }
            }
        }
        // Serving-MTP: extend the draft KV over this step's prefill columns, then chain fresh
        // drafts from every emitting survivor. The plan that RAN this step is the one whose
        // `cur` holds the pre-final-norm residuals the merge/seed copies read.
        if mtp_on && (!pairs.is_empty() || !rearms.is_empty()) {
            let plan = if use_wide {
                self.wide.as_ref().expect("wide plan")
            } else if band >= self.plan.k {
                &self.plan
            } else {
                self.bands.get(&band).expect("band plan built by ensure_band")
            };
            let m = self.mtp.as_ref().expect("mtp_on");
            for chunk in pairs.chunks(MtpEngine::K_DRAFT_BATCH) {
                m.draft_pairs(ctx, plan, chunk)?;
            }
            let k_next = self.spec_depth_for(self.live_occupancy());
            if k_next > 0 {
                Self::arm_drafts(
                    m,
                    self.pld,
                    &mut self.slots,
                    &mut self.stats,
                    ctx,
                    plan,
                    rearms,
                    k_next,
                )?;
            }
        }
        self.reap_finished();
        Ok(out)
    }

    /// Release every finished sequence's slot + KV blocks (+ image rows).
    fn reap_finished(&mut self) {
        for slot in 0..self.slots.len() {
            let finished = matches!(
                self.slots[slot],
                Some(Sequence {
                    state: SeqState::Finished,
                    ..
                })
            );
            if finished {
                let mut seq = self.slots[slot].take().expect("finished seq");
                self.stats.finished += 1;
                self.release_sequence_blocks(&mut seq);
            }
        }
    }

    /// Serving-MTP verify round: pack a (1 + k)-column span per draft-ready sequence, run ONE
    /// wide-plan step, accept each span's longest draft prefix (plus the model's own next token),
    /// then seed + re-draft every survivor from its last accepted column. Returns `None` when no
    /// round was packed — the caller falls through to the plain path.
    ///
    /// KV soundness on rejection is positional: a rejected draft's K/V rows sit at positions the
    /// sequence has not reached, and the next span (or plain step) overwrites them before any
    /// attention reads that far. Attention-only stacks need nothing else — which is why spec is
    /// refused at construction for recurrent architectures.
    fn step_spec(&mut self, gpu: &Lfm2Gpu, ctx: &GpuCtx) -> Result<Option<Vec<Emission>>> {
        // Fairness: while prompts are waiting, alternate verify and plain rounds — a verified
        // sequence re-arms immediately and would otherwise win every step and starve admissions.
        let prefill_pending = !self.queue.is_empty()
            || self
                .slots
                .iter()
                .flatten()
                .any(|s| matches!(s.state, SeqState::Prefill { .. }));
        if prefill_pending && self.last_spec {
            return Ok(None);
        }
        let n_slots = self.slots.len();
        let ready =
            |q: &Sequence| matches!(q.state, SeqState::Decode) && q.drafts_ready && !q.pending_drafts.is_empty();
        // The round's verify length = the first ready sequence's drafted length in scan order
        // (uniform per round; every parked sequence is eventually first, so no depth starves).
        let Some(k_verify) = (0..n_slots)
            .map(|i| (self.scan_cursor + i) % n_slots)
            .find_map(|s| {
                self.slots[s]
                    .as_ref()
                    .filter(|q| ready(q))
                    .map(|q| q.pending_drafts.len())
            })
        else {
            return Ok(None);
        };
        let span = 1 + k_verify;
        let budget = self
            .wide
            .as_ref()
            .map(|w| w.k.min(self.spec_verify_cols.max(span)))
            .unwrap_or(self.plan.k);
        let max_spans = budget / span;
        if max_spans == 0 {
            // Cannot happen when the drafts came from this scheduler (depth is capped at build);
            // drop them defensively rather than stall parked forever.
            for seq in self.slots.iter_mut().flatten() {
                seq.pending_drafts.clear();
                seq.drafts_ready = false;
            }
            return Ok(None);
        }
        let mut cols: Vec<BatchCol> = Vec::new();
        let mut spans: Vec<SpanOwner> = Vec::new();
        let mut cursor = self.scan_cursor;
        for i in 0..n_slots {
            if spans.len() == max_spans {
                break;
            }
            let slot = (self.scan_cursor + i) % n_slots;
            let Some(seq) = self.slots[slot].as_mut() else {
                continue;
            };
            if !(matches!(seq.state, SeqState::Decode)
                && seq.drafts_ready
                && seq.pending_drafts.len() == k_verify)
            {
                continue;
            }
            let pos0 = seq.tokens.len() - 1;
            // Blocks for the whole span plus the NEXT round's head position.
            if !Self::ensure_blocks(&mut self.pool, &mut self.radix, seq, pos0 + span + 1) {
                // Pool starved: drop the drafts so the sequence rejoins the plain path, where
                // the preemption ladder can free blocks. Drafts are always discardable.
                seq.pending_drafts.clear();
                seq.drafts_ready = false;
                continue;
            }
            if seq.blocks.len() > seq.table_synced {
                Self::write_btab_mirror(self.mtp.as_ref(), gpu, ctx, (slot + 1) as u32, &seq.blocks);
                seq.table_synced = seq.blocks.len();
            }
            let base = cols.len();
            cols.push(BatchCol {
                token: *seq.tokens.last().expect("decode has tokens"),
                pos: pos0 as u32,
                btrow: (slot + 1) as u32,
                need_logit: true,
                mpos: Self::mpos_at(seq, pos0),
                embed_row: -1,
            });
            for (d, &t) in seq.pending_drafts.iter().enumerate() {
                let p = pos0 + 1 + d;
                cols.push(BatchCol {
                    token: t,
                    pos: p as u32,
                    btrow: (slot + 1) as u32,
                    need_logit: true,
                    mpos: Self::mpos_at(seq, p),
                    embed_row: -1,
                });
            }
            spans.push(SpanOwner {
                slot,
                base,
                drafts: std::mem::take(&mut seq.pending_drafts),
            });
            seq.drafts_ready = false;
            cursor = (slot + 1) % n_slots;
        }
        if spans.is_empty() {
            return Ok(None);
        }
        self.scan_cursor = cursor;
        // Plan for the round, sized to the LIVE columns (a step costs ~its plan width no matter
        // how few columns are real): pow2 verify band with the wide GEMV families at ≥ one KC16
        // tile, the narrow decode band below that (where KC16 pads and loses) or without
        // subgroups. Decode attention either way — the bitwise contract.
        let band = self.decode_band(cols.len());
        // Narrow bands cap at plan.k, so a round wider than the decode plan MUST take a verify
        // band (packing already kept cols ≤ plan.k when no wide plan exists); below one KC16
        // tile the narrow band wins, so verify bands start at 16.
        let use_verify =
            self.wide.is_some() && (cols.len() > self.plan.k || cols.len() >= 16);
        let vband = cols.len().next_power_of_two().max(16);
        if use_verify {
            self.verify_bands
                .entry(vband)
                .or_insert_with(|| gpu.make_batch_plan_verify(ctx, vband));
        } else {
            self.ensure_band(gpu, ctx, band);
        }
        if std::env::var("OSFKB_PLAN_TRACE").is_ok() {
            eprintln!(
                "[plan] SPEC spans={} k_verify={k_verify} cols={} plan={}",
                spans.len(),
                cols.len(),
                if use_verify { "verify" } else { "band" }
            );
        }
        let _t_verify = std::time::Instant::now();
        let chosen = if use_verify {
            gpu.batch_step(ctx, self.verify_bands.get(&vband).expect("just built"), &cols)?
        } else {
            gpu.batch_step(ctx, self.band_plan(band), &cols)?
        };
        let _t_verify_us = _t_verify.elapsed().as_micros();
        self.stats.steps += 1;
        self.stats.spec_rounds += 1;
        self.last_spec = true;
        let mut out = Vec::new();
        let mut rearms: Vec<Rearm> = Vec::new();
        for sp in &spans {
            let seq = self.slots[sp.slot].as_mut().expect("span slot occupied");
            let mut j = 0usize;
            while j < k_verify && chosen[sp.base + j] == sp.drafts[j] {
                j += 1;
            }
            self.stats.spec_drafted += k_verify as u64;
            self.stats.spec_accepted += j as u64;
            if std::env::var("OSFKB_SPEC_TRACE").is_ok() {
                eprintln!(
                    "[spec] slot {} pos0 {} accept {j}/{k_verify} drafts {:?} chosen {:?}",
                    sp.slot,
                    cols[sp.base].pos,
                    sp.drafts,
                    &chosen[sp.base..sp.base + 1 + k_verify]
                );
            }
            // Emit the accepted prefix plus the model's own token at the first divergence (on a
            // full accept that bonus is the span's last verify output).
            let mut finished = false;
            for d in 0..=j {
                if finished {
                    break;
                }
                let token = chosen[sp.base + d];
                seq.tokens.push(token);
                seq.emitted.push(token);
                self.stats.tokens_out += 1;
                let finish = decide_finish(&self.eos, seq, token);
                finished = finish.is_some();
                out.push(Emission {
                    id: seq.id,
                    token,
                    finish,
                    logprobs: None,
                });
            }
            if finished {
                seq.state = SeqState::Finished;
            } else {
                let col = sp.base + j;
                rearms.push(Rearm {
                    req: DraftReq {
                        btrow: (sp.slot + 1) as u32,
                        first_token: chosen[col],
                        pos0: cols[col].pos as usize,
                        mpos0: cols[col].mpos,
                    },
                    seed: col,
                    slot: sp.slot,
                    remaining: seq.max_tokens.saturating_sub(seq.emitted.len()),
                });
            }
            // Same guard as the admit-side match: NEVER publish a multimodal prompt — image
            // placeholder ids make different images byte-identical, so cached vision blocks
            // both poison the index and (being duplicates) leak on release.
            if self.radix_enabled && seq.vision.is_none() {
                let full = seq.tokens.len() / PAGE_BLK;
                if full > seq.published {
                    self.radix
                        .insert(&seq.tokens[..full * PAGE_BLK], &seq.blocks[..full]);
                    seq.published = full;
                }
            }
        }
        let k_next = self.spec_depth_for(self.live_occupancy());
        let _t_arm = std::time::Instant::now();
        if k_next > 0 && !rearms.is_empty() {
            let m = self.mtp.as_ref().expect("spec round ran");
            let plan = if use_verify {
                self.verify_bands.get(&vband).expect("built above")
            } else if band >= self.plan.k {
                &self.plan
            } else {
                self.bands.get(&band).expect("band plan built by ensure_band")
            };
            Self::arm_drafts(
                m,
                self.pld,
                &mut self.slots,
                &mut self.stats,
                ctx,
                plan,
                rearms,
                k_next,
            )?;
        }
        if std::env::var("OSFKB_SPEC_TRACE").is_ok() {
            eprintln!(
                "[spec] round: verify {} us, arm {} us (cols={}, emitted={})",
                _t_verify_us,
                _t_arm.elapsed().as_micros(),
                cols.len(),
                out.len()
            );
        }
        self.reap_finished();
        Ok(Some(out))
    }

    /// Chained decode over the given slots (all in Decode state; see [`Self::step`]).
    fn step_chained(
        &mut self,
        gpu: &Lfm2Gpu,
        ctx: &GpuCtx,
        active: &[usize],
        band: usize,
    ) -> Result<Vec<Emission>> {
        // Chain length: bounded by the remaining budget of the nearest-to-finish sequence so we
        // rarely overshoot, and by the pool (each sequence needs blocks covering pos + steps).
        let mut steps = CHAIN_MAX;
        for &slot in active {
            let seq = self.slots[slot].as_ref().expect("active slot");
            steps = steps.min(seq.max_tokens.saturating_sub(seq.emitted.len()).max(1));
        }
        let mut cols = Vec::with_capacity(active.len());
        for &slot in active {
            let seq = self.slots[slot].as_mut().expect("active slot");
            let pos = seq.tokens.len() as u32 - 1;
            if !Self::ensure_blocks(
                &mut self.pool,
                &mut self.radix,
                seq,
                pos as usize + steps + 1,
            ) {
                // Pool pressure: fall back to a single unchained step next call.
                steps = 1;
            }
            if seq.blocks.len() > seq.table_synced {
                gpu.write_btab_row(ctx, (slot + 1) as u32, &seq.blocks);
                seq.table_synced = seq.blocks.len();
            }
            cols.push(BatchCol {
                token: *seq.tokens.last().expect("decode has tokens"),
                pos,
                btrow: (slot + 1) as u32,
                need_logit: true,
                // A generated token is text, but on a multimodal sequence its position resumes at
                // `next_pos`, NOT at the token index — the image consumed fewer positions than slots.
                mpos: Self::mpos_at(seq, pos as usize),
                embed_row: -1,
            });
        }
        let _t_chain = std::time::Instant::now();
        let chosen = gpu.batch_step_chained(ctx, self.band_plan(band), &cols, steps)?;
        self.stats.chain_us += _t_chain.elapsed().as_micros() as u64;
        self.stats.steps += steps as u64;
        self.stats.chained_steps += steps as u64;
        let mut out = Vec::new();
        for (ci, &slot) in active.iter().enumerate() {
            let seq = self.slots[slot].as_mut().expect("active slot");
            for srow in chosen.iter().take(steps) {
                if matches!(seq.state, SeqState::Finished) {
                    break; // overshoot truncated: tokens past EOS/max are discarded
                }
                let token = srow[ci];
                seq.tokens.push(token);
                seq.emitted.push(token);
                self.stats.tokens_out += 1;
                let finish = decide_finish(&self.eos, seq, token);
                out.push(Emission {
                    id: seq.id,
                    token,
                    finish,
                    // Chained decode is only taken by chain-safe sequences, which never request
                    // logprobs (that forces the per-column read-back path).
                    logprobs: None,
                });
                if finish.is_some() {
                    seq.state = SeqState::Finished;
                }
            }
            // Same guard as the admit-side match: NEVER publish a multimodal prompt — image
            // placeholder ids make different images byte-identical, so cached vision blocks
            // both poison the index and (being duplicates) leak on release.
            if self.radix_enabled && seq.vision.is_none() {
                let full = seq.tokens.len() / PAGE_BLK;
                if full > seq.published {
                    self.radix
                        .insert(&seq.tokens[..full * PAGE_BLK], &seq.blocks[..full]);
                    seq.published = full;
                }
            }
        }
        for slot in 0..self.slots.len() {
            let finished = matches!(
                self.slots[slot],
                Some(Sequence {
                    state: SeqState::Finished,
                    ..
                })
            );
            if finished {
                let mut seq = self.slots[slot].take().expect("finished seq");
                self.stats.finished += 1;
                self.release_sequence_blocks(&mut seq);
            }
        }
        Ok(out)
    }

    /// Drive until every submitted request finishes; returns each request's emitted tokens.
    pub fn run_to_completion(
        &mut self,
        gpu: &Lfm2Gpu,
        ctx: &GpuCtx,
    ) -> Result<HashMap<u64, Vec<u32>>> {
        let mut results: HashMap<u64, Vec<u32>> = HashMap::new();
        let mut idle = 0;
        while self.pending() > 0 {
            let steps_before = self.stats.steps;
            let ems = self.step(gpu, ctx)?;
            // Mid-prefill steps legitimately emit nothing — a stall is a step that ran NO columns
            // at all (pool starved with work still pending).
            if self.stats.steps == steps_before {
                idle += 1;
                if idle > 3 {
                    return Err(anyhow!(
                        "scheduler stalled with {} pending (pool starved?)",
                        self.pending()
                    ));
                }
            } else {
                idle = 0;
            }
            for e in ems {
                results.entry(e.id).or_default().push(e.token);
            }
        }
        Ok(results)
    }

    /// Disable radix prefix reuse (recurrent-state architectures).
    pub fn set_radix(&mut self, on: bool) {
        self.radix_enabled = on;
    }

    pub fn cache_entries(&self) -> usize {
        self.radix.len()
    }

    /// Free KV blocks right now — the leak canary (steady-state idle should equal capacity).
    pub fn kv_free_blocks(&self) -> usize {
        self.pool.available()
    }
}

/// The wide-prefill budget a one-shot generation should ask for: enough columns that the prompt
/// rides the batched plan, capped where `batch_step`'s `need_logit` mask (a u64) runs out.
/// Zero without subgroups — the wide plan needs them, and `new_with_wide` would ignore it anyway.
pub fn default_wide_budget(ctx: &GpuCtx) -> usize {
    if ctx.subgroups { MAX_BATCH_COLS } else { 0 }
}

/// One-shot generation **through the batched-prefill path** — the correct entry point for a CLI
/// or a single-request service.
///
/// [`Lfm2Gpu::decode`] is the M=1 solo loop: it ingests the prompt ONE TOKEN AT A TIME through the
/// decode GEMV kernels, so a long prompt costs one full forward pass per token. That is why a
/// 2,883-token prompt on the 4B claim-extractor took **147 s** (~20 tok/s) — the model never
/// touched the prefill path at all. The scheduler packs prompt tokens as COLUMNS per dispatch
/// (and, with a wide budget and subgroups, through the KC=16 wide plan + the tiled Q4 GEMM), which
/// is the same trap and the same fix the scoreboard's `prefill` row hit: 16.3 → 120.1 tok/s.
///
/// `decode` remains the byte-exact oracle the tests compare against; this is what callers want.
pub fn generate_once(
    gpu: &Lfm2Gpu,
    ctx: &GpuCtx,
    prompt: Vec<u32>,
    max_tokens: usize,
    params: RequestParams,
    eos: Vec<u32>,
) -> Result<(Vec<u32>, GenStats)> {
    let mut sched = Scheduler::new_with_wide(gpu, ctx, 1, default_wide_budget(ctx), eos);
    let prompt_len = prompt.len();
    let id = sched.submit_with_params(prompt, max_tokens, params)?;

    let t0 = std::time::Instant::now();
    let mut ttft = None;
    let mut out = Vec::with_capacity(max_tokens);
    loop {
        for em in sched.step(gpu, ctx)? {
            if em.id != id {
                continue;
            }
            if ttft.is_none() {
                // The first emission lands only after the whole prompt is ingested, so this wall
                // time IS the prefill — the number the 147 s regression was hiding in.
                ttft = Some(t0.elapsed().as_secs_f64());
            }
            out.push(em.token);
            if em.done() {
                let total = t0.elapsed().as_secs_f64();
                let prefill_s = ttft.unwrap_or(total);
                // The first token comes out of prefill; only the rest are decode.
                let decoded = out.len().saturating_sub(1) as f64;
                let stats = GenStats {
                    prefill_s,
                    prefill_tok_s: prompt_len as f64 / prefill_s.max(1e-9),
                    decode_tok_s: decoded / (total - prefill_s).max(1e-9),
                };
                return Ok((out, stats));
            }
        }
        // The request is neither queued nor in a slot, yet never emitted a terminal token — that
        // can only be a scheduler bug, and looping forever would hide it.
        anyhow::ensure!(
            sched.pending() > 0,
            "scheduler went idle without completing request {id}"
        );
    }
}

/// What a one-shot generation cost, split at the prefill/decode boundary — because those two are
/// bound by completely different things and a single tok/s number hides which one regressed.
#[derive(Debug, Clone, Copy)]
pub struct GenStats {
    pub prefill_s: f64,
    pub prefill_tok_s: f64,
    pub decode_tok_s: f64,
}

fn matched_plus_one(seq: &Sequence) -> usize {
    match seq.state {
        SeqState::Prefill { done } => done + 1,
        _ => seq.tokens.len(),
    }
}

/// Decide why (if at all) `seq` terminates after emitting `token` (already pushed to `emitted`).
/// EOS wins over a stop-token id, which wins over the length budget, so a natural stop that also
/// lands exactly on `max_tokens` reports `"stop"` rather than `"length"`.
fn decide_finish(eos: &[u32], seq: &Sequence, token: u32) -> Option<FinishReason> {
    if eos.contains(&token) {
        Some(FinishReason::Eos)
    } else if seq.params.stop_token_ids.contains(&token) {
        Some(FinishReason::StopTokenId(token))
    } else if seq.emitted.len() >= seq.max_tokens {
        Some(FinishReason::Length)
    } else {
        None
    }
}

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

    #[test]
    fn pool_alloc_release_roundtrip_and_trash_reserved() {
        let mut p = KvPool::new(8); // 7 usable + trash (block 7)
        let mut got = Vec::new();
        while let Some(b) = p.alloc() {
            got.push(b);
        }
        assert_eq!(got.len(), 7, "trash block must never be handed out");
        assert!(!got.contains(&7));
        for b in got {
            assert_eq!(p.unref(b), 0);
            p.release(b);
        }
        assert_eq!(p.available(), 7);
    }

    #[test]
    fn radix_matches_longest_full_block_prefix_exactly() {
        let mut pool = KvPool::new(64);
        let mut r = RadixIndex::new();
        let toks: Vec<u32> = (0..48).collect(); // 3 full blocks
        let blocks: Vec<u32> = (0..3).map(|_| pool.alloc().unwrap()).collect();
        for &b in &blocks {
            pool.unref(b); // cached, unreferenced
        }
        assert_eq!(r.insert(&toks, &blocks), 3);
        // Full match on the same tokens.
        let (m, n) = r.match_prefix(&mut pool, &toks);
        assert_eq!((m.clone(), n), (blocks.clone(), 48));
        for &b in &m {
            pool.unref(b);
        }
        // Diverging in block 2 → only 1 block matches.
        let mut t2 = toks.clone();
        t2[20] = 999;
        let (m2, n2) = r.match_prefix(&mut pool, &t2);
        assert_eq!((m2.len(), n2), (1, 16));
        pool.unref(m2[0]);
        // Sub-block prefix (10 tokens) → no full block → no match.
        let (m3, n3) = r.match_prefix(&mut pool, &toks[..10]);
        assert_eq!((m3.len(), n3), (0, 0));
    }

    /// Two identical prompts prefilled CONCURRENTLY each publish their own blocks for the same
    /// chunk hashes; the index keeps the FIRST as canonical. Release must consult
    /// `owned_blocks()` — `published` alone marked the duplicate as radix-owned and leaked it
    /// with refcount 0 and no owner (the 16-way identical-page livelock: the pool drained until
    /// admission spun forever with nothing runnable).
    #[test]
    fn duplicate_publish_is_not_owned_and_must_be_freed() {
        let mut pool = KvPool::new(64);
        let mut r = RadixIndex::new();
        let toks: Vec<u32> = (0..32).collect(); // 2 full blocks
        let first: Vec<u32> = (0..2).map(|_| pool.alloc().unwrap()).collect();
        let dup: Vec<u32> = (0..2).map(|_| pool.alloc().unwrap()).collect();
        assert_eq!(r.insert(&toks, &first), 2);
        assert_eq!(r.insert(&toks, &dup), 0, "duplicate chunks must not displace the canonical blocks");
        let owned = r.owned_blocks();
        assert!(first.iter().all(|b| owned.contains(b)));
        assert!(
            dup.iter().all(|b| !owned.contains(b)),
            "duplicate blocks are NOT index-owned — release must return them to the free list"
        );
        // The fixed release rule: unref to 0, then free anything the index does not own.
        let before = pool.available();
        for &b in &dup {
            assert_eq!(pool.unref(b), 0);
            if !owned.contains(&b) {
                pool.release(b);
            }
        }
        assert_eq!(pool.available(), before + dup.len(), "duplicates must actually come back");
    }

    #[test]
    fn radix_evicts_lru_leaves_but_never_referenced_blocks() {
        let mut pool = KvPool::new(4); // 3 usable
        let mut r = RadixIndex::new();
        let a: Vec<u32> = (0..16).collect();
        let b: Vec<u32> = (100..116).collect();
        let ba = pool.alloc().unwrap();
        let bb = pool.alloc().unwrap();
        r.insert(&a, &[ba]);
        r.insert(&b, &[bb]);
        pool.unref(ba);
        // `a` is older; both unreferenced → evicting frees `a` first.
        pool.unref(bb);
        let _ = r.match_prefix(&mut pool, &b); // touch b (LRU bump + reference)
        r.evict(&mut pool, 2);
        assert_eq!(pool.available(), 2, "one evictable (a) + one pre-free");
        assert_eq!(r.len(), 1, "b survives (referenced + recent)");
        // Referenced entries are never evicted even under pressure.
        r.evict(&mut pool, 10);
        assert_eq!(r.len(), 1);
    }

    #[test]
    fn radix_chain_verifies_chunks_against_hash_collisions() {
        let mut pool = KvPool::new(8);
        let mut r = RadixIndex::new();
        let a: Vec<u32> = (0..16).collect();
        let ba = pool.alloc().unwrap();
        r.insert(&a, &[ba]);
        pool.unref(ba);
        // Forge an entry under a's hash with different tokens: match must reject on chunk compare.
        let h = chain_hash(0, &a);
        r.map.get_mut(&h).unwrap().chunk[0] = 42;
        let (m, n) = r.match_prefix(&mut pool, &a);
        assert_eq!((m.len(), n), (0, 0));
    }
}