inferencelayer 0.2.14

//! OpenAI-compatible server for the CUDA engine, with continuous batching.
//!
//! Why a separate binary from `lfm2-serve`: that server's `GenState` carries a wgpu `GpuCtx` plus
//! the vision/embedding/rerank surface, and on a CUDA-only host `GpuCtx::new()` fails outright (the
//! Scaleway GPU image ships no Vulkan userspace). Rather than destabilise a working server, this
//! exposes exactly the endpoints the official benchmark harness drives — `/v1/completions`,
//! `/v1/chat/completions`, `/v1/models`, `/health` — so vLLM, SGLang and inferencelayer can all be
//! measured by the SAME tool. Measuring engines with different tools is what makes a benchmark
//! unpublishable.
//!
//! Continuous batching: one thread owns the GPU. Arrivals are prefilled into a free KV slot and
//! marked active; every step advances all active slots (each at its OWN position); finished
//! sequences release their slot mid-flight so a waiting request can take it.
use anyhow::{Context as _, Result};
use axum::Json;
use axum::extract::State;
use axum::response::IntoResponse;
use axum::response::sse::{Event, Sse};
use axum::routing::{get, post};
use inferencelayer::cuda_llm::CudaQwen3;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::mpsc;
use tokio::sync::mpsc as amp;

#[derive(Deserialize)]
struct CompletionReq {
    #[serde(default)]
    model: Option<String>,
    prompt: PromptField,
    #[serde(default = "default_max")]
    max_tokens: usize,
    #[serde(default)]
    stream: bool,
    #[serde(default)]
    ignore_eos: bool,
    #[serde(default)]
    temperature: f32,
}
fn default_max() -> usize {
    16
}

/// The harness sends either text or raw token ids (a list of ints), and the strict protocol needs
/// the id form so both engines get byte-identical prompts.
///
/// Hand-deserialized: `#[serde(untagged)]` buffers the value into a probe tree and tries variants
/// against it, which for a 2000-int prompt array is real time on the TTFT path. A visitor peeks
/// the JSON type once and parses in place.
enum PromptField {
    Text(String),
    Ids(Vec<u32>),
    Texts(Vec<String>),
}

impl<'de> Deserialize<'de> for PromptField {
    fn deserialize<D: serde::Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
        struct V;
        impl<'de> serde::de::Visitor<'de> for V {
            type Value = PromptField;
            fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                f.write_str("a string, an array of ints, or an array of strings")
            }
            fn visit_str<E: serde::de::Error>(self, v: &str) -> Result<PromptField, E> {
                Ok(PromptField::Text(v.to_string()))
            }
            fn visit_string<E: serde::de::Error>(self, v: String) -> Result<PromptField, E> {
                Ok(PromptField::Text(v))
            }
            fn visit_seq<A: serde::de::SeqAccess<'de>>(
                self,
                mut a: A,
            ) -> Result<PromptField, A::Error> {
                #[derive(Deserialize)]
                #[serde(untagged)]
                enum Elem {
                    I(u32),
                    S(String),
                }
                let mut ids = Vec::with_capacity(a.size_hint().unwrap_or(0));
                let mut texts: Option<Vec<String>> = None;
                while let Some(e) = a.next_element::<Elem>()? {
                    match (e, texts.as_mut()) {
                        (Elem::I(i), None) => ids.push(i),
                        (Elem::S(s), None) if ids.is_empty() => texts = Some(vec![s]),
                        (Elem::S(s), Some(t)) => t.push(s),
                        _ => return Err(serde::de::Error::custom("mixed prompt array")),
                    }
                }
                Ok(match texts {
                    Some(t) => PromptField::Texts(t),
                    None => PromptField::Ids(ids),
                })
            }
        }
        d.deserialize_any(V)
    }
}

#[derive(Deserialize)]
struct ChatReq {
    #[serde(default)]
    model: Option<String>,
    messages: Vec<ChatMsg>,
    #[serde(default = "default_max")]
    max_tokens: usize,
    #[serde(default)]
    stream: bool,
    #[serde(default)]
    ignore_eos: bool,
}
#[derive(Deserialize)]
struct ChatMsg {
    role: String,
    content: String,
}

#[derive(Serialize)]
struct Choice {
    index: usize,
    text: String,
    finish_reason: Option<String>,
    /// Non-standard: the raw generated token ids. The benchmark/validation protocol is ids-in /
    /// ids-out; text alone cannot be compared against a reference tokenizer-losslessly.
    ids: Vec<u32>,
}
#[derive(Serialize)]
struct Usage {
    prompt_tokens: usize,
    completion_tokens: usize,
    total_tokens: usize,
}
#[derive(Serialize)]
struct CompletionResp {
    id: String,
    object: &'static str,
    model: String,
    choices: Vec<Choice>,
    usage: Usage,
}

/// A unit of work handed to the engine thread.
struct Job {
    ids: Vec<u32>,
    max_tokens: usize,
    ignore_eos: bool,
    reply: amp::UnboundedSender<Ev>,
}

/// Engine -> handler. One `Delta` per generated token so the harness can measure ITL honestly.
enum Ev {
    /// The raw sampled id. Detokenisation happens in the per-request task, NOT in the engine loop:
    /// the loop is the only thread that can advance the GPU, so any work it does there is time no
    /// sequence is decoding. With 64 slots the tokenizer calls alone were ~7 ms per step.
    Token(u32),
    Done {
        completion_tokens: usize,
        prompt_tokens: usize,
    },
    Error(String),
}

/// Incremental detokenisation state, one per request.
///
/// Decoding the whole generation on every token is O(n^2); this keeps prefix/read offsets and
/// decodes a short window twice, taking the difference — cost independent of generation length.
/// It also only ever emits whole characters: a token can complete a multi-byte character that the
/// shorter decode rendered as U+FFFD, so the split point need not be a char boundary. Slicing on it
/// blind panicked mid-emoji.
struct Detok {
    /// None when the checkpoint ships no HF tokenizer.json (Mistral's tekken.json) — the id
    /// stream still flows, events just carry empty text. The benchmark protocol sends and counts
    /// ids, so nothing measured depends on detokenised text.
    tok: Option<Arc<tokenizers::Tokenizer>>,
    ids: Vec<u32>,
    prefix_off: usize,
    read_off: usize,
}

impl Detok {
    fn new(tok: Option<Arc<tokenizers::Tokenizer>>) -> Self {
        Self {
            tok,
            ids: Vec::new(),
            prefix_off: 0,
            read_off: 0,
        }
    }

    fn push(&mut self, id: u32) -> String {
        self.ids.push(id);
        let Some(tok) = self.tok.as_ref() else {
            return String::new();
        };
        let prev = tok
            .decode(&self.ids[self.prefix_off..self.read_off], true)
            .unwrap_or_default();
        let now = tok
            .decode(&self.ids[self.prefix_off..], true)
            .unwrap_or_default();
        if now.len() > prev.len() && now.starts_with(&prev) {
            let mut end = now.len();
            while end > prev.len() && !now.is_char_boundary(end) {
                end -= 1;
            }
            if end > prev.len() {
                let out = now[prev.len()..end].to_string();
                self.prefix_off = self.read_off;
                self.read_off = self.ids.len();
                return out;
            }
        }
        // partial multi-byte character: hold the window open and try again next token
        self.read_off = self.ids.len();
        String::new()
    }
}

#[derive(Clone)]
struct App {
    submit: mpsc::Sender<Job>,
    tok: Option<Arc<tokenizers::Tokenizer>>,
    model_name: String,
}

/// Per-slot state on the engine thread.
struct Slot {
    reply: amp::UnboundedSender<Ev>,
    ids: Vec<u32>,
    prompt_tokens: usize,
    budget: usize,
    ignore_eos: bool,
}

/// A long prompt mid-prefill under mixed scheduling: its chunks run INSIDE mixed steps (all live
/// decode rows + one chunk per step), so decoders never stall behind it and its first token lands
/// as soon as its own prefill completes — not after every prefill queued behind it.
struct PfCursor {
    slot: usize,
    ids: Vec<u32>,
    /// Tokens already prefilled; the next chunk starts here.
    pos: usize,
    max_tokens: usize,
    ignore_eos: bool,
    reply: amp::UnboundedSender<Ev>,
}

fn main() -> Result<()> {
    let a: Vec<String> = std::env::args().collect();
    if a.len() < 2 {
        eprintln!("usage: cuda-serve <model-dir> [--port N] [--slots N] [--max-seq N]");
        std::process::exit(2);
    }
    let dir = std::path::PathBuf::from(&a[1]);
    let flag = |name: &str, d: usize| -> usize {
        a.iter()
            .position(|x| x == name)
            .and_then(|i| a.get(i + 1))
            .and_then(|v| v.parse().ok())
            .unwrap_or(d)
    };
    let (port, slots, max_seq) = (
        flag("--port", 8300),
        flag("--slots", 32),
        flag("--max-seq", 4096),
    );

    // Optional: Mistral checkpoints carry tekken.json, which the tokenizers crate cannot read.
    // The benchmark protocol is ids-in/ids-out, so the server runs fine without detokenisation.
    let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
        .ok()
        .map(Arc::new);
    if tok.is_none() {
        eprintln!("no tokenizer.json — serving raw ids (text fields will be empty)");
    }
    let eos = read_eos(&dir);
    let model_name = dir.file_name().unwrap().to_string_lossy().to_string();
    let (tx, rx) = mpsc::channel::<Job>();

    // engine thread: owns the GPU
    {
        let dir = dir.clone();
        let tok = tok.clone();
        std::thread::spawn(move || {
            if let Err(e) = engine_loop(&dir, slots, max_seq, rx, tok, eos) {
                eprintln!("engine thread died: {e:?}");
                std::process::exit(1);
            }
        });
    }

    let app = axum::Router::new()
        .route("/v1/completions", post(completions))
        .route("/v1/chat/completions", post(chat))
        .route("/v1/models", get(models))
        .route("/health", get(|| async { "ok" }))
        .with_state(App {
            submit: tx,
            tok,
            model_name,
        });

    // A small worker pool on purpose. The default (one worker per core) plus 64 SSE tasks and
    // their detokenisers made ~70 runnable threads on 24 cores every step — and the ENGINE thread
    // sat in the runqueue behind them, adding whole scheduler quanta to every decode step. Six
    // workers drain 64 light streams comfortably and leave the cores to the engine.
    let rt = tokio::runtime::Builder::new_multi_thread()
        .worker_threads(6)
        .enable_all()
        .build()?;
    rt.block_on(async move {
        let l = tokio::net::TcpListener::bind(("0.0.0.0", port as u16)).await?;
        eprintln!("cuda-serve on :{port} ({slots} slots, max_seq {max_seq})");
        // Every response here is a small latency-sensitive write (SSE events, one JSON body);
        // without NODELAY they can sit in the kernel behind Nagle waiting on the client's ACK.
        // axum 0.8 has no Serve::tcp_nodelay; the Listener impl below is its supported hook.
        axum::serve(NodelayListener(l), app).await?;
        Ok::<_, anyhow::Error>(())
    })
}

/// `TcpListener` that sets TCP_NODELAY on every accepted connection.
struct NodelayListener(tokio::net::TcpListener);

impl axum::serve::Listener for NodelayListener {
    type Io = tokio::net::TcpStream;
    type Addr = std::net::SocketAddr;
    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
        loop {
            match self.0.accept().await {
                Ok((io, addr)) => {
                    let _ = io.set_nodelay(true);
                    return (io, addr);
                }
                Err(_) => continue,
            }
        }
    }
    fn local_addr(&self) -> std::io::Result<Self::Addr> {
        self.0.local_addr()
    }
}

fn read_eos(dir: &std::path::Path) -> Vec<u32> {
    let mut v = Vec::new();
    if let Ok(s) = std::fs::read_to_string(dir.join("config.json"))
        && let Ok(j) = serde_json::from_str::<serde_json::Value>(&s)
    {
        match j.get("eos_token_id") {
            Some(serde_json::Value::Number(n)) => v.push(n.as_u64().unwrap_or(0) as u32),
            Some(serde_json::Value::Array(a)) => {
                v.extend(a.iter().filter_map(|x| x.as_u64()).map(|x| x as u32))
            }
            _ => {}
        }
    }
    v
}

/// The scheduler. Admit → step all active slots → emit → release finished.
fn engine_loop(
    dir: &std::path::Path,
    slots: usize,
    max_seq: usize,
    rx: mpsc::Receiver<Job>,
    tok: Option<Arc<tokenizers::Tokenizer>>,
    eos: Vec<u32>,
) -> Result<()> {
    let mut eng = CudaQwen3::load(dir, max_seq)?;
    eng.init_batch(slots, max_seq)?;
    // Capture the sequential decode graph NOW, on an idle stream. Capturing lazily on the first
    // M=1 request raced the admission burst's async tail (ILLEGAL_ADDRESS unless
    // CUDA_LAUNCH_BLOCKING=1).
    eng.sync()?;
    eprintln!("model loaded, {slots} slots ready");

    let mut live: Vec<Option<Slot>> = (0..slots).map(|_| None).collect();
    let mut queue: VecDeque<Job> = VecDeque::new();
    // Tokens sampled by the step in flight, emitted while the NEXT step runs; the mask records
    // which slots were live when that step ran (a slot admitted afterwards has no token in it).
    let mut pending: Option<Vec<u32>> = None;
    let mut prev_mask: Vec<bool> = vec![false; slots];
    // Exactly one live request, sitting in slot 0, decoding on the engine's SEQUENTIAL graph —
    // 2.2 ms/step against the width-1 batch path's 3.1. Slot 0's KV region is offset 0, which is
    // the sequential path's, so the state is shared: entering needs nothing (the prefill already
    // set d_pos and d_token), and leaving is one occupy at the current position. SEQ1=0 disables.
    let mut seq_mode = false;
    // DEFAULT ON since the bisect terminal: the seq and batch arms agree to 0.05% and differ only
    // in tie-breaks at q4-flattened razor sites — the same property vLLM/SGLang exhibit across
    // batch sizes. SEQ1=0 opts out (restores cross-width tie-break identity). Forced off under
    // Q4FREE=1: the seq entry path runs one EAGER step, which multiplies the ORIGINAL q4 quants
    // that Q4FREE just freed.
    let seq_enabled = std::env::var("SEQ1").as_deref() != Ok("0") && eng.seq_capable();
    // Admission accumulation window (see the admit block). ADMIT_MIN/ADMIT_HOLD env-tunable.
    let admit_min: usize = std::env::var("ADMIT_MIN")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(8);
    // hold=1: with wave prefill, waiting extra steps to fatten the wave no longer pays — a
    // second wave is one ~2.4 ms step away, while each held step is straight TTFT for the
    // stragglers (M=8 burst p50 19-25 ms at hold=2, 14.6-16 ms at hold=1; M=64 unchanged).
    let admit_hold: u32 = std::env::var("ADMIT_HOLD")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(1);
    // Wave cap. The loadgen's TTFT percentiles exposed what conc.py never measured: a 64-burst
    // was admitted as ONE wave — 63 prefill graphs replaying serially on the GPU (~3.5 ms each
    // at T=23) behind a single sync, so every first token released together at ~235 ms (min 4.0,
    // p50/p90/max all 235.x). Capping the wave interleaves sub-waves with decode steps: first
    // tokens stream as a staircase instead of one cliff. The real fix (wave-batched prefill,
    // one [sum-T] pass) comes next; this bounds the damage structurally.
    // With wave prefill a big wave is ONE cheap pass, so cap only where prefills serialise.
    // Full-width admission: the TTFT staircase comes from WAVE_SUB (sub-waves that emit as
    // they land), NOT from capping admission. Capping cost a decode step between sub-waves
    // (~1% aggregate: Llama 5491 -> 5447, Mistral 2117 -> 2095); sub-waves cost nothing.
    let admit_cap: usize = std::env::var("ADMIT_CAP")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(
            if std::env::var("WAVE").as_deref() != Ok("0") && eng.wave_capable() {
                64
            } else {
                8
            },
        );
    let mut held_steps: u32 = 0;
    // MIXED PREFILL/DECODE (MIX=0 disables; dense bf16-mirror models only). Prompts longer than
    // the wave cap become per-request CURSORS instead of prefilling to completion inside
    // admission; while a cursor is outstanding, each scheduler iteration runs ONE mixed step —
    // every live decode row AND the cursor's next chunk in the same forward pass. The decode
    // rows ride the chunk's weight streaming for free, and first tokens staircase per request
    // instead of releasing after the whole prefill backlog (vLLM's iteration shape). Cursors
    // park on the HIGHEST free slot so they sit above the decode width (see free_slot_high).
    let mix_on = std::env::var("MIX").as_deref() != Ok("0") && eng.mix_capable();
    // OVERLAP mode (MIX=2, experimental — MEASURED WORSE than row-merge): a long prefill
    // replays its captured one-shot graph on the engine's second stream while the decode graph
    // free-runs on the main stream. The premise (decode rides free) fails because BOTH sides
    // are ~15 GB/step HBM workloads — run concurrently they stretch each other (Llama 8x2k
    // 441.9 vs row-merge 454.5, prefill passes 93 ms vs 73), where row-merge's ONE weight read
    // serves both. Kept for A/B and for workloads where the prefill is compute-bound enough to
    // leave bandwidth headroom. Outputs are byte-identical to solo (same one-shot graph).
    let mix_overlap = mix_on && std::env::var("MIX").as_deref() == Ok("2") && eng.overlap_capable();
    let mut flight: Option<(PfCursor, cudarc::driver::CudaEvent)> = None;
    // Prefill-token budget per mixed step. At 2048 a 2k prompt is ONE chunk at pos0=0 — no
    // prefix gather, no extra attention triangle, i.e. one-shot cost. Only prompts above the
    // budget pay the chunk arm's gather on their later chunks.
    let mix_budget: usize = std::env::var("MIX_BUDGET")
        .ok()
        .and_then(|v| v.parse().ok())
        .unwrap_or(2048)
        .max(1);
    // Prompts at or under this length keep the wave path (batched same-length prefill beats
    // everything for shorts); above it they cursor. Mirrors WAVE_MAXT unless overridden.
    let mix_minlen: usize = std::env::var("MIX_MINLEN")
        .ok()
        .and_then(|v| v.parse().ok())
        .or_else(|| std::env::var("WAVE_MAXT").ok().and_then(|v| v.parse().ok()))
        .unwrap_or(64);
    let mut prefilling: VecDeque<PfCursor> = VecDeque::new();
    // Dedicated emitter: the engine thread hands each step's (sender, events) batch over and
    // never runs a tokio waker itself. Send order per receiver is preserved — the admission
    // first-token is sent inline BEFORE any batched token for that request exists.
    let (emit_tx, emit_rx) = mpsc::channel::<Vec<(amp::UnboundedSender<Ev>, Ev, Option<Ev>)>>();
    std::thread::spawn(move || {
        while let Ok(batch) = emit_rx.recv() {
            for (tx, ev, fin) in batch {
                let _ = tx.send(ev);
                if let Some(f) = fin {
                    let _ = tx.send(f);
                }
            }
        }
    });
    // STEP_PROF=1: aggregate per-phase wall time every 64 steps, to stderr.
    let prof = std::env::var("STEP_PROF").is_ok();
    let (mut t_admit, mut t_launch, mut t_emit, mut t_sync, mut t_read) =
        (0f64, 0f64, 0f64, 0f64, 0f64);
    let mut nstep = 0u64;
    macro_rules! ph {
        ($acc:ident, $e:expr) => {{
            if prof {
                let t0 = std::time::Instant::now();
                let r = $e;
                $acc += t0.elapsed().as_secs_f64();
                r
            } else {
                $e
            }
        }};
    }

    loop {
        // 0. a step's tokens may still be pending after the last live slot retired mid-emit;
        // flush before blocking so no request's final token waits on the NEXT arrival.
        if live.iter().all(|s| s.is_none()) {
            pending = None;
        }
        // 1. collect arrivals — block only when there is nothing to do
        let mut cold = false;
        if live.iter().all(|s| s.is_none())
            && queue.is_empty()
            && prefilling.is_empty()
            && flight.is_none()
        {
            match rx.recv() {
                Ok(j) => {
                    queue.push_back(j);
                    cold = true;
                }
                Err(_) => return Ok(()), // all senders gone
            }
        }
        while let Ok(j) = rx.try_recv() {
            queue.push_back(j);
        }
        // 1a. COLD-BURST COALESCE: the GPU is idle and one request just landed. A concurrent
        // burst's siblings are typically <1 ms behind it; admitting the head solo splits the
        // burst into TWO full weight-streaming prefills (solo staged + the rest as a wave) and
        // every sibling's TTFT eats both. Waiting a moment costs only the head's latency —
        // the GPU has nothing else to do — and turns 1+(N-1) into one N-wave.
        // Silence-triggered: stop 300 us after the LAST arrival (a lone request pays ~0.3 ms,
        // not the full window), hard-capped at ADMIT_COLD_US so a drip can't stall admission.
        if cold {
            let cold_us: u64 = std::env::var("ADMIT_COLD_US")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(4000);
            let sil_us: u64 = std::env::var("COLD_SILENCE_US")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(800);
            let t0 = std::time::Instant::now();
            let mut last = t0;
            while t0.elapsed().as_micros() < cold_us as u128
                && last.elapsed().as_micros() <= sil_us as u128
            {
                match rx.recv_timeout(std::time::Duration::from_micros(100)) {
                    Ok(j) => {
                        queue.push_back(j);
                        last = std::time::Instant::now();
                    }
                    Err(mpsc::RecvTimeoutError::Timeout) => {}
                    Err(mpsc::RecvTimeoutError::Disconnected) => break,
                }
            }
        }

        // 1b. COMPACT the live set. Under equal budgets retirement is arrival-ordered, so
        // without this the highest-ever slot pins the batch width for the entire drain — every
        // step then pays for sequences that already left (the M>=32 rows lost ~20-30% to it).
        // A retired slot is refilled by MOVING the highest live sequence down (~0.5 ms of KV
        // dtod against a permanently narrower step). Safe here: the stream is quiescent between
        // iterations, and the pending-emit bookkeeping moves with the slot.
        loop {
            let Some(hi) = live.iter().rposition(|s| s.is_some()) else {
                break;
            };
            // A destination must be genuinely free: a slot can be live-None yet RESERVED by a
            // mid-prefill cursor, and moving a sequence onto it would clobber the prompt KV the
            // cursor's chunks are laying down.
            let Some(lo) = live[..hi].iter().enumerate().position(|(i, s)| {
                s.is_none()
                    && !prefilling.iter().any(|c| c.slot == i)
                    && flight.as_ref().map(|(c, _)| c.slot) != Some(i)
            }) else {
                break;
            };
            let used = live[hi].as_ref().map(|s| s.ids.len() + 1).unwrap_or(0);
            eng.move_slot(hi, lo, used)?;
            live[lo] = live[hi].take();
            prev_mask[lo] = prev_mask[hi];
            prev_mask[hi] = false;
            if let Some(p) = pending.as_mut() {
                p[lo] = p[hi];
            }
        }

        // 2. admit while slots are free — ALL prefill graphs replay back-to-back on the stream
        // and ONE sync drains the burst. The serial form (replay, sync, read, repeat) cost 231 ms
        // to admit a 64-request wave — ~3.6 ms each — all of it in front of the first decode step.
        //
        // ACCUMULATION WINDOW: when the decode set is busy and arrivals TRICKLE (conc.py's 64
        // GIL-bound client threads spread over ~100 ms), admitting each arrival alone stalls the
        // whole decode set once per request — ~40 singleton waves cost the M=64 row ~240 ms of
        // wall against vLLM, whose scheduler folds prefills into the running iteration. Hold
        // arrivals for at most ADMIT_HOLD decode steps (default 2, ~2.5-8 ms) or until ADMIT_MIN
        // (default 8) are waiting; an idle decode set always admits immediately, so single-request
        // TTFT is untouched. This is scheduling policy, not measurement policy: same harness,
        // same requests, tokens counted the same way.
        if prof {
            nstep = nstep.wrapping_add(1);
        }
        let t_a0 = std::time::Instant::now();
        let live_now = live.iter().filter(|s| s.is_some()).count();
        let admit_now = if queue.is_empty() || flight.is_some() {
            // An in-flight overlap prefill owns the Prefill buffers; every admission path
            // (wave, staged, chunk) writes them, so arrivals wait out the flight (<= one
            // prefill's wall). Cursor creation waits with them.
            held_steps = 0;
            false
        } else if live_now == 0 || queue.len() >= admit_min || held_steps >= admit_hold {
            held_steps = 0;
            true
        } else {
            held_steps += 1;
            false
        };
        // Pull the whole waiting burst, upload every prompt in ONE host->device copy, then let
        // each prefill replay + bookkeeping queue on the stream with NO synchronizing call —
        // cudarc's pageable memcpy_htod waits for prior stream work, which is what serialized
        // the old loop at ~3.6 ms per admission.
        let mut burst: Vec<(usize, Job)> = Vec::new();
        // Long prompts can't wave-batch: they prefill serially at T-proportional cost, so a
        // 64-burst of 2k prompts would rebuild the TTFT staircase (64 x ~36 ms) the wave path
        // exists to kill. Cap them per admission wave; the next wave is only one step away.
        // Long prompts prefill serially (see WAVE_MAXT), so cap how many enter one admission
        // wave or a 64-burst of 2k prompts rebuilds the TTFT staircase.
        let mut big = 0usize;
        let mut longs: Vec<(usize, Job)> = Vec::new();
        while admit_now && !queue.is_empty() && burst.len() < admit_cap {
            // Mixed scheduling: long prompts don't prefill inside admission at all — they take
            // the HIGHEST free slot (above the decode width, where no decode-row kernel can
            // touch their slab) and become cursors, chunked through the mixed steps below.
            if mix_on && queue.front().unwrap().ids.len() > mix_minlen {
                let Some(slot) = eng.free_slot_high() else {
                    break;
                };
                let j = queue.pop_front().unwrap();
                eng.reserve_slot(slot);
                longs.push((slot, j));
                continue;
            }
            if queue.front().unwrap().ids.len() > 64 {
                if big >= 8 {
                    break;
                }
                big += 1;
            }
            let Some(slot) = eng.free_slot() else { break };
            let j = queue.pop_front().unwrap();
            if j.ids.is_empty() {
                let _ = j.reply.send(Ev::Error("empty prompt".into()));
                continue;
            }
            // reserve host-side so free_slot cannot hand this slot out twice within the burst
            eng.reserve_slot(slot);
            burst.push((slot, j));
        }
        // Stage cursor prompts on device NOW (admission is where pageable uploads live) and
        // queue them. Overlap mode needs no per-slot staging (the flight htod's straight into
        // p.tok on stream2), only the quiesce; row-merge mode stages chunks per slot.
        for (slot, j) in longs {
            let staged = if mix_overlap {
                eng.quiesce_slot(slot)
            } else {
                eng.mix_stage(slot, &j.ids)
            };
            match staged {
                Ok(()) => prefilling.push_back(PfCursor {
                    slot,
                    pos: 0,
                    ids: j.ids,
                    max_tokens: j.max_tokens,
                    ignore_eos: j.ignore_eos,
                    reply: j.reply,
                }),
                Err(e) => {
                    let _ = j.reply.send(Ev::Error(format!("mix stage: {e}")));
                    eng.release_slot(slot)?;
                }
            }
        }
        if !burst.is_empty() {
            // Admitted requests are grouped into BATCHES, one per prefill sub-wave, so each
            // sub-wave's first tokens are emitted as soon as that sub-wave lands instead of
            // after the whole burst. Same total prefill work (no extra decode step between
            // sub-waves, which is what capping the wave cost), staircase TTFT for free.
            let sub_wave: usize = std::env::var("WAVE_SUB")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(32);
            let mut admitted: Vec<(usize, Job)> = Vec::new();
            // Emit ONE sub-wave's first tokens: sync, read d_firsts, stream each token, and
            // move the request into the live set. A macro, not a closure — it needs &mut eng,
            // &mut live and the emitter at once, and its `continue` must bind to its own loop.
            macro_rules! emit_batch { ($batch:expr) => {{
                let firsts = eng.first_tokens()?;   // one sync + readback per sub-wave
                for (slot, j) in $batch {
                    let plen = j.ids.len();
                    let first = firsts[slot];
                    if std::env::var("EMIT_TRACE").is_ok() {
                        let seen = eng.peek_prefill_tokens(4).unwrap_or_default();
                        eprintln!("[adm] slot={slot} plen={plen} first={first} p.tok[0..4]={seen:?} sent={:?}",
                                  &j.ids[..4.min(j.ids.len())]);
                    }
                    let _ = j.reply.send(Ev::Token(first));
                    if j.max_tokens <= 1 || (!j.ignore_eos && eos.contains(&first)) {
                        let _ = j.reply.send(Ev::Done { completion_tokens: 1, prompt_tokens: plen });
                        eng.release_slot(slot)?;
                        continue;
                    }
                    eng.set_slot_token_async(slot, first)?;
                    let mut ids = j.ids.clone();
                    ids.push(first);
                    live[slot] = Some(Slot {
                        reply: j.reply,
                        ids,
                        prompt_tokens: plen,
                        budget: j.max_tokens,
                        ignore_eos: j.ignore_eos,
                    });
                }
            }}; }
            // WAVE PREFILL: same-length groups of short prompts run as ONE [B*T]-row pass —
            // the serial per-request graph replays put every first token of a 64-burst at
            // ~235 ms (each 23-token prefill is ~3.5 ms of GPU, and they queued back to back).
            // Leftovers (singletons, long prompts, wave-incapable models) keep the staged path.
            let wave_on = std::env::var("WAVE").as_deref() != Ok("0") && eng.wave_capable();
            let mut rest: Vec<(usize, Job)> = Vec::new();
            if wave_on {
                let mut groups: std::collections::HashMap<usize, Vec<(usize, Job)>> =
                    std::collections::HashMap::new();
                for (slot, j) in burst {
                    groups.entry(j.ids.len()).or_default().push((slot, j));
                }
                // WAVE_ROWS caps a single wave's [B*T] pass so the prefill scratch stays
                // bounded; long prompts now wave too (per-segment flash attention inside),
                // which is what removes the serial-prefill staircase for 2k-token bursts.
                let wave_rows: usize = std::env::var("WAVE_ROWS")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(16384);
                // WAVE_MAXT: longest prompt that goes through a wave. MEASURED: waving 8x2000
                // was WORSE than 8 serial graph-replayed prefills (TTFT 178 vs 158 ms,
                // aggregate 1053 vs 1100) — at T=2000 a single request's GEMMs are already
                // efficient, so the wave adds prefill-scratch growth and gives up the per-length
                // captured graph while amortizing nothing. Short prompts are the opposite case.
                let wave_maxt: usize = std::env::var("WAVE_MAXT")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(64);
                for (tseg, grp) in groups {
                    if tseg == 0 || tseg > wave_maxt || grp.len() < 2 {
                        rest.extend(grp);
                        continue;
                    }
                    // split into successive waves: bounded by the scratch cap AND by the
                    // sub-wave size that gives the TTFT staircase
                    let per_wave = (wave_rows / tseg.max(1)).max(1).min(sub_wave);
                    let mut grp = grp;
                    while grp.len() > per_wave {
                        let tail = grp.split_off(per_wave);
                        let flat: Vec<u32> = grp
                            .iter()
                            .flat_map(|(_, j)| j.ids.iter().copied())
                            .collect();
                        let slots_v: Vec<usize> = grp.iter().map(|(s, _)| *s).collect();
                        match eng.prefill_wave(&flat, tseg, &slots_v) {
                            Ok(()) => {
                                let mut batch = Vec::with_capacity(grp.len());
                                for (slot, j) in grp {
                                    eng.occupy_slot_async(slot, tseg)?;
                                    batch.push((slot, j));
                                }
                                emit_batch!(batch);
                            }
                            Err(e) => {
                                eprintln!("wave prefill failed ({e}); falling back to staged");
                                rest.extend(grp);
                            }
                        }
                        grp = tail;
                    }
                    let flat: Vec<u32> = grp
                        .iter()
                        .flat_map(|(_, j)| j.ids.iter().copied())
                        .collect();
                    let slots_v: Vec<usize> = grp.iter().map(|(s, _)| *s).collect();
                    match eng.prefill_wave(&flat, tseg, &slots_v) {
                        Ok(()) => {
                            let mut batch = Vec::with_capacity(grp.len());
                            for (slot, j) in grp {
                                eng.occupy_slot_async(slot, tseg)?;
                                batch.push((slot, j));
                            }
                            emit_batch!(batch);
                        }
                        Err(e) => {
                            eprintln!("wave prefill failed ({e}); falling back to staged");
                            rest.extend(grp);
                        }
                    }
                }
            } else {
                rest = burst;
            }
            // CHUNKED PREFILL (CHUNK=<n>, opt-in): long prompts are prefilled n tokens at a
            // time via prefill_chunk (strip attention against the slab). Each chunk is a
            // separate launch, so the scheduler can interleave decode steps between them —
            // which is the point; this first cut runs them back-to-back and only proves the
            // path is numerically identical to a one-shot prefill.
            // CHUNK>0 is CORRECT (byte-identical to one-shot) but ~9x slower per prefill —
            // see attn_strip's header. Default 0 until the kernel is mma-based.
            let chunk: usize = std::env::var("CHUNK")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0);
            if chunk > 0 {
                let mut still: Vec<(usize, Job)> = Vec::new();
                let mut cbatch: Vec<(usize, Job)> = Vec::new();
                for (slot, j) in std::mem::take(&mut rest) {
                    // No length guard on purpose: CHUNK >= prompt length runs the whole prompt
                    // as ONE chunk, which is the discriminator that localised the defect —
                    // single-chunk output equals many-chunk output, so attn_strip is wrong and
                    // the chunk-to-chunk carry is fine. Keep this behaviour for the next debugger.
                    if j.ids.len() < 2 {
                        still.push((slot, j));
                        continue;
                    }
                    let n = j.ids.len();
                    let mut pos = 0usize;
                    let mut ok = true;
                    while pos < n {
                        let end = (pos + chunk).min(n);
                        let last = end == n;
                        if let Err(e) = eng.prefill_chunk(&j.ids[pos..end], slot, pos, last) {
                            eprintln!("chunked prefill failed ({e}); falling back to one-shot");
                            ok = false;
                            break;
                        }
                        pos = end;
                    }
                    if ok {
                        eng.occupy_slot_async(slot, n)?;
                        cbatch.push((slot, j));
                    } else {
                        still.push((slot, j));
                    }
                }
                if !cbatch.is_empty() {
                    emit_batch!(cbatch);
                }
                rest = still;
            }
            let concat: Vec<u32> = rest
                .iter()
                .flat_map(|(_, j)| j.ids.iter().copied())
                .collect();
            if !concat.is_empty() {
                eng.stage_burst(&concat)?;
            }
            let mut off = 0usize;
            for (slot, j) in rest {
                let n = j.ids.len();
                match eng.prefill_slot_staged(&j.ids, off, slot) {
                    Ok(()) => {}
                    Err(e) => {
                        let _ = j.reply.send(Ev::Error(format!("prefill: {e}")));
                        eng.release_slot(slot)?;
                        off += n;
                        continue;
                    }
                }
                eng.occupy_slot_async(slot, n)?;
                admitted.push((slot, j));
                off += n;
            }
            if !admitted.is_empty() {
                let batch = std::mem::take(&mut admitted);
                emit_batch!(batch);
            }
        }

        if prof {
            t_admit += t_a0.elapsed().as_secs_f64();
        }
        if live.iter().all(|s| s.is_none()) && prefilling.is_empty() && flight.is_none() {
            seq_mode = false;
            continue;
        }

        // ---- sequential fast path: one live request, slot 0, nothing queued ----
        let live_count = live.iter().filter(|s| s.is_some()).count();
        if seq_enabled
            && live_count == 1
            && live[0].is_some()
            && queue.is_empty()
            && prefilling.is_empty()
            && flight.is_none()
        {
            if !seq_mode {
                // Entering: slot 0's prefill (or the transition below) left d_pos/d_token at the
                // sequence head; the sequential graph self-feeds from there. Captured LAZILY here
                // — after a real prefill, matching the bench's (working) capture condition; the
                // boot-time capture produced degenerate steps, and the crash that originally
                // motivated boot capture predates the orphaned-graph keystone fix.
                if !eng.has_seq_graph() {
                    // ONE EAGER STEP FIRST, then capture. Every capture taken without a prior
                    // eager step (boot, cold entry) replayed as garbage; every capture taken
                    // after one worked — step()'s first call materializes lazy state, and an
                    // allocation made INSIDE capture becomes a graph-owned node that
                    // AUTO_FREE_ON_LAUNCH frees out from under later replays. The eager step's
                    // token is real — emit it through the normal bookkeeping below by falling
                    // through with the graph left uncaptured this iteration.
                    eng.step(0, 0)?;
                    eng.sync()?;
                    let t = eng.last_token()?;
                    let sl = live[0].as_mut().unwrap();
                    sl.ids.push(t);
                    let produced = sl.ids.len() - sl.prompt_tokens;
                    let fin = ((!sl.ignore_eos && eos.contains(&t)) || produced >= sl.budget).then(
                        || Ev::Done {
                            completion_tokens: produced,
                            prompt_tokens: sl.prompt_tokens,
                        },
                    );
                    let done = fin.is_some();
                    // via the emitter: an inline send could overtake this request's still-queued
                    // batch-path tokens
                    let _ = emit_tx.send(vec![(sl.reply.clone(), Ev::Token(t), fin)]);
                    if done {
                        eng.release_slot(0)?;
                        live[0] = None;
                        continue;
                    }
                    eng.capture()?;
                    seq_mode = true;
                    continue;
                }
                // pending tokens from the batch path belong to this same request — flush them
                // through the normal lagged emit before switching cadence.
                if let (Some(prev), Some(sl)) = (pending.take(), live[0].as_mut()) {
                    if prev_mask[0] {
                        let t = prev[0];
                        sl.ids.push(t);
                        let produced = sl.ids.len() - sl.prompt_tokens;
                        let fin = ((!sl.ignore_eos && eos.contains(&t)) || produced >= sl.budget)
                            .then(|| Ev::Done {
                                completion_tokens: produced,
                                prompt_tokens: sl.prompt_tokens,
                            });
                        let done = fin.is_some();
                        let _ = emit_tx.send(vec![(sl.reply.clone(), Ev::Token(t), fin)]);
                        if done {
                            eng.release_slot(0)?;
                            live[0] = None;
                            continue;
                        }
                    }
                    // the sequential graph feeds from d_token = last sampled; the batch path kept
                    // it in b.tok[0] — reseed the head explicitly.
                    eng.seed_seq_head(*sl.ids.last().unwrap(), sl.ids.len() - 1)?;
                }
                seq_mode = true;
            }
            if std::env::var("EMIT_TRACE").is_ok() {
                let (pp, tt) = (eng.peek_pos()?, eng.last_token()?);
                let n = live[0].as_ref().map(|s| s.ids.len()).unwrap_or(0);
                eprintln!("[seq pre] d_pos={pp} d_token={tt} ids.len={n}");
            }
            eng.run_graph(1)?;
            eng.sync()?;
            let t = eng.last_token()?;
            let sl = live[0].as_mut().unwrap();
            sl.ids.push(t);
            let produced = sl.ids.len() - sl.prompt_tokens;
            let fin =
                ((!sl.ignore_eos && eos.contains(&t)) || produced >= sl.budget).then(|| Ev::Done {
                    completion_tokens: produced,
                    prompt_tokens: sl.prompt_tokens,
                });
            let done = fin.is_some();
            let _ = emit_tx.send(vec![(sl.reply.clone(), Ev::Token(t), fin)]);
            if done {
                eng.release_slot(0)?;
                live[0] = None;
                seq_mode = false;
            }
            continue;
        }
        if seq_mode {
            // A second request arrived (or the single one moved off slot 0): hand the sequence
            // back to the batch machinery at its current position.
            if let Some(sl) = live[0].as_ref() {
                eng.occupy_slot(0, *sl.ids.last().unwrap(), sl.ids.len() - 1)?;
            }
            seq_mode = false;
        }

        // Live-context signal for the engine's context-keyed kernel policies (the tensor
        // attention arm and its graph key). Max over live sequences; decays as they retire.
        eng.set_ctx_live(
            live.iter()
                .flatten()
                .map(|s| s.ids.len())
                .max()
                .unwrap_or(0),
        );

        // Step k-1's tokens, emitted while step k runs. ONE definition shared by the pure-decode
        // and mixed-step branches — the bookkeeping must be identical or the two cadences drift.
        macro_rules! flush_pending {
            () => {{
                if let Some(prev) = pending.take() {
                    let mut out: Vec<(amp::UnboundedSender<Ev>, Ev, Option<Ev>)> =
                        Vec::with_capacity(live.len());
                    for (slot, cell) in live.iter_mut().enumerate() {
                        let Some(sl) = cell.as_mut() else { continue };
                        if !prev_mask[slot] {
                            continue;
                        } // admitted after step k-1: no token yet
                        let t = prev[slot];
                        sl.ids.push(t);
                        let produced = sl.ids.len() - sl.prompt_tokens;
                        if produced <= 4 && std::env::var("EMIT_TRACE").is_ok() {
                            eprintln!("[emit] slot={slot} produced={produced} tok={t}");
                        }
                        let hit_eos = !sl.ignore_eos && eos.contains(&t);
                        let done = hit_eos || produced >= sl.budget;
                        let fin = done.then(|| Ev::Done {
                            completion_tokens: produced,
                            prompt_tokens: sl.prompt_tokens,
                        });
                        out.push((sl.reply.clone(), Ev::Token(t), fin));
                        if done {
                            eng.release_slot(slot)?;
                            *cell = None;
                        }
                    }
                    if !out.is_empty() {
                        let _ = emit_tx.send(out);
                    }
                }
            }};
        }
        // A cursor's prefill just completed: emit its first token and move it into the live set
        // (the admission flow's shape: occupy, sync+read d_firsts, inline first-token send).
        // prev_mask stays false — its first DECODE token comes from the next step, not this one.
        macro_rules! finish_cursor {
            ($cur:expr) => {{
                let mut cur = $cur;
                let plen = cur.ids.len();
                eng.occupy_slot_async(cur.slot, plen)?;
                let firsts = eng.first_tokens()?;
                let first = firsts[cur.slot];
                let _ = cur.reply.send(Ev::Token(first));
                eng.mix_unstage(cur.slot);
                if cur.max_tokens <= 1 || (!cur.ignore_eos && eos.contains(&first)) {
                    let _ = cur.reply.send(Ev::Done {
                        completion_tokens: 1,
                        prompt_tokens: plen,
                    });
                    eng.release_slot(cur.slot)?;
                } else {
                    eng.set_slot_token_async(cur.slot, first)?;
                    let mut ids = std::mem::take(&mut cur.ids);
                    ids.push(first);
                    live[cur.slot] = Some(Slot {
                        reply: cur.reply,
                        ids,
                        prompt_tokens: plen,
                        budget: cur.max_tokens,
                        ignore_eos: cur.ignore_eos,
                    });
                    prev_mask[cur.slot] = false;
                }
            }};
        }

        // 2a. OVERLAPPED PREFILL: launch the head cursor's one-shot graph on stream2, then keep
        // taking ordinary decode steps on the main stream until its event completes. The decode
        // set never sees the prefill except as bandwidth contention; the prefill is the same
        // graph a solo request takes, so outputs are byte-identical to solo.
        if mix_overlap {
            if flight.is_none() && !prefilling.is_empty() {
                let cur = prefilling.pop_front().unwrap();
                match eng.prefill_overlap_start(&cur.ids, cur.slot) {
                    Ok(ev) => flight = Some((cur, ev)),
                    Err(e) => {
                        eprintln!("overlap prefill failed ({e}); dropping request");
                        let _ = cur.reply.send(Ev::Error(format!("prefill: {e}")));
                        eng.release_slot(cur.slot)?;
                    }
                }
            }
            if let Some((_, ev)) = flight.as_ref() {
                let no_live = live.iter().all(|s| s.is_none());
                if ev.is_complete() || (no_live && queue.is_empty()) {
                    let (cur, ev) = flight.take().unwrap();
                    ev.synchronize()?; // no-op when already complete
                    // Order the main stream after the flight before the slot decodes from its
                    // freshly written slab.
                    eng.wait_overlap(&ev)?;
                    finish_cursor!(cur);
                    continue;
                }
                if no_live {
                    // arrivals are waiting behind the flight but there is nothing to decode:
                    // poll instead of spinning phantom empty-width steps against the prefill
                    std::thread::sleep(std::time::Duration::from_micros(200));
                    continue;
                }
                // fall through: one normal decode step runs UNDER the in-flight prefill
            }
        }

        // 2b. MIXED STEP: a prefill cursor is outstanding. With live decoders, run ONE step
        // carrying every decode row AND the cursor's next chunk — the decode rows ride the
        // chunk's weight streaming for free, so decode neither stalls nor costs extra wall.
        // With no decoders and a fresh cursor, the head takes the ordinary graphed one-shot
        // prefill, which is strictly faster than an eager chunk pass (captured graph + tuned
        // plans) — the cold burst's FIRST request goes through here.
        if mix_on && !mix_overlap && !prefilling.is_empty() {
            let m = live
                .iter()
                .rposition(|s| s.is_some())
                .map(|i| i + 1)
                .unwrap_or(0)
                .min(slots);
            if m == 0 && prefilling.front().map(|c| c.pos) == Some(0) {
                let cur = prefilling.pop_front().unwrap();
                let r = eng
                    .stage_burst(&cur.ids)
                    .and_then(|()| eng.prefill_slot_staged(&cur.ids, 0, cur.slot));
                match r {
                    Ok(()) => finish_cursor!(cur),
                    Err(e) => {
                        let _ = cur.reply.send(Ev::Error(format!("prefill: {e}")));
                        eng.mix_unstage(cur.slot);
                        eng.release_slot(cur.slot)?;
                    }
                }
                continue;
            }
            let (c_slot, c_pos, take, c_last) = {
                let cur = prefilling.front().unwrap();
                let rem = cur.ids.len() - cur.pos;
                let take = rem.min(mix_budget);
                (cur.slot, cur.pos, take, take == rem)
            };
            if let Err(e) = eng.step_mixed(c_slot, c_pos, take, c_last, m) {
                let cur = prefilling.pop_front().unwrap();
                eprintln!("mixed step failed ({e}); dropping request");
                let _ = cur.reply.send(Ev::Error(format!("mixed prefill: {e}")));
                eng.mix_unstage(cur.slot);
                eng.release_slot(cur.slot)?;
                continue;
            }
            // decode-row bookkeeping, pipelined exactly like a batch step
            flush_pending!();
            eng.sync()?;
            if m > 0 {
                pending = Some(eng.batch_last_tokens()?);
                for (slot, cell) in live.iter().enumerate() {
                    prev_mask[slot] = cell.is_some();
                }
            }
            if c_last {
                let cur = prefilling.pop_front().unwrap();
                finish_cursor!(cur);
            } else {
                prefilling.front_mut().unwrap().pos += take;
            }
            continue;
        }

        // 3. one decode step, at a width matched to the live set.
        //
        // Slots pack low (free_slot returns the lowest free index), so the live set always fits in
        // [0, highest+1). Running the full slot count regardless meant one request paid 64
        // sequences' worth of GEMM: the server served M=1 at 251 tok/s where a 1-wide batch does
        // 326. Widths are rounded up to a power of two so the graph cache stays small — a capture
        // bakes in its grid dimensions, so each width needs its own.
        // EXACT width, not next_power_of_two: compaction (above) keeps the live set dense, so
        // `highest` IS the live count, and pow2 rounding was charging up to 2x per step across
        // the whole arrival ramp. One lazily-captured graph per width — up to `slots` graphs,
        // absorbed by the warm pass.
        let highest = live
            .iter()
            .rposition(|s| s.is_some())
            .map(|i| i + 1)
            .unwrap_or(1);
        let want = highest.min(slots);
        if eng.width() != want {
            eng.set_width(want)?;
        }
        let decode_prof = std::env::var("DECODE_PROF").is_ok();
        if !decode_prof && !eng.has_batch_graph() {
            eng.capture_batch()?;
        }
        // PIPELINED EMIT. The batch graph samples into device memory and feeds itself, so the
        // HOST never sits between steps: launch step k, then do all of step k-1's bookkeeping
        // WHILE the GPU runs, then sync and read k. The channel SENDS go to a dedicated emitter
        // thread: tokio's unbounded_send runs waker.wake() INLINE on the sender, so 64 task
        // wake-ups per step (plus the scheduler noise they drag in) were still landing on the
        // engine thread — ~2.3 ms of the 6.2 ms HTTP step at M=64 against 3.9 in-process.
        ph!(
            t_launch,
            if decode_prof {
                eng.step_batch()?
            } else {
                eng.run_graph_batch(1)?
            }
        );
        let t_e0 = std::time::Instant::now();
        flush_pending!();
        if prof {
            t_emit += t_e0.elapsed().as_secs_f64();
        }
        ph!(t_sync, eng.sync()?);
        pending = Some(ph!(t_read, eng.batch_last_tokens()?));
        for (slot, cell) in live.iter().enumerate() {
            prev_mask[slot] = cell.is_some();
        }
        if prof && nstep % 64 == 0 {
            eprintln!(
                "[prof] admit {:.1}ms launch {:.1} emit {:.1} sync {:.1} read {:.1} (per 64 steps)",
                t_admit * 1e3,
                t_launch * 1e3,
                t_emit * 1e3,
                t_sync * 1e3,
                t_read * 1e3
            );
            t_admit = 0.;
            t_launch = 0.;
            t_emit = 0.;
            t_sync = 0.;
            t_read = 0.;
        }
    }
}

async fn ids_for(app: &App, p: &PromptField) -> Vec<u32> {
    match p {
        PromptField::Ids(v) => v.clone(),
        PromptField::Text(s) => app
            .tok
            .as_ref()
            .and_then(|t| t.encode(s.as_str(), false).ok())
            .map(|e| e.get_ids().to_vec())
            .unwrap_or_default(),
        PromptField::Texts(v) => app
            .tok
            .as_ref()
            .and_then(|t| t.encode(v.first().cloned().unwrap_or_default(), false).ok())
            .map(|e| e.get_ids().to_vec())
            .unwrap_or_default(),
    }
}

async fn completions(
    State(app): State<App>,
    Json(req): Json<CompletionReq>,
) -> axum::response::Response {
    let ids = ids_for(&app, &req.prompt).await;
    let _ = req.temperature; // greedy only; the harness pins temperature 0
    run(app, ids, req.max_tokens, req.stream, req.ignore_eos).await
}

async fn chat(State(app): State<App>, Json(req): Json<ChatReq>) -> axum::response::Response {
    // minimal ChatML render; the benchmark path uses /v1/completions with ids
    let mut s = String::new();
    for m in &req.messages {
        s.push_str(&format!(
            "<|im_start|>{}\n{}<|im_end|>\n",
            m.role, m.content
        ));
    }
    s.push_str("<|im_start|>assistant\n");
    let ids = app
        .tok
        .as_ref()
        .and_then(|t| t.encode(s.as_str(), false).ok())
        .map(|e| e.get_ids().to_vec())
        .unwrap_or_default();
    run(app, ids, req.max_tokens, req.stream, req.ignore_eos).await
}

async fn run(
    app: App,
    ids: Vec<u32>,
    max_tokens: usize,
    stream: bool,
    ignore_eos: bool,
) -> axum::response::Response {
    let (tx, mut rx) = amp::unbounded_channel::<Ev>();
    if app
        .submit
        .send(Job {
            ids,
            max_tokens,
            ignore_eos,
            reply: tx,
        })
        .is_err()
    {
        return (axum::http::StatusCode::INTERNAL_SERVER_ERROR, "engine gone").into_response();
    }
    let model = app.model_name.clone();

    if stream {
        return Sse::new(SseStream {
            rx,
            model,
            ended: false,
            dt: Detok::new(app.tok.clone()),
        })
        .into_response();
    }

    let mut text = String::new();
    let mut dt = Detok::new(app.tok.clone());
    let mut out_ids: Vec<u32> = Vec::new();
    let (mut ct, mut pt) = (0usize, 0usize);
    while let Some(ev) = rx.recv().await {
        match ev {
            Ev::Token(id) => {
                out_ids.push(id);
                text.push_str(&dt.push(id));
            }
            Ev::Done {
                completion_tokens,
                prompt_tokens,
            } => {
                ct = completion_tokens;
                pt = prompt_tokens;
                break;
            }
            Ev::Error(e) => return (axum::http::StatusCode::BAD_REQUEST, e).into_response(),
        }
    }
    Json(CompletionResp {
        id: "cmpl-inferencelayer".into(),
        object: "text_completion",
        model,
        choices: vec![Choice {
            index: 0,
            text,
            finish_reason: Some("length".into()),
            ids: out_ids,
        }],
        usage: Usage {
            prompt_tokens: pt,
            completion_tokens: ct,
            total_tokens: pt + ct,
        },
    })
    .into_response()
}

/// Hand-rolled SSE stream (mirrors `serve.rs`; the tree deliberately avoids an async-stream dep).
/// One event per generated token, so the harness's inter-token latency is real.
struct SseStream {
    rx: amp::UnboundedReceiver<Ev>,
    model: String,
    ended: bool,
    dt: Detok,
}

impl futures_core::Stream for SseStream {
    type Item = Result<Event, std::convert::Infallible>;
    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        use std::task::Poll;
        if self.ended {
            return Poll::Ready(None);
        }
        match self.rx.poll_recv(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(None) => {
                self.ended = true;
                Poll::Ready(Some(Ok(Event::default().data("[DONE]"))))
            }
            Poll::Ready(Some(Ev::Token(id))) => {
                // ONE event per token, even when the text is empty (a partial multi-byte
                // character): that is what the OpenAI streaming contract promises and what makes
                // an inter-token-latency measurement comparable across engines.
                let t = self.dt.push(id);
                let j = serde_json::json!({
                    "object": "text_completion",
                    "model": self.model,
                    "choices": [{"index": 0, "text": t, "finish_reason": serde_json::Value::Null}]
                });
                Poll::Ready(Some(Ok(Event::default().data(j.to_string()))))
            }
            Poll::Ready(Some(Ev::Done { .. })) => {
                self.ended = true;
                Poll::Ready(Some(Ok(Event::default().data("[DONE]"))))
            }
            Poll::Ready(Some(Ev::Error(e))) => {
                self.ended = true;
                let j = serde_json::json!({"error": e});
                Poll::Ready(Some(Ok(Event::default().data(j.to_string()))))
            }
        }
    }
}

async fn models(State(app): State<App>) -> impl IntoResponse {
    Json(serde_json::json!({
        "object": "list",
        "data": [{"id": app.model_name, "object": "model", "owned_by": "kortexya"}]
    }))
}