inferencelayer 0.2.8

//! Micro-batched pipeline-parallel serving: the continuous-batching scheduler wired onto the
//! shard chain. This is what turns a p-stage pipeline from `1/p` utilization (M=1: one token
//! walks the chain while p−1 devices idle) into a full pipeline: requests are partitioned into
//! **micro-batch groups** (one per stage) and group g runs on stage s while group g−1 runs on
//! stage s+1 — every device busy every tick.
//!
//! Invariants that keep it sound and bitwise:
//! - A sequence never has two rounds in flight (round r+1 needs round r's token) — enforced by
//!   an explicit per-slot in-flight mark, NOT by pinning sequences to groups: any launching
//!   group packs its micro-batch from ALL ready sequences (rotating-cursor fairness). Static
//!   slot-mod-group pinning left ~1/3 of steps running 1-6 columns at full step cost once the
//!   completion tail frayed (measured 950 steps vs the 640 a packed run needs).
//! - The coordinator owns the block allocator; block-table updates travel IN-BAND ahead of the
//!   step that needs them (TCP ordering per stage connection), so every stage's paged KV layout
//!   is identical. Each stage writes only its own layers' K/V at the same (block, offset).
//! - Every column replays the same compiled batch pipelines as a solo run (per stage), and the
//!   M=1 shard gate established cross-process kernel determinism — so batched × sharded output
//!   is asserted bitwise-equal to the unsharded solo run (see `tests/shard_serve.rs`).
use std::collections::{HashMap, VecDeque};
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::mpsc::{Receiver, Sender, channel};

use anyhow::{Result, anyhow};

use crate::forward::{BatchCol, BatchPlan, MAX_SLOTS, PAGE_BLK, PAGE_ROW, StageBatchOut};
use crate::replica::FleetPlan;
use crate::server::KvPool;
use crate::shard::{Fleet, ShardClient, chain_workers, negotiate_fleet, read_done, validate_chain};
use crate::{GpuCtx, Lfm2Gpu, Weights};

/// Which transport carries the µbatches between stages (see `shard.rs` module docs).
enum Transport<'a> {
    /// Coordinator-relayed (default): one forwarding thread per stage; every boundary crosses
    /// the coordinator's network twice.
    Hub,
    /// TRUE P2P: workers forward directly to the next peer; only tokens return, on the
    /// coordinator's sink (`ret` = explicit `host:port`, else auto-bound + peer-derived).
    P2p { ret: Option<&'a str> },
}

enum SeqState {
    Prefill { done: usize },
    Decode,
}

struct Seq {
    id: u64,
    slot: usize,
    tokens: Vec<u32>,
    prompt_len: usize,
    blocks: Vec<u32>,
    state: SeqState,
    max_tokens: usize,
    emitted: Vec<u32>,
    /// Serving-MTP: drafts for the next verify span (len == the occupancy-gated depth at the
    /// round that drafted them; may differ from the current depth across a threshold).
    pending_drafts: Vec<u32>,
    drafts_ready: bool,
}

/// A micro-batch in flight: block-table updates ride ahead of the step on each stage.
struct MbMsg {
    group: usize,
    /// Serving-MTP flags (see [`ShardClient::bstep_grouped`]): the LAST stage spontaneously
    /// accepts/seeds/drafts and appends the drafts to its token response — zero extra trips.
    spec_on: bool,
    spec_verify: bool,
    spec_k: usize,
    /// Route to the WIDE (KC=16) prefill plan on every stage.
    wide: bool,
    /// Slots to zero (fresh admissions) — applied on every stage BEFORE the step (in-band).
    zslots: Vec<u32>,
    /// DN rollbacks `(slot, col, snapshot_group)` — applied on every stage before this step.
    /// Entries ride the FIRST µbatch sent after they are queued (any group): in-band stream
    /// order guarantees they land after the µbatch that made the snapshot and before any later
    /// verify of that slot; each entry names the group whose plan holds its snapshot.
    restores: Vec<(u32, u32, u32)>,
    btab: Vec<(u32, Vec<u32>)>,
    /// Per column: (pos, btrow, need_logit, token) — tokens feed a stage_first worker directly.
    cols: Vec<(u32, u32, u32, u32)>,
    hidden: Vec<f32>,
}

struct MbDone {
    group: usize,
    tokens: Vec<u32>,
}

/// What a group has in flight (see [`PipelineServe::inflight`]).
enum Round {
    /// `k` = draft depth the last stage appended per emitting column (0 = none this round).
    Plain(Vec<(usize, bool)>, usize),
    /// `k_verify` = each span's drafted length (uniform per round); `k_next` = the re-draft
    /// depth the last stage appended per span (0 = drain: accept without re-drafting).
    Spec(Vec<SpanOwner>, usize, usize),
}

/// One sequence's verify span within a spec µbatch.
struct SpanOwner {
    slot: usize,
    base: usize,
    drafts: Vec<u32>,
}

pub use crate::sampling::prompt_lookup_drafts;

/// Pipeline-parallel micro-batched scheduler (see module docs). Greedy-only.
/// Live per-remote-stage counters, updated lock-free by each stage's forwarding thread and read
/// by the coordinator for the `/stats` dashboard. Each thread touches only its OWN index, so the
/// updates never contend.
#[derive(Default)]
struct StageMeter {
    steps: AtomicU64,
    cols: AtomicU64,
    nanos: AtomicU64,
}

/// A snapshot of one pipeline stage's identity + live throughput — the orchestrator's view of
/// what each GPU donor is doing. Returned by [`PipelineServe::stage_stats`].
pub struct StageStat {
    /// Stage index (0 = coordinator-local / headless, then one per remote donor in chain order).
    pub stage: usize,
    /// `"<Backend>/<adapter>"` for a remote donor, or `"local <backend>"`/`"headless"` for stage 0.
    pub backend: String,
    /// First layer this stage owns (inclusive).
    pub start: usize,
    /// One past the last layer.
    pub end: usize,
    /// Micro-batch steps this stage processed since connect (0 for the untimed local stage).
    pub steps: u64,
    /// Columns (≈tokens) routed through this stage.
    pub cols: u64,
    /// Mean round-trip time the coordinator waited on this stage, in ms (0 for local; Hub only —
    /// under TRUE-P2P the inner boundaries never touch the coordinator, so it can't time them).
    pub avg_ms: f64,
}

pub struct PipelineServe {
    /// Local stage 0 (None = HEADLESS coordinator: pure orchestrator, zero GPU work — the
    /// pipeline-filling configuration; workers carry every layer, tokens ride the first hop).
    local: Option<(GpuCtx, Lfm2Gpu, BatchPlan, Option<BatchPlan>)>,
    prefill_k: usize,
    dn_local: bool,
    pool: KvPool,
    to_first: Sender<MbMsg>,
    done_rx: Receiver<MbDone>,
    _threads: Vec<std::thread::JoinHandle<()>>,
    n_groups: usize,
    mb_k: usize,
    slots: Vec<Option<Seq>>,
    queue: VecDeque<Seq>,
    next_id: u64,
    eos: Vec<u32>,
    /// Per group: the round in flight (None = idle). Plain = one column per owner;
    /// Spec = verify spans (slot, base column, drafted tokens); Draft = slots awaiting drafts.
    inflight: Vec<Option<Round>>,
    /// Serving-MTP BASE draft depth (env `OSFKB_SERVE_MTP`, 0 = off). The per-round depth is
    /// occupancy-gated — see [`Self::spec_depth_for`].
    spec_k: usize,
    /// Occupancy schedule: decode concurrency ≤ `spec_c_full` runs the base depth; ≤
    /// `spec_c_half` runs `min(base, 4)`; above it speculation is OFF (on MoE the verify cost
    /// scales with batch × span — past the knee it is a net loss). Envs `OSFKB_SPEC_C_FULL` /
    /// `OSFKB_SPEC_C_HALF`; defaults 4 / 8 are PROVISIONAL until the measured A3B knee (M0-b)
    /// pins them.
    spec_c_full: usize,
    spec_c_half: usize,
    /// Spec verify rounds launched (observability + the occupancy-gate test hook).
    spec_rounds: u64,
    /// 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 in the sequence's own
    /// prompt+history. Free acceptance on copy/echo spans; wrong drafts only cost speed.
    pld: bool,
    /// Rounds whose drafts came from the prompt lookup (observability + the PLD test hook).
    pld_hits: u64,
    /// Slots freshly admitted since the last launched micro-batch (recurrent-state zeroing rides
    /// in-band to every stage with the next step).
    fresh_slots: Vec<u32>,
    /// Slots with a round in flight — the ONE ordering invariant (see module docs). A launching
    /// group packs from every slot not marked here.
    slot_busy: Vec<bool>,
    /// Rotating scan start so all ready sequences advance uniformly (completion stays bunched,
    /// which is what keeps the tail's micro-batches full).
    scan_cursor: usize,
    /// DN rollbacks awaiting the next outbound µbatch (serving-MTP partial accepts); each
    /// carries the group whose plan holds its snapshot.
    pending_restores: Vec<(u32, u32, u32)>,
    /// P2P transport only: control connections for stages 1..n (chained worker→worker), kept
    /// for the explicit shutdown that unwinds the chain. Empty under hub-spoke.
    extra_clients: Vec<ShardClient>,
    /// Live per-remote-stage meters (index 0 = first worker stage). Updated by the Hub forwarding
    /// threads; empty of updates under P2P (inner boundaries bypass the coordinator).
    stage_meters: Arc<Vec<StageMeter>>,
    /// Per-stage identity for the dashboard: index 0 = coordinator (local/headless), then one per
    /// remote donor in chain order. `(backend, start, end)`.
    stage_desc: Vec<(String, usize, usize)>,
}

impl PipelineServe {
    /// Load stage 0 (`0..split`) locally, connect + initialize the worker chain, and spawn one
    /// forwarding thread per remote stage. `mb_k` = columns per micro-batch (per group).
    pub fn connect(
        dir: &Path,
        split: usize,
        workers: &[&str],
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
    ) -> Result<Self> {
        Self::connect_impl(dir, split, workers, mb_k, prefill_k, eos, Transport::Hub)
    }

    /// [`Self::connect`] over the TRUE-P2P transport: µbatches enter at worker 0, each stage
    /// forwards its residual (and the in-band control ops) straight to the next peer, and only
    /// tokens return on the coordinator's sink. Validates that the worker ranges tile the model
    /// and prints the (possibly heterogeneous) topology before any tokens flow.
    pub fn connect_p2p(
        dir: &Path,
        split: usize,
        workers: &[&str],
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
        ret: Option<&str>,
    ) -> Result<Self> {
        Self::connect_impl(
            dir,
            split,
            workers,
            mb_k,
            prefill_k,
            eos,
            Transport::P2p { ret },
        )
    }

    /// [`Self::connect`] over a pre-registered [`Fleet`] (reverse-joined workers — see
    /// `accept_fleet`): the workers dialed US, so no worker addresses are needed.
    pub fn connect_fleet(
        dir: &Path,
        split: usize,
        fleet: Fleet,
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
    ) -> Result<Self> {
        Self::init_with(dir, split, fleet, mb_k, prefill_k, eos, Transport::Hub)
    }

    /// [`Self::connect_p2p`] over a pre-registered [`Fleet`]. NAT note: a joined worker's
    /// DATA address must still be reachable from its predecessor — put NAT'd machines first
    /// in the join order (stage 1 receives via its own dialed-out control socket).
    pub fn connect_fleet_p2p(
        dir: &Path,
        split: usize,
        fleet: Fleet,
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
        ret: Option<&str>,
    ) -> Result<Self> {
        Self::init_with(
            dir,
            split,
            fleet,
            mb_k,
            prefill_k,
            eos,
            Transport::P2p { ret },
        )
    }

    fn connect_impl(
        dir: &Path,
        split: usize,
        workers: &[&str],
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
        transport: Transport<'_>,
    ) -> Result<Self> {
        let fleet = Fleet {
            clients: workers
                .iter()
                .map(|a| ShardClient::connect(a))
                .collect::<Result<Vec<_>>>()?,
            data_addrs: workers.iter().map(|a| a.to_string()).collect(),
        };
        Self::init_with(dir, split, fleet, mb_k, prefill_k, eos, transport)
    }

    fn init_with(
        dir: &Path,
        split: usize,
        mut fleet: Fleet,
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
        transport: Transport<'_>,
    ) -> Result<Self> {
        anyhow::ensure!(!fleet.clients.is_empty(), "need at least one remote stage");
        anyhow::ensure!(
            fleet.clients.len() == fleet.data_addrs.len(),
            "fleet clients/data addresses mismatch"
        );
        anyhow::ensure!((1..=MAX_SLOTS).contains(&mb_k), "mb_k in 1..=MAX_SLOTS");
        anyhow::ensure!(prefill_k <= MAX_SLOTS, "prefill_k ≤ MAX_SLOTS");
        // split == 0 → HEADLESS: the coordinator never touches a GPU; worker 0 owns the
        // embedding (stage_first) and the first hop carries token ids, not hidden states.
        let local = if split > 0 {
            let ctx = GpuCtx::new()?;
            let w = Weights::load_shard(&ctx, dir, 0, split)?;
            let gpu = Lfm2Gpu::new(&ctx, w);
            let bp = gpu.make_batch_plan(&ctx, mb_k);
            let pk = if prefill_k > 0 && !ctx.subgroups {
                0
            } else {
                prefill_k
            };
            let bpw = (pk > 0).then(|| gpu.make_batch_plan_wide(&ctx, pk));
            Some((ctx, gpu, bp, bpw))
        } else {
            None
        };
        let prefill_k = match &local {
            Some((_, _, _, bpw)) if bpw.is_none() => 0,
            _ => prefill_k,
        };
        let dn_local = local
            .as_ref()
            .map(|(_, g, _, _)| g.w.cfg.layer_is_attn.iter().any(|a| !a))
            .unwrap_or(false);
        // Pool sizing must match the WORKERS' pool. That invariant is what forces a CONSTANT here
        // rather than anything derived from local VRAM: the coordinator hands out block IDs and the
        // workers index their own KV buffers by them, so every process must end up with the same
        // block COUNT. A per-process VRAM probe would silently desync them. Every process reads this
        // same env with this same default, so they agree by construction.
        //
        // The old default of 4096 tokens was starving the fleet. It is not a VRAM budget — it is a
        // flat token count — and 48 concurrent requests × (17-token prompt + 128 generated) = 6960
        // tokens simply does not fit in it, so most of the working set queued instead of running.
        // Measured on the 35B-A3B, 4-stage on 4× V100, 48 concurrent, mb_k 8 (locked clocks, quiet
        // box):
        //
        //     KV_POOL_TOKENS  4096  →  448.8 tok/s      (the old default)
        //     KV_POOL_TOKENS 16384  →  540.7 tok/s      ← +20%, and past the 462.4 crown
        //     KV_POOL_TOKENS 32768  →  536.9 tok/s      (plateau — 16384 is the knee)
        //
        // This `KvPool` is only the block-ID ALLOCATOR. The GPU memory those ids actually index is
        // the ENGINE's KV buffer, sized by `EngineOpts::sized_from_env` — so the two MUST agree, or
        // the allocator hands out block ids past the end of the buffer and we get silent
        // out-of-bounds KV writes on every process.
        //
        // They used to agree only by coincidence: this site did its own `OSFKB_KV_POOL_TOKENS` read
        // with its own `4096` default, and the engine did a separate one defaulting to `MAX_T`,
        // which happens to be 4096 too. Two independently-maintained numbers that must be equal is
        // a bug waiting to be committed — and I nearly committed it, by raising this one alone to
        // capture a measured +20% (see below). Deriving BOTH from the same function makes the
        // divergence unrepresentable. Behaviour-preserving today; the point is that it stays so.
        //
        // Every process — coordinator and workers — resolves this identically because they all call
        // the same function and read the same env, which is what makes the fleet's shared block ids
        // meaningful in the first place.
        //
        // The default (MAX_T = 4096 tokens) is small, and it costs real throughput: 48 concurrent ×
        // (17-token prompt + 128 generated) = 6960 tokens does not fit, so most of the working set
        // queues instead of running. Measured on the 35B-A3B, 4 stages on 4× V100, 48 concurrent,
        // mb_k 8 (locked clocks, quiet box):
        //
        //     OSFKB_KV_POOL_TOKENS= 4096  →  448.8 tok/s   (the default)
        //     OSFKB_KV_POOL_TOKENS=16384  →  540.7 tok/s   ← +20%, past the 462.4 crown
        //     OSFKB_KV_POOL_TOKENS=32768  →  536.9 tok/s   (plateau — 16384 is the knee)
        //
        // Raising the default is NOT free: it is a 4× larger KV allocation on every process, and it
        // would also hit SOLO decode, which serves one sequence and would burn that VRAM for
        // nothing. Doing it properly means a fleet-vs-solo split plus coordinator/worker skew
        // validation (an old worker at 4096 against a new coordinator at 16384 is the same desync,
        // across processes). Filed as S-C1 — not smuggled in here.
        let pool_blocks = crate::forward::EngineOpts::sized_from_env().pool_tokens / PAGE_BLK;
        let pool = KvPool::new(pool_blocks);
        let n_stages = fleet.clients.len();
        let n_groups_early = std::env::var("OSFKB_PIPE_GROUPS")
            .ok()
            .and_then(|v| v.parse::<usize>().ok())
            .unwrap_or(n_stages + 1)
            .max(n_stages + 1);
        for c in &mut fleet.clients {
            c.chain_clear()?; // shed any stale P2P routing before configuring this session
        }
        // Fleet negotiation BEFORE binit: auto stages get their layer range assigned here,
        // and batch plans can only build against a loaded engine.
        let infos = negotiate_fleet(&mut fleet.clients, split, dir)?;
        eprintln!("{}", validate_chain(split, &infos)?);
        // Dashboard identity: stage 0 (coordinator) then each remote donor in chain order, plus a
        // lock-free meter per remote stage the forwarding threads will update.
        let mut stage_desc: Vec<(String, usize, usize)> = vec![(
            match &local {
                Some((ctx, ..)) => format!("local {}", ctx.backend),
                None => "headless".to_string(),
            },
            0,
            split,
        )];
        for inf in &infos {
            stage_desc.push((inf.backend.clone(), inf.start, inf.end));
        }
        let stage_meters = Arc::new(
            (0..n_stages)
                .map(|_| StageMeter::default())
                .collect::<Vec<_>>(),
        );
        for c in &mut fleet.clients {
            c.binit_grouped(mb_k, prefill_k, n_groups_early)?;
        }
        // Ship the donor-display vocab ONCE before serving (fire-and-forget, off the hot path): it
        // lets each donor render the text it processes. A missing/unrecognized tokenizer.json just
        // means donors show ids or nothing — never a serving error.
        if let Some(blob) = std::fs::read(dir.join("tokenizer.json"))
            .ok()
            .and_then(|b| crate::vocab::build_vocab_blob(&b))
        {
            for c in &mut fleet.clients {
                let _ = c.send_vocab(&blob);
            }
        }
        let Fleet {
            mut clients,
            data_addrs,
        } = fleet;
        let (to_first, mut prev_rx) = channel::<MbMsg>();
        let (done_tx, done_rx) = channel::<MbDone>();
        let mut threads = Vec::with_capacity(clients.len() + 1);
        let mut extra_clients = Vec::new();
        match transport {
            Transport::Hub => {
                // Chain: main → stage1 → … → stageN → done. Each thread owns its client; a
                // middle stage forwards the residual stream, the last stage emits tokens.
                let n = clients.len();
                for (i, mut client) in clients.into_iter().enumerate() {
                    let last = i + 1 == n;
                    let (tx_next, rx_next) = channel::<MbMsg>();
                    let rx = std::mem::replace(&mut prev_rx, rx_next);
                    let done_tx = done_tx.clone();
                    let meters = stage_meters.clone();
                    threads.push(std::thread::spawn(move || {
                        while let Ok(msg) = rx.recv() {
                            for slot in &msg.zslots {
                                if client.zslot(*slot).is_err() {
                                    return;
                                }
                            }
                            for (slot, col, snap_group) in &msg.restores {
                                if client
                                    .dn_restore_grouped(*slot, *col as usize, *snap_group as usize)
                                    .is_err()
                                {
                                    return;
                                }
                            }
                            for (row, blocks) in &msg.btab {
                                if client.btab(*row, blocks).is_err() {
                                    return;
                                }
                            }
                            // Meter this stage's round trip for /stats (this thread is the sole
                            // writer of meters[i], so Relaxed ordering is sufficient).
                            let ncols = msg.cols.len() as u64;
                            let t0 = std::time::Instant::now();
                            let stepped = client.bstep_grouped(
                                &msg.cols,
                                &msg.hidden,
                                msg.wide,
                                msg.group,
                                (msg.spec_on, msg.spec_verify, msg.spec_k),
                            );
                            let m = &meters[i];
                            m.steps.fetch_add(1, Ordering::Relaxed);
                            m.cols.fetch_add(ncols, Ordering::Relaxed);
                            m.nanos
                                .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
                            match stepped {
                                Ok(StageBatchOut::Hidden(h)) => {
                                    if last {
                                        return; // protocol error: last stage must produce tokens
                                    }
                                    let _ = tx_next.send(MbMsg {
                                        group: msg.group,
                                        spec_on: msg.spec_on,
                                        spec_verify: msg.spec_verify,
                                        spec_k: msg.spec_k,
                                        wide: msg.wide,
                                        zslots: msg.zslots,
                                        restores: msg.restores,
                                        btab: msg.btab,
                                        cols: msg.cols,
                                        hidden: h,
                                    });
                                }
                                Ok(StageBatchOut::Tokens(t)) => {
                                    if !last {
                                        return; // protocol error: only the last stage has a head
                                    }
                                    let _ = done_tx.send(MbDone {
                                        group: msg.group,
                                        tokens: t,
                                    });
                                }
                                Err(_) => return,
                            }
                        }
                        client.shutdown();
                    }));
                }
            }
            Transport::P2p { ret } => {
                // Topology already negotiated + validated above; wire the chain, then run
                // exactly TWO coordinator threads: a sender that feeds worker 0 (control ops
                // in-band ahead of each µbatch — each stage applies AND forwards them) and a
                // sink reader that demultiplexes finished groups. Inter-stage residual
                // traffic never touches the coordinator.
                let refs: Vec<&str> = data_addrs.iter().map(String::as_str).collect();
                // `OSFKB_WEBRTC_HOPS=i,j,…` FORCES those worker→worker hops onto a WebRTC data
                // channel (deterministic demo / config). Independently, `auto_webrtc` (on when this
                // coordinator carries the `webrtc` endpoint) falls back to WebRTC for any OTHER hop
                // whose direct TCP/WS dial fails (peers behind NAT). Empty + no auto ⇒ all-TCP.
                let webrtc_hops: Vec<usize> = std::env::var("OSFKB_WEBRTC_HOPS")
                    .ok()
                    .map(|s| s.split(',').filter_map(|x| x.trim().parse().ok()).collect())
                    .unwrap_or_default();
                let mut sink = chain_workers(
                    &mut clients,
                    &refs,
                    ret,
                    &webrtc_hops,
                    cfg!(feature = "webrtc"),
                )?;
                let mut c0 = clients.remove(0);
                extra_clients = clients;
                let rx = prev_rx;
                threads.push(std::thread::spawn(move || {
                    while let Ok(msg) = rx.recv() {
                        for slot in &msg.zslots {
                            if c0.zslot(*slot).is_err() {
                                return;
                            }
                        }
                        for (slot, col, snap_group) in &msg.restores {
                            if c0
                                .dn_restore_grouped(*slot, *col as usize, *snap_group as usize)
                                .is_err()
                            {
                                return;
                            }
                        }
                        for (row, blocks) in &msg.btab {
                            if c0.btab(*row, blocks).is_err() {
                                return;
                            }
                        }
                        if c0
                            .bstep_send_grouped(
                                &msg.cols,
                                &msg.hidden,
                                msg.wide,
                                msg.group,
                                (msg.spec_on, msg.spec_verify, msg.spec_k),
                            )
                            .is_err()
                        {
                            return;
                        }
                    }
                    c0.shutdown();
                }));
                threads.push(std::thread::spawn(move || {
                    loop {
                        match read_done(&mut sink) {
                            Ok((group, tokens)) => {
                                if done_tx.send(MbDone { group, tokens }).is_err() {
                                    return;
                                }
                            }
                            Err(_) => return, // chain unwound (last stage cleared its hop)
                        }
                    }
                }));
            }
        }
        // Groups = in-flight micro-batches. n_stages+1 fills the pipe only if the coordinator
        // relaunches instantly; real relaunch latency adds bubbles, so OVERSUBSCRIBE: more
        // groups smooth the return jitter (sweep on the 35B; sequences spread thinner per
        // group, so mb_k budgets stay full only while concurrency >= groups * mb_k).
        let n_groups = n_groups_early;
        Ok(Self {
            local,
            prefill_k,
            dn_local,
            pool,
            to_first,
            done_rx,
            _threads: threads,
            n_groups,
            mb_k,
            slots: (0..MAX_SLOTS).map(|_| None).collect(),
            queue: VecDeque::new(),
            next_id: 1,
            eos,
            inflight: (0..n_groups).map(|_| None).collect(),
            spec_k: std::env::var("OSFKB_SERVE_MTP")
                .ok()
                .and_then(|v| v.parse().ok())
                .unwrap_or(0),
            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),
            spec_rounds: 0,
            pld: std::env::var("OSFKB_PLD").ok().as_deref() != Some("0"),
            pld_hits: 0,
            fresh_slots: Vec::new(),
            slot_busy: vec![false; MAX_SLOTS],
            scan_cursor: 0,
            pending_restores: Vec::new(),
            extra_clients,
            stage_meters,
            stage_desc,
        })
    }

    /// Per-stage routing snapshot for the orchestrator dashboard — what each GPU donor is doing
    /// (identity + live steps / tokens / mean round-trip latency). Reads the lock-free stage
    /// meters, so it's cheap to poll between turns.
    pub fn stage_stats(&self) -> Vec<StageStat> {
        self.stage_desc
            .iter()
            .enumerate()
            .map(|(j, (backend, start, end))| {
                let (steps, cols, avg_ms) = if j == 0 {
                    (0, 0, 0.0) // coordinator-local stage — not round-trip-timed
                } else {
                    match self.stage_meters.get(j - 1) {
                        Some(m) => {
                            let steps = m.steps.load(Ordering::Relaxed);
                            let avg_ms = if steps > 0 {
                                m.nanos.load(Ordering::Relaxed) as f64 / steps as f64 / 1e6
                            } else {
                                0.0
                            };
                            (steps, m.cols.load(Ordering::Relaxed), avg_ms)
                        }
                        None => (0, 0, 0.0),
                    }
                };
                StageStat {
                    stage: j,
                    backend: backend.clone(),
                    start: *start,
                    end: *end,
                    steps,
                    cols,
                    avg_ms,
                }
            })
            .collect()
    }

    /// Override the occupancy schedule (see the field docs): decode concurrency ≤ `full` runs
    /// the base draft depth, ≤ `half` runs `min(base, 4)`, above it speculation is off.
    pub fn set_spec_schedule(&mut self, full: usize, half: usize) {
        self.spec_c_full = full;
        self.spec_c_half = half;
    }

    /// Spec verify rounds launched so far (observability; the occupancy-gate test hook).
    pub fn spec_rounds(&self) -> u64 {
        self.spec_rounds
    }

    /// Rounds whose drafts came from the prompt lookup instead of the MTP head.
    pub fn pld_hits(&self) -> u64 {
        self.pld_hits
    }

    /// 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; counting only decoders would let a burst of admissions
    /// draft at "low occupancy" and immediately regret it.
    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
        }
    }

    /// Enqueue a request (prompt token ids); greedy decoding.
    pub fn submit(&mut self, prompt: Vec<u32>, max_tokens: usize) -> Result<u64> {
        anyhow::ensure!(!prompt.is_empty(), "empty prompt");
        anyhow::ensure!(
            prompt.len() + max_tokens <= PAGE_ROW * PAGE_BLK,
            "request exceeds the per-sequence token cap"
        );
        let id = self.next_id;
        self.next_id += 1;
        self.queue.push_back(Seq {
            id,
            slot: usize::MAX,
            tokens: prompt.clone(),
            prompt_len: prompt.len(),
            blocks: Vec::new(),
            state: SeqState::Prefill { done: 0 },
            max_tokens,
            emitted: Vec::new(),
            pending_drafts: Vec::new(),
            drafts_ready: false,
        });
        Ok(id)
    }

    /// Worst-case KV blocks a sequence will ever need: its whole prompt plus every token it may
    /// still emit. This is what must be RESERVED, not what it holds right now — blocks are grown
    /// lazily by `ensure_blocks`, so a sequence that fits today can starve tomorrow.
    fn worst_case_blocks(seq: &Seq) -> usize {
        (seq.prompt_len + seq.max_tokens).div_ceil(PAGE_BLK)
    }

    /// Blocks the already-admitted sequences have not yet allocated but will still demand.
    fn committed_blocks(&self) -> usize {
        self.slots
            .iter()
            .flatten()
            .map(|s| Self::worst_case_blocks(s).saturating_sub(s.blocks.len()))
            .sum()
    }

    /// Admit queued requests into free slots — but ONLY as many as the KV pool can actually carry
    /// to completion.
    ///
    /// This used to admit unconditionally into any free slot (there are `MAX_SLOTS` = 192 of them)
    /// without ever consulting the pool. With more concurrent requests than the pool has blocks for,
    /// every one of them got a slot, `launch_group` then failed to grow ANY of them
    /// (`ensure_blocks` → pool exhaustion), nothing was in flight, and `run_streaming` hard-errored
    /// with "pipeline stalled with N pending". Measured on the 35B-A3B 4-stage fleet: 8 concurrent
    /// served fine at 282.8 tok/s, 48 concurrent DEADLOCKED. A serving path must degrade — run what
    /// fits and queue the rest — not admit work it cannot finish and then die.
    ///
    /// Reserving the WORST case (prompt + every token still to emit) is what makes it deadlock-free:
    /// once admitted, a sequence can always grow to completion, so the pipeline can always drain and
    /// free blocks for the queue behind it.
    fn admit(&mut self) {
        while !self.queue.is_empty() {
            let Some(slot) = self.slots.iter().position(Option::is_none) else {
                return;
            };
            let need = Self::worst_case_blocks(self.queue.front().expect("queue non-empty"));
            let free = self
                .pool
                .available()
                .saturating_sub(self.committed_blocks());
            if need > free {
                // Nothing more fits. If NOTHING is resident either, this request can never run in
                // this pool at all — fail loudly rather than spin, because no amount of draining
                // will make room.
                if self.slots.iter().all(Option::is_none) {
                    return; // run_streaming's `launched`/`pending` check reports it
                }
                return; // drain what is resident; the queue is retried every loop
            }
            let mut seq = self.queue.pop_front().expect("queue non-empty");
            seq.slot = slot;
            // Fresh recurrent state locally; remote stages get it in-band with the first step.
            if let Some((ctx, gpu, _, _)) = &self.local
                && self.dn_local
            {
                gpu.zero_dn_slot(ctx, slot + 1);
            }
            self.fresh_slots.push((slot + 1) as u32);
            self.slots[slot] = Some(seq);
        }
    }

    /// Grow `seq.blocks` to cover `need` positions; records the full row for broadcast when it
    /// grew. Returns false on pool exhaustion.
    fn ensure_blocks(
        pool: &mut KvPool,
        seq: &mut Seq,
        need: usize,
        updates: &mut Vec<(u32, Vec<u32>)>,
    ) -> bool {
        let want = need.div_ceil(PAGE_BLK);
        let mut grew = false;
        while seq.blocks.len() < want {
            match pool.alloc() {
                Some(b) => {
                    seq.blocks.push(b);
                    grew = true;
                }
                None => return false,
            }
        }
        if grew {
            let row = (seq.slot + 1) as u32;
            updates.retain(|(r, _)| *r != row);
            updates.push((row, seq.blocks.clone()));
        }
        true
    }

    /// Build group `g`'s next micro-batch, run stage 0, ship it. Returns false if `g` has no work.
    /// `g` is only the in-flight token: the micro-batch packs from EVERY ready (non-busy) slot,
    /// scanning from a rotating cursor so sequences advance uniformly.
    fn launch_group(&mut self, g: usize) -> Result<bool> {
        // Spec verify first; no spec-ready sequences ⇒ fall through to the plain packer
        // (prefill chunks and sequences awaiting their first draft ride the normal path).
        if self.spec_k > 0
            && let Some(sent) = self.launch_group_spec(g)?
        {
            return Ok(sent);
        }
        let mut cols_meta: Vec<(u32, u32, u32, u32)> = Vec::new();
        let mut cols: Vec<BatchCol> = Vec::new();
        let mut owners: Vec<(usize, bool)> = Vec::new();
        let mut updates: Vec<(u32, Vec<u32>)> = Vec::new();
        let n_slots = self.slots.len();
        // Decode columns first, packed from every ready sequence — but RESERVE a slice for
        // pending prefill chunks: with global packing a full decode pool otherwise owns every
        // column and admissions serialize into a thin completion tail (measured: the dip to
        // ncols 13 mid-run + ~100 trailing 1-4-column steps).
        let prefill_pending = (0..n_slots).any(|s| {
            !self.slot_busy[s]
                && matches!(
                    self.slots[s].as_ref().map(|q| &q.state),
                    Some(SeqState::Prefill { .. })
                )
        });
        let decode_budget = if prefill_pending {
            self.mb_k - (self.mb_k / 4).max(1)
        } else {
            self.mb_k
        };
        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 self.slot_busy[slot] {
                continue;
            }
            let Some(seq) = self.slots[slot].as_mut() else {
                continue;
            };
            if let SeqState::Decode = seq.state {
                if self.spec_k > 0 && seq.drafts_ready {
                    // Parked for a spec span (anti-convoy hold) — never re-decode it plainly.
                    continue;
                }
                let pos = seq.tokens.len() as u32 - 1;
                if !Self::ensure_blocks(&mut self.pool, seq, pos as usize + 1, &mut updates) {
                    continue;
                }
                cols.push(BatchCol::text(
                    *seq.tokens.last().expect("decode has tokens"),
                    pos,
                    (slot + 1) as u32,
                    true,
                ));
                cols_meta.push((
                    pos,
                    (slot + 1) as u32,
                    1,
                    *seq.tokens.last().expect("decode"),
                ));
                owners.push((slot, true));
                self.slot_busy[slot] = true;
                cursor = (slot + 1) % n_slots;
            }
        }
        if !cols.is_empty() {
            self.scan_cursor = cursor;
        }
        // Then prefill chunks — from as MANY prefilling sequences as fit (one sequence per
        // micro-batch serialized 160 admissions into 160 launches; a wide launch carried 12 of
        // its 64 columns). A launch with NO decode work routes through the WIDE (KC=16) plan at
        // prefill_k width — 4× weight amortization and 4× fewer per-step syncs on the prompt
        // phase (the TTFT path).
        let wide = cols.is_empty() && self.prefill_k > 0;
        let budget = if wide { self.prefill_k } else { self.mb_k };
        for slot in 0..n_slots {
            if cols.len() >= budget {
                break;
            }
            if self.slot_busy[slot] {
                continue;
            }
            let Some(seq) = self.slots[slot].as_mut() else {
                continue;
            };
            if let SeqState::Prefill { done } = seq.state {
                let chunk_start = cols.len();
                let mut p = done;
                while cols.len() < budget && p < seq.prompt_len {
                    if !Self::ensure_blocks(&mut self.pool, seq, p + 1, &mut updates) {
                        break;
                    }
                    let last = p + 1 == seq.prompt_len;
                    cols.push(BatchCol::text(
                        seq.tokens[p],
                        p as u32,
                        (slot + 1) as u32,
                        last,
                    ));
                    cols_meta.push((p as u32, (slot + 1) as u32, u32::from(last), seq.tokens[p]));
                    owners.push((slot, last));
                    p += 1;
                }
                seq.state = if p == seq.prompt_len {
                    SeqState::Decode // flips on the returned token below
                } else {
                    SeqState::Prefill { done: p }
                };
                if cols.len() > chunk_start {
                    self.slot_busy[slot] = true;
                }
            }
        }
        let wide = wide && !cols.is_empty();
        if cols.is_empty() {
            return Ok(false);
        }
        // Occupancy-gated draft depth for THIS round's spontaneous drafts (0 = none; the
        // decode output itself is depth-independent, so gating only moves throughput).
        let k_eff = self.spec_depth_for(self.live_occupancy());
        // Local stage 0 (when present): apply block tables, run, ship the residual stream.
        // HEADLESS: ship the tokens directly — worker 0 embeds them (4 B/col vs h·4 B/col),
        // and the coordinator does ZERO GPU work, so all groups' stage-0 serialization vanishes.
        let hidden = if let Some((ctx, gpu, bp, bpw)) = &self.local {
            for (row, blocks) in &updates {
                gpu.write_btab_row(ctx, *row, blocks);
            }
            let plan = if wide {
                bpw.as_ref().expect("wide routed only when built")
            } else {
                bp
            };
            match gpu.batch_stage_step(ctx, plan, &cols, None)? {
                StageBatchOut::Hidden(h) => h,
                StageBatchOut::Tokens(_) => return Err(anyhow!("stage 0 must not be last")),
            }
        } else {
            Vec::new()
        };
        self.to_first
            .send(MbMsg {
                group: g,
                spec_on: k_eff > 0 && !wide,
                spec_verify: false,
                spec_k: k_eff,
                wide,
                zslots: std::mem::take(&mut self.fresh_slots),
                restores: std::mem::take(&mut self.pending_restores),
                btab: updates,
                cols: cols_meta,
                hidden,
            })
            .map_err(|_| anyhow!("stage chain died"))?;
        self.inflight[g] = Some(Round::Plain(owners, if wide { 0 } else { k_eff }));
        Ok(true)
    }

    /// Spec verify packer: spans of (1 + drafted-len) same-slot columns per draft-ready
    /// sequence. Spans in one round share their drafted length (depth changes across an
    /// occupancy threshold produce transiently smaller rounds, never mixed arithmetic).
    /// Returns None when no sequence is ready (caller falls back to the plain packer).
    fn launch_group_spec(&mut self, g: usize) -> Result<Option<bool>> {
        let n_slots = self.slots.len();
        // The round's verify length = the first ready sequence's drafted length (scan order —
        // every parked sequence is eventually first, so no depth starves).
        let ready_len = |q: &Seq| matches!(q.state, SeqState::Decode) && q.drafts_ready;
        let Some(k_verify) = (0..n_slots)
            .map(|i| (self.scan_cursor + i) % n_slots)
            .find_map(|s| {
                (!self.slot_busy[s])
                    .then(|| self.slots[s].as_ref())
                    .flatten()
                    .filter(|q| ready_len(q))
                    .map(|q| q.pending_drafts.len())
            })
        else {
            return Ok(None);
        };
        let span = 1 + k_verify;
        let budget = self.mb_k / span;
        if budget == 0 {
            return Ok(None);
        }
        // Anti-convoy: launching the moment ONE span is ready locks the pipeline into tiny
        // µbatches (each done re-readies only its own spans). Hold while other groups are in
        // flight and this launch would be under-filled relative to what will free up.
        let ready = (0..n_slots)
            .filter(|s| {
                !self.slot_busy[*s]
                    && self.slots[*s]
                        .as_ref()
                        .is_some_and(|q| ready_len(q) && q.pending_drafts.len() == k_verify)
            })
            .count();
        if ready == 0 {
            return Ok(None);
        }
        let any_inflight = self.inflight.iter().any(Option::is_some);
        if any_inflight && ready < budget {
            // Hold: fall through to the plain packer (prefill keeps flowing); the ready spans
            // stay parked (the plain packer skips draft-ready sequences below).
            return Ok(None);
        }
        let mut spans: Vec<SpanOwner> = Vec::new();
        let mut cols_meta: Vec<(u32, u32, u32, u32)> = Vec::new();
        let mut updates: Vec<(u32, Vec<u32>)> = Vec::new();
        let mut cursor = self.scan_cursor;
        for i in 0..n_slots {
            if spans.len() == budget {
                break;
            }
            let slot = (self.scan_cursor + i) % n_slots;
            if self.slot_busy[slot] {
                continue;
            }
            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;
            if !Self::ensure_blocks(&mut self.pool, seq, pos0 + span + 1, &mut updates) {
                continue;
            }
            let base = cols_meta.len();
            let tok = *seq.tokens.last().expect("decode has tokens");
            cols_meta.push((pos0 as u32, (slot + 1) as u32, 1, tok));
            for (d, t) in seq.pending_drafts.iter().enumerate() {
                cols_meta.push(((pos0 + 1 + d) as u32, (slot + 1) as u32, 1, *t));
            }
            spans.push(SpanOwner {
                slot,
                base,
                drafts: seq.pending_drafts.clone(),
            });
            seq.drafts_ready = false;
            self.slot_busy[slot] = true;
            cursor = (slot + 1) % n_slots;
        }
        if spans.is_empty() {
            return Ok(None);
        }
        self.scan_cursor = cursor;
        anyhow::ensure!(
            self.local.is_none(),
            "serving-MTP requires a headless coordinator"
        );
        // Re-draft depth for the accepted positions: occupancy-gated. 0 = drain — the worker
        // accepts + seeds but appends no drafts, and the sequence rejoins the plain packer.
        let k_next = self.spec_depth_for(self.live_occupancy());
        self.to_first
            .send(MbMsg {
                group: g,
                spec_on: false,
                spec_verify: true,
                spec_k: k_next,
                wide: false,
                zslots: std::mem::take(&mut self.fresh_slots),
                restores: std::mem::take(&mut self.pending_restores),
                btab: updates,
                cols: cols_meta,
                hidden: Vec::new(),
            })
            .map_err(|_| anyhow!("stage chain died"))?;
        self.spec_rounds += 1;
        self.inflight[g] = Some(Round::Spec(spans, k_verify, k_next));
        Ok(Some(true))
    }

    fn apply_done(&mut self, done: MbDone) -> Vec<(u64, u32, bool)> {
        let round = self.inflight[done.group].take().expect("group in flight");
        match round {
            Round::Plain(owners, k) => self.apply_plain(done, owners, k),
            Round::Spec(spans, k_verify, k_next) => self.apply_spec(done, spans, k_verify, k_next),
        }
    }

    fn apply_plain(
        &mut self,
        done: MbDone,
        owners: Vec<(usize, bool)>,
        k: usize,
    ) -> Vec<(u64, u32, bool)> {
        let group = done.group;
        let ncols = owners.len();
        let mut out = Vec::new();
        let pld_on = self.pld;
        let mut pld_hits = 0u64;
        // Under spec serving the last stage appended k drafts per EMITTING column, in column
        // order (see the worker's spontaneous-draft block).
        let mut next_draft = ncols;
        for (ci, (slot, emits)) in owners.into_iter().enumerate() {
            self.slot_busy[slot] = false;
            if !emits {
                continue;
            }
            let token = done.tokens[ci];
            let drafts = if k > 0 && next_draft + k <= done.tokens.len() {
                let d = done.tokens[next_draft..next_draft + k].to_vec();
                next_draft += k;
                Some(d)
            } else {
                None
            };
            let seq = self.slots[slot].as_mut().expect("owner slot occupied");
            seq.tokens.push(token);
            seq.emitted.push(token);
            seq.state = SeqState::Decode;
            let finished = self.eos.contains(&token) || seq.emitted.len() >= seq.max_tokens;
            out.push((seq.id, token, finished));
            if finished {
                let mut seq = self.slots[slot].take().expect("finished seq");
                for b in seq.blocks.drain(..) {
                    self.pool.unref(b);
                    self.pool.release(b);
                }
            } else if let Some(d) = drafts {
                let seq = self.slots[slot].as_mut().expect("live");
                // PLD replace-form: the sequence's own recent structure beats the MTP head on
                // copy/echo spans; the verify round judges either the same way.
                let look = if pld_on {
                    prompt_lookup_drafts(&seq.tokens, d.len())
                } else {
                    Vec::new()
                };
                if look.is_empty() {
                    seq.pending_drafts = d;
                } else {
                    seq.pending_drafts = look;
                    pld_hits += 1;
                }
                seq.drafts_ready = true;
            }
            let _ = group;
        }
        self.pld_hits += pld_hits;
        out
    }

    fn apply_spec(
        &mut self,
        done: MbDone,
        spans: Vec<SpanOwner>,
        k_verify: usize,
        k_next: usize,
    ) -> Vec<(u64, u32, bool)> {
        let group = done.group;
        let k = k_verify;
        let ncols: usize = spans.len() * (1 + k);
        let mut out = Vec::new();
        let mut pld_hits = 0u64;
        for (si, sp) in spans.into_iter().enumerate() {
            self.slot_busy[sp.slot] = false;
            let mut j = 0usize;
            while j < k && done.tokens[sp.base + j] == sp.drafts[j] {
                j += 1;
            }
            let seq = self.slots[sp.slot].as_mut().expect("span slot occupied");
            let mut finished = false;
            for t in done.tokens[sp.base..=sp.base + j].iter() {
                if finished {
                    break;
                }
                seq.tokens.push(*t);
                seq.emitted.push(*t);
                finished = self.eos.contains(t) || seq.emitted.len() >= seq.max_tokens;
                out.push((seq.id, *t, finished));
            }
            if j < k {
                self.pending_restores.push((
                    (sp.slot + 1) as u32,
                    (sp.base + j) as u32,
                    group as u32,
                ));
            }
            if finished {
                let mut seq = self.slots[sp.slot].take().expect("finished seq");
                for b in seq.blocks.drain(..) {
                    self.pool.unref(b);
                    self.pool.release(b);
                }
            } else {
                // The last stage already seeded + (when k_next > 0) drafted at base+j; drafts
                // ride the response at the NEXT depth (span order, after the verify tokens).
                // k_next == 0 = drain: no drafts appended, the sequence rejoins plain packing.
                if k_next > 0 {
                    let d0 = ncols + si * k_next;
                    let pld_on = self.pld;
                    let seq = self.slots[sp.slot].as_mut().expect("live");
                    let look = if pld_on {
                        prompt_lookup_drafts(&seq.tokens, k_next)
                    } else {
                        Vec::new()
                    };
                    if look.is_empty() {
                        seq.pending_drafts = done.tokens[d0..d0 + k_next].to_vec();
                    } else {
                        seq.pending_drafts = look;
                        pld_hits += 1;
                    }
                    seq.drafts_ready = true;
                }
            }
        }
        self.pld_hits += pld_hits;
        out
    }

    /// Drive until every submitted request finishes; returns each request's emitted tokens.
    /// Keeps `n_stages + 1` micro-batches in flight — the pipeline-filling property.
    pub fn run_to_completion(&mut self) -> Result<HashMap<u64, Vec<u32>>> {
        self.run_streaming(|_, _, _| {})
    }

    /// Like [`Self::run_to_completion`], but invokes `on_token(request_id, token, finished)` for
    /// every token the instant it is produced — the hook for live/streamed output.
    pub fn run_streaming(
        &mut self,
        mut on_token: impl FnMut(u64, u32, bool),
    ) -> Result<HashMap<u64, Vec<u32>>> {
        let mut results: HashMap<u64, Vec<u32>> = HashMap::new();
        loop {
            self.admit();
            let mut launched = false;
            for g in 0..self.n_groups {
                if self.inflight[g].is_none() {
                    launched |= self.launch_group(g)?;
                }
            }
            let pending = self.queue.len() + self.slots.iter().flatten().count();
            let any_inflight = self.inflight.iter().any(Option::is_some);
            if pending == 0 && !any_inflight {
                return Ok(results);
            }
            if !any_inflight {
                // With a KV-aware `admit`, reaching here means the pool cannot carry even ONE queued
                // request to completion — draining will never make room, so say what is actually
                // wrong instead of the old bare "stalled", which sent me hunting a scheduler bug
                // when the answer was capacity.
                anyhow::ensure!(
                    launched,
                    "pipeline stalled with {pending} pending: the KV pool ({} blocks × {PAGE_BLK} \
                     tokens) cannot hold even one of them to completion. Raise --vram-gb / the KV \
                     fraction, or lower --max-tokens.",
                    self.pool.capacity(),
                );
                continue;
            }
            let done = self
                .done_rx
                .recv()
                .map_err(|_| anyhow!("stage chain died"))?;
            for (id, token, finished) in self.apply_done(done) {
                results.entry(id).or_default().push(token);
                on_token(id, token, finished);
            }
        }
    }

    /// End the worker sessions (they keep listening for the next coordinator): dropping the
    /// head of the channel chain drains each thread, whose exit path sends OP_SHUTDOWN. Under
    /// P2P, stages 1..n get their shutdown on the kept control connections FIRST — each clears
    /// its next hop, so the chain (and finally the sink) unwinds stage by stage.
    pub fn shutdown_workers(mut self) {
        for c in &mut self.extra_clients {
            c.shutdown();
        }
        let (dead_tx, _) = channel::<MbMsg>();
        drop(std::mem::replace(&mut self.to_first, dead_tx));
        for t in self._threads.drain(..) {
            let _ = t.join();
        }
    }
}

/// REPLICATION serving: several independent [`PipelineServe`] replicas + a request dispatcher —
/// the throughput axis of the fleet (see `replica.rs` for how the worker pool is *allocated* to
/// replicas; this *runs* them). Each replica is a full pipeline that independently holds the whole
/// model; requests spread across them. The coordinator is a pure router here (each replica built
/// headless, `split = 0`), so it does no local GPU work and the replicas run truly concurrently.
pub struct ReplicatedServe {
    replicas: Vec<PipelineServe>,
    stage_counts: Vec<usize>,
}

impl ReplicatedServe {
    /// Build one [`PipelineServe`] per replica from a pre-accepted [`Fleet`], partitioned per the
    /// allocator's `plan` (its worker indices point into `fleet.clients`). Each replica then
    /// auto-splits ITS workers across the model's layers via the usual negotiation. Spare workers
    /// (in `plan.spares`) are dropped — their control sockets close, so they re-join and idle.
    pub fn connect(
        dir: &Path,
        plan: &FleetPlan,
        fleet: Fleet,
        mb_k: usize,
        prefill_k: usize,
        eos: Vec<u32>,
    ) -> Result<Self> {
        let mut clients: Vec<Option<ShardClient>> = fleet.clients.into_iter().map(Some).collect();
        let addrs = fleet.data_addrs;
        let mut replicas = Vec::with_capacity(plan.replicas.len());
        let mut stage_counts = Vec::with_capacity(plan.replicas.len());
        for (r, rep) in plan.replicas.iter().enumerate() {
            let sub_clients: Vec<ShardClient> = rep
                .workers
                .iter()
                .map(|&i| {
                    clients
                        .get_mut(i)
                        .and_then(Option::take)
                        .ok_or_else(|| anyhow!("worker {i} assigned to two replicas"))
                })
                .collect::<Result<_>>()?;
            let sub_addrs: Vec<String> = rep.workers.iter().map(|&i| addrs[i].clone()).collect();
            eprintln!(
                "replica {}/{}: {} stage(s), {:.1} GB",
                r + 1,
                plan.replicas.len(),
                sub_clients.len(),
                rep.vram as f64 / 1e9
            );
            stage_counts.push(sub_clients.len());
            let subfleet = Fleet {
                clients: sub_clients,
                data_addrs: sub_addrs,
            };
            replicas.push(PipelineServe::connect_fleet(
                dir,
                0,
                subfleet,
                mb_k,
                prefill_k,
                eos.clone(),
            )?);
        }
        Ok(Self {
            replicas,
            stage_counts,
        })
    }

    /// Number of replicas built.
    pub fn n_replicas(&self) -> usize {
        self.replicas.len()
    }

    /// Stage count of each replica (for the topology print).
    pub fn stage_counts(&self) -> &[usize] {
        &self.stage_counts
    }

    /// Serve many prompts, spread round-robin across the replicas which run CONCURRENTLY (one OS
    /// thread drives each replica; the coordinator is headless, so there's no local-GPU
    /// contention). Returns per-input completions IN INPUT ORDER. Wall-clock ≈ (prompts / replicas)
    /// × per-prompt time, so throughput scales ~with the replica count. Consumes `self` and shuts
    /// the workers down at the end (one-shot batch serving).
    pub fn run_prompts(self, prompts: Vec<Vec<u32>>, max_tokens: usize) -> Result<Vec<Vec<u32>>> {
        let r = self.replicas.len().max(1);
        let n = prompts.len();
        let mut buckets: Vec<Vec<(usize, Vec<u32>)>> = (0..r).map(|_| Vec::new()).collect();
        for (i, p) in prompts.into_iter().enumerate() {
            buckets[i % r].push((i, p));
        }
        let handles: Vec<_> = self
            .replicas
            .into_iter()
            .zip(buckets)
            .map(|(mut serve, bucket)| {
                std::thread::spawn(move || -> Result<Vec<(usize, Vec<u32>)>> {
                    let mut rids: Vec<(u64, usize)> = Vec::with_capacity(bucket.len());
                    for (orig, p) in bucket {
                        rids.push((serve.submit(p, max_tokens)?, orig));
                    }
                    let done = serve.run_to_completion()?;
                    let out = rids
                        .into_iter()
                        .map(|(rid, orig)| (orig, done.get(&rid).cloned().unwrap_or_default()))
                        .collect::<Vec<_>>();
                    serve.shutdown_workers();
                    Ok(out)
                })
            })
            .collect();
        let mut out: Vec<Vec<u32>> = vec![Vec::new(); n];
        for h in handles {
            let part = h
                .join()
                .map_err(|_| anyhow!("replica serving thread panicked"))??;
            for (orig, toks) in part {
                out[orig] = toks;
            }
        }
        Ok(out)
    }
}