fix(neuron): TP unload with Arc-shared pool mutex (#98)

cfg(cuda)-only: the unload path moved the pool out of the (previously
owned) mutex; with TpLoadedModel.pool behind Arc for the engine's
owned guard, recover sole ownership via Arc::try_unwrap —
Arc::try_unwrap(tp) succeeding already implies the engine is idle, so
this is the expected case; the fallback (engine guard mid-release
race) unloads through the lock and lets the pool drop with its last
reference.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

crates/neuron/src/harness/candle.rs

@@ -480,7 +480,7 @@ pub struct TpLoadedModel {
     /// so this Mutex no longer covers the leader's KV cache; it just
     /// serialises subprocess RPC traffic on the pool's
     /// `Vec<Worker>` channels.
-    pub pool: tokio::sync::Mutex<super::tp::WorkerPool>,
+    pub pool: Arc<tokio::sync::Mutex<super::tp::WorkerPool>>,
     /// Bounded admission scheduler (#53), mirroring the single-GPU path.
     /// Gated before the pool lock so an overloaded TP model returns fast
     /// backpressure instead of an unbounded, untimed wait.
@@ -545,6 +545,13 @@ pub struct TpLoadedModel {
     /// Mint for pool-wide snapshot ids. Plain counter; uniqueness only
     /// needs to hold per model lifetime (snapshots die with the model).
     pub next_snapshot_id: std::sync::atomic::AtomicU64,
+    /// Lockstep batched decode engine (#98) — the TP mirror of
+    /// `LoadedModel::engine`. Set once at load (after the `Arc` is
+    /// built, so the engine can hold a `Weak` back to this model);
+    /// unset when `max_in_flight` is 1 or batching is killed. Text
+    /// chat streams route through it; the engine holds the pool mutex
+    /// while it has active slots.
+    pub engine: std::sync::OnceLock<super::engine::EngineHandle>,
     /// Cached tightest free VRAM (MiB) for the control plane (#53) — see
     /// [`LoadedModel::last_free_mb`]. Read by `derived_limit` so `GET /models`
     /// never fans a VRAM query out to the (inference-saturated) TP workers.
@@ -1657,7 +1664,7 @@ pub(crate) async fn chunked_prefill_via_worker(
 /// `start_offset` skips a restored cached prefix, as in
 /// [`chunked_prefill_local`].
 #[cfg(feature = "cuda")]
-async fn chunked_prefill_tp(
+pub(crate) async fn chunked_prefill_tp(
     pool: &mut super::tp::WorkerPool,
     model_id: &str,
     leader_handle: super::device_worker::TpHandle,
@@ -3468,16 +3475,18 @@ impl Harness for CandleHarness {
                 Some(super::engine::EngineHandle::spawn(
                     super::engine::EngineConfig {
                         model_id: spec.model_id.clone(),
-                        worker: Arc::clone(w),
-                        handle: h,
                         tokenizer: tokenizer.clone(),
-                        prefix_cache: prefix_cache.clone(),
-                        prefill_rate: Arc::clone(&prefill_rate),
                         reasoning_tokens: reasoning_tokens.clone(),
                         tool_call_tokens: tool_call_tokens.clone(),
-                        poisoned: Arc::clone(&poisoned),
-                        inference_lock: Arc::clone(&inference_lock),
                         max_slots: self.admission_cfg.max_in_flight,
+                        backend: super::engine::BackendConfig::Single {
+                            worker: Arc::clone(w),
+                            handle: h,
+                            prefix_cache: prefix_cache.clone(),
+                            prefill_rate: Arc::clone(&prefill_rate),
+                            poisoned: Arc::clone(&poisoned),
+                            inference_lock: Arc::clone(&inference_lock),
+                        },
                     },
                 ))
             }
@@ -3588,12 +3597,35 @@ impl Harness for CandleHarness {
                         "TP unload: DropTp RPC failed (leader model may leak in worker slab)"
                     );
                 }
-                let mut pool = tp.pool.into_inner();
-                if let Err(e) = pool.unload_model(model_id).await {
-                    tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
-                }
-                if let Err(e) = pool.shutdown().await {
-                    tracing::warn!(model = %model_id, error = %e, "TP pool shutdown failed");
+                // The pool mutex is Arc-shared with the batch engine's
+                // active-phase guard (#98). `Arc::try_unwrap(tp)`
+                // succeeding above means the engine is idle (it holds
+                // `Arc<TpLoadedModel>` whenever it holds the pool
+                // guard), so sole ownership is the expected case; the
+                // fallback covers the narrow race where the engine's
+                // guard is mid-release.
+                match Arc::try_unwrap(tp.pool) {
+                    Ok(pool_mutex) => {
+                        let mut pool = pool_mutex.into_inner();
+                        if let Err(e) = pool.unload_model(model_id).await {
+                            tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
+                        }
+                        if let Err(e) = pool.shutdown().await {
+                            tracing::warn!(model = %model_id, error = %e, "TP pool shutdown failed");
+                        }
+                    }
+                    Err(pool_arc) => {
+                        tracing::warn!(
+                            model = %model_id,
+                            "TP unload: pool mutex still referenced (engine guard \
+                             mid-release); unloading without explicit pool shutdown — \
+                             worker children reap when the last reference drops"
+                        );
+                        let mut pool = pool_arc.lock().await;
+                        if let Err(e) = pool.unload_model(model_id).await {
+                            tracing::warn!(model = %model_id, error = %e, "TP unload RPC failed");
+                        }
+                    }
                 }
             }
         }
@@ -3748,7 +3780,7 @@ impl CandleHarness {
             model_id: spec.model_id.clone(),
             tokenizer,
             devices: devices.clone(),
-            pool: TMutex::new(pool),
+            pool: StdArc::new(TMutex::new(pool)),
             admission: super::admission::AdmissionController::new(&self.admission_cfg),
             leader_handle,
             leader_device: leader_device.clone(),
@@ -3778,7 +3810,33 @@ impl CandleHarness {
             derived_input_cap: AtomicUsize::new(0),
             last_free_mb: AtomicU64::new(0),
             next_snapshot_id: std::sync::atomic::AtomicU64::new(1),
+            engine: std::sync::OnceLock::new(),
         });
+        // Batched decode engine (#98): spawned when the operator raised
+        // max_in_flight above 1 on a snapshot-capable TP model. The
+        // engine holds only a Weak — an idle engine never keeps an
+        // unloaded model alive.
+        if tp_loaded.prefix_cache.is_some()
+            && self.admission_cfg.max_in_flight > 1
+            && super::engine::batching_enabled()
+        {
+            tracing::info!(
+                model = %spec.model_id,
+                max_slots = self.admission_cfg.max_in_flight,
+                "batched decode engine enabled for TP model (#98)"
+            );
+            let handle = super::engine::EngineHandle::spawn(super::engine::EngineConfig {
+                model_id: spec.model_id.clone(),
+                tokenizer: tp_loaded.tokenizer.clone(),
+                reasoning_tokens: tp_loaded.reasoning_tokens.clone(),
+                tool_call_tokens: tp_loaded.tool_call_tokens.clone(),
+                max_slots: self.admission_cfg.max_in_flight,
+                backend: super::engine::BackendConfig::Tp {
+                    tp: StdArc::downgrade(&tp_loaded),
+                },
+            });
+            let _ = tp_loaded.engine.set(handle);
+        }
         if tp_loaded.prefix_cache.is_some() {
             tracing::info!(
                 model = %spec.model_id,
@@ -4098,6 +4156,38 @@ impl CandleHarness {
             .map_err(InferenceError::from)?;
 
         let tool_schemas = build_tool_schemas(&request);
+        // Batched decode engine (#98): text streams multiplex through
+        // the per-model engine instead of serializing on the pool
+        // mutex per request. Vision requests keep the direct path and
+        // serialize against the engine via the pool lock it holds
+        // while active.
+        if vision_route.is_none()
+            && let Some(engine) = tp.engine.get()
+        {
+            engine
+                .submit(super::engine::EngineRequest {
+                    prompt_tokens,
+                    max_new,
+                    temperature,
+                    top_p,
+                    seed,
+                    eos_id,
+                    tool_schemas,
+                    tx,
+                    admit,
+                    span,
+                })
+                .await
+                .map_err(InferenceError::Other)?;
+            let reasoning_markers = tp.reasoning_tokens.clone();
+            return Ok(InferenceStream {
+                events: event_rx,
+                id: projector_id,
+                created,
+                model_id: projector_model_id,
+                reasoning_markers,
+            });
+        }
         let tp_for_task = Arc::clone(&tp);
         tokio::spawn(
             async move {
@@ -5911,7 +6001,7 @@ fn restore_or_clear_local(
 /// (some restored, some not) — the clear fallback resets every rank,
 /// restoring consistency.
 #[cfg(feature = "cuda")]
-async fn restore_or_clear_tp(
+pub(crate) async fn restore_or_clear_tp(
     pool: &mut super::tp::WorkerPool,
     tp: &TpLoadedModel,
     prompt_tokens: &[u32],
@@ -5968,7 +6058,7 @@ async fn restore_or_clear_tp(
 /// On any rank failing, drops the id everywhere (idempotent) so no
 /// rank leaks a half-stored snapshot.
 #[cfg(feature = "cuda")]
-async fn store_prefix_snapshot_tp(
+pub(crate) async fn store_prefix_snapshot_tp(
     pool: &mut super::tp::WorkerPool,
     tp: &TpLoadedModel,
     prompt_tokens: Vec<u32>,

crates/neuron/src/harness/device_worker/dispatch.rs

@@ -476,6 +476,50 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
                 let _ = reply.send(());
             }
             #[cfg(feature = "cuda")]
+            Job::TpForwardLogitsBatch {
+                handle,
+                tokens,
+                prefix_lens,
+                padded_len,
+                step,
+                reply,
+            } => {
+                let result = tp_forward_logits_batch(
+                    &mut state,
+                    handle,
+                    &tokens,
+                    &prefix_lens,
+                    padded_len,
+                    step,
+                );
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpAssembleKvBatch {
+                handle,
+                seqs,
+                reply,
+            } => {
+                let result = tp_assemble_kv_batch(&mut state, handle, &seqs);
+                if result.is_ok() {
+                    trim_device_pool(&state);
+                }
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
+            Job::TpExtractKvRows {
+                handle,
+                rows,
+                padded_len,
+                steps,
+                snapshot_ids,
+                reply,
+            } => {
+                let result =
+                    tp_extract_kv_rows(&mut state, handle, &rows, padded_len, steps, &snapshot_ids);
+                let _ = reply.send(result);
+            }
+            #[cfg(feature = "cuda")]
             Job::TpForwardLogits {
                 handle,
                 tokens,
@@ -991,6 +1035,114 @@ fn tp_forward_logits(
     Ok(values)
 }
 
+/// TP-equivalent of [`forward_logits_batch`] on the leader's shard
+/// (#98). The caller has already fanned the matching
+/// `GenerateStepBatch` out to the subprocess ranks.
+#[cfg(feature = "cuda")]
+fn tp_forward_logits_batch(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    tokens: &[u32],
+    prefix_lens: &[usize],
+    padded_len: usize,
+    step: usize,
+) -> anyhow::Result<Vec<Vec<f32>>> {
+    use candle_core::{DType, Tensor};
+
+    let b = tokens.len();
+    anyhow::ensure!(b > 0, "TpForwardLogitsBatch: empty batch");
+    anyhow::ensure!(
+        prefix_lens.len() == b,
+        "TpForwardLogitsBatch: {} prefix_lens for batch of {b}",
+        prefix_lens.len()
+    );
+
+    let input = Tensor::from_vec(tokens.to_vec(), (b, 1), &state.device)?;
+    let model = state
+        .tp_models
+        .get_mut(&handle)
+        .ok_or_else(|| anyhow::anyhow!("TpForwardLogitsBatch: no model for handle {}", handle.0))?;
+
+    let positions: Vec<usize> = prefix_lens.iter().map(|&len| len + step).collect();
+    let total_len = padded_len + step + 1;
+    let mask = model.batch_decode_mask(prefix_lens, padded_len, total_len)?;
+    let logits = model.forward_batch_decode(&input, &positions, mask.as_ref())?;
+
+    let logits = logits.to_dtype(DType::F32)?;
+    let (rows, _, vocab) = logits.dims3()?;
+    anyhow::ensure!(
+        rows == b,
+        "TpForwardLogitsBatch: model returned {rows} logits rows for batch of {b}"
+    );
+    let flat: Vec<f32> = logits.flatten_all()?.to_vec1()?;
+    Ok(flat.chunks(vocab).map(<[f32]>::to_vec).collect())
+}
+
+/// TP-equivalent of [`assemble_kv_batch`] against the TP slab (#98),
+/// keyed by pool-minted snapshot ids.
+#[cfg(feature = "cuda")]
+fn tp_assemble_kv_batch(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    seqs: &[(u64, usize)],
+) -> anyhow::Result<usize> {
+    let DeviceWorkerState {
+        tp_models,
+        tp_kv_snapshots,
+        ..
+    } = state;
+    let mut pairs = Vec::with_capacity(seqs.len());
+    for (id, len) in seqs {
+        let snap = tp_kv_snapshots.get(&(handle, *id)).ok_or_else(|| {
+            anyhow::anyhow!(
+                "TpAssembleKvBatch: no snapshot {id} for handle {}",
+                handle.0
+            )
+        })?;
+        pairs.push((snap, *len));
+    }
+    let batch = crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs)?;
+    let model = tp_models
+        .get_mut(&handle)
+        .ok_or_else(|| anyhow::anyhow!("TpAssembleKvBatch: no model for handle {}", handle.0))?;
+    model.restore_kv_cache(&batch.snapshot)?;
+    Ok(batch.padded_len)
+}
+
+/// TP-equivalent of [`extract_kv_rows`] against the TP slab (#98),
+/// storing each extracted row under the caller's pre-minted pool id.
+/// Returns total snapshot bytes (budget accounting).
+#[cfg(feature = "cuda")]
+fn tp_extract_kv_rows(
+    state: &mut DeviceWorkerState,
+    handle: TpHandle,
+    rows: &[(usize, usize)],
+    padded_len: usize,
+    steps: usize,
+    snapshot_ids: &[u64],
+) -> anyhow::Result<u64> {
+    anyhow::ensure!(
+        snapshot_ids.len() == rows.len(),
+        "TpExtractKvRows: {} rows vs {} snapshot_ids",
+        rows.len(),
+        snapshot_ids.len()
+    );
+    let live = state
+        .tp_models
+        .get(&handle)
+        .ok_or_else(|| anyhow::anyhow!("TpExtractKvRows: no model for handle {}", handle.0))?
+        .snapshot_kv_cache()?;
+    let mut total = 0u64;
+    for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
+        let snap = crate::harness::arch::qwen3_5::snapshot::extract_row(
+            &live, row, prefix_len, padded_len, steps,
+        )?;
+        total += snap.size_bytes();
+        state.tp_kv_snapshots.insert((handle, id), snap);
+    }
+    Ok(total)
+}
+
 /// Image-bearing leader forward (rank 0). Preprocesses each source
 /// `image_data_uris` entry through the same deterministic
 /// `preprocess_data_uri` every rank runs, uploads to the leader's
@@ -1400,6 +1552,18 @@ fn drain_poisoned(job: Job, device_index: u32) {
             let _ = reply.send(Err(err()));
         }
         #[cfg(feature = "cuda")]
+        Job::TpForwardLogitsBatch { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpAssembleKvBatch { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
+        Job::TpExtractKvRows { reply, .. } => {
+            let _ = reply.send(Err(err()));
+        }
+        #[cfg(feature = "cuda")]
         Job::TpDropKvSnapshot { reply, .. } => {
             // Bookkeeping-only — unit reply so eviction never wedges
             // on a poisoned worker (same shape as DropKvSnapshot).

crates/neuron/src/harness/device_worker/jobs.rs

@@ -348,6 +348,41 @@ pub enum Job {
         offset: usize,
         reply: oneshot::Sender<Result<Vec<f32>>>,
     },
+    /// Leader half of a TP batched decode step (#98) — mirrors the
+    /// single-GPU `ForwardLogitsBatch` against the TP slab. The caller
+    /// (`WorkerPool::generate_step_batch`) fans out the matching
+    /// `GenerateStepBatch` RPC to subprocess ranks first so the
+    /// row-parallel collectives pair up.
+    #[cfg(feature = "cuda")]
+    TpForwardLogitsBatch {
+        handle: TpHandle,
+        tokens: Vec<u32>,
+        prefix_lens: Vec<usize>,
+        padded_len: usize,
+        step: usize,
+        reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
+    },
+    /// Leader half of a TP batch assembly (#98) — mirrors
+    /// `AssembleKvBatch` against the TP slab, with **pool-minted**
+    /// snapshot ids (the same ids every subprocess rank stores under).
+    #[cfg(feature = "cuda")]
+    TpAssembleKvBatch {
+        handle: TpHandle,
+        seqs: Vec<(u64, usize)>,
+        reply: oneshot::Sender<Result<usize>>,
+    },
+    /// Leader half of a TP row extraction (#98) — mirrors
+    /// `ExtractKvRows` against the TP slab, storing each extracted row
+    /// under the **pre-minted** id in `snapshot_ids` (one per row).
+    #[cfg(feature = "cuda")]
+    TpExtractKvRows {
+        handle: TpHandle,
+        rows: Vec<(usize, usize)>,
+        padded_len: usize,
+        steps: usize,
+        snapshot_ids: Vec<u64>,
+        reply: oneshot::Sender<Result<u64>>,
+    },
     /// Image-bearing leader (rank 0) forward for the single-shot vision
     /// prefill. The handler preprocesses each `image_data_uris` entry
     /// (the same deterministic path every rank runs), encodes through

crates/neuron/src/harness/device_worker/mod.rs

@@ -874,6 +874,112 @@ impl DeviceWorkerHandle {
         }
     }
 
+    /// Leader half of a TP batched decode step (#98) — per-row
+    /// `[vocab]` logits back for per-slot sampling. The caller fans
+    /// the matching `GenerateStepBatch` out to the subprocess ranks.
+    #[cfg(feature = "cuda")]
+    pub async fn tp_forward_logits_batch(
+        &self,
+        handle: TpHandle,
+        tokens: Vec<u32>,
+        prefix_lens: Vec<usize>,
+        padded_len: usize,
+        step: usize,
+    ) -> Result<Vec<Vec<f32>>, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::TpForwardLogitsBatch {
+                handle,
+                tokens,
+                prefix_lens,
+                padded_len,
+                step,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// Leader half of a TP batch assembly (#98). Returns the padded
+    /// uniform KV length. Snapshot ids are pool-minted by the caller.
+    #[cfg(feature = "cuda")]
+    pub async fn tp_assemble_kv_batch(
+        &self,
+        handle: TpHandle,
+        seqs: Vec<(u64, usize)>,
+    ) -> Result<usize, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::TpAssembleKvBatch {
+                handle,
+                seqs,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
+    /// Leader half of a TP row extraction (#98). Stores each extracted
+    /// row under the caller's pre-minted pool id; returns total bytes.
+    #[cfg(feature = "cuda")]
+    pub async fn tp_extract_kv_rows(
+        &self,
+        handle: TpHandle,
+        rows: Vec<(usize, usize)>,
+        padded_len: usize,
+        steps: usize,
+        snapshot_ids: Vec<u64>,
+    ) -> Result<u64, WorkerError> {
+        if self.poisoned.load(Ordering::Acquire) {
+            return Err(WorkerError::Poisoned {
+                device_index: self.device_index,
+            });
+        }
+        let (reply_tx, reply_rx) = oneshot::channel();
+        self.tx
+            .send(Job::TpExtractKvRows {
+                handle,
+                rows,
+                padded_len,
+                steps,
+                snapshot_ids,
+                reply: reply_tx,
+            })
+            .map_err(|_| WorkerError::Gone {
+                device_index: self.device_index,
+            })?;
+        match reply_rx.await {
+            Ok(result) => result.map_err(WorkerError::from),
+            Err(_) => Err(WorkerError::Gone {
+                device_index: self.device_index,
+            }),
+        }
+    }
+
     /// Image-bearing TP leader forward (single-shot vision prefill).
     /// Routes `Job::TpForwardLogitsWithImages` onto the worker thread;
     /// the handler preprocesses + encodes + splices + forwards and

crates/neuron/src/harness/engine.rs

@@ -51,7 +51,7 @@ use super::candle::{
     stable_snapshot_cut, store_prefix_snapshot_via_worker,
 };
 use super::context_limit::PrefillRateEma;
-use super::device_worker::{ArchHandle, DeviceWorkerHandle};
+use super::device_worker::{ArchHandle, DeviceWorkerHandle, jobs};
 use crate::wire::event::{
     FinishReason, FinishTiming, InferenceEvent, ReasoningTokenPair, ToolCallTokenPair,
 };
@@ -69,28 +69,327 @@ pub fn batching_enabled() -> bool {
 }
 
 /// Everything the engine needs that is per-model (not per-request).
-/// Deliberately does NOT hold `Arc<LoadedModel>` — the engine task
-/// must not keep the model alive (the task exits when the model drops
-/// its `EngineHandle` and the channel closes).
 pub struct EngineConfig {
     pub model_id: String,
-    pub worker: Arc<DeviceWorkerHandle>,
-    pub handle: ArchHandle,
     pub tokenizer: Tokenizer,
-    pub prefix_cache: Option<Arc<ModelPrefixCache>>,
-    pub prefill_rate: Arc<PrefillRateEma>,
     pub reasoning_tokens: Option<ReasoningTokenPair>,
     pub tool_call_tokens: Option<ToolCallTokenPair>,
-    /// Shared with `LoadedModel.poisoned` so a device fault inside the
-    /// engine fast-rejects subsequent requests at the harness boundary.
-    pub poisoned: Arc<AtomicBool>,
-    /// Shared with `LoadedModel.inference_lock`. Held for the whole
-    /// active phase (first join → last slot finished) so the
-    /// non-engine forward paths (vision, non-streaming chat) can never
-    /// clobber the live batched cache state mid-decode. Released while
-    /// idle.
-    pub inference_lock: Arc<tokio::sync::Mutex<()>>,
     pub max_slots: usize,
+    pub backend: BackendConfig,
+}
+
+/// The two serving substrates an engine can drive. Deliberately does
+/// NOT keep the owning model alive while idle — the engine task must
+/// exit when the model unloads (the `EngineHandle` sender drops and
+/// the channel closes; the TP `Weak` stops upgrading).
+pub enum BackendConfig {
+    /// Single-GPU device-worker path (`LoadedModel`). Holds Arc'd
+    /// clones of the model's shared pieces rather than the model.
+    Single {
+        worker: Arc<DeviceWorkerHandle>,
+        handle: ArchHandle,
+        prefix_cache: Option<Arc<ModelPrefixCache>>,
+        prefill_rate: Arc<PrefillRateEma>,
+        /// Shared with `LoadedModel.poisoned` so a device fault inside
+        /// the engine fast-rejects subsequent requests at the harness
+        /// boundary.
+        poisoned: Arc<AtomicBool>,
+        /// Shared with `LoadedModel.inference_lock`. Held for the whole
+        /// active phase (first join → last slot finished) so the
+        /// non-engine forward paths (vision, non-streaming chat) can
+        /// never clobber the live batched cache state mid-decode.
+        /// Released while idle.
+        inference_lock: Arc<tokio::sync::Mutex<()>>,
+    },
+    /// Tensor-parallel path (`TpLoadedModel`). The `Weak` upgrades for
+    /// the duration of an active phase only: while the engine has
+    /// slots it holds the `Arc` (an unload then completes when the
+    /// batch drains), while idle it holds nothing and unload is
+    /// immediate.
+    #[cfg(feature = "cuda")]
+    Tp {
+        tp: std::sync::Weak<super::candle::TpLoadedModel>,
+    },
+}
+
+/// Held while the engine has active slots: the exclusive right to the
+/// model's forward path plus (for TP) the upgraded model handle.
+enum ActiveSession {
+    Single {
+        _guard: tokio::sync::OwnedMutexGuard<()>,
+    },
+    #[cfg(feature = "cuda")]
+    Tp {
+        tp: Arc<super::candle::TpLoadedModel>,
+        pool: tokio::sync::OwnedMutexGuard<super::tp::WorkerPool>,
+    },
+}
+
+impl BackendConfig {
+    /// Enter the active phase: take the model's serialization lock
+    /// (single-GPU `inference_lock` / TP pool mutex). `None` when the
+    /// model has been unloaded (TP `Weak` no longer upgrades).
+    async fn acquire(&self) -> Option<ActiveSession> {
+        match self {
+            BackendConfig::Single { inference_lock, .. } => Some(ActiveSession::Single {
+                _guard: Arc::clone(inference_lock).lock_owned().await,
+            }),
+            #[cfg(feature = "cuda")]
+            BackendConfig::Tp { tp } => {
+                let tp = tp.upgrade()?;
+                let pool = Arc::clone(&tp.pool).lock_owned().await;
+                Some(ActiveSession::Tp { tp, pool })
+            }
+        }
+    }
+
+    /// Record a device fault on the owning model's poison flag.
+    fn mark_poisoned(&self) {
+        match self {
+            BackendConfig::Single { poisoned, .. } => poisoned.store(true, Ordering::Release),
+            #[cfg(feature = "cuda")]
+            BackendConfig::Tp { tp } => {
+                if let Some(tp) = tp.upgrade() {
+                    tp.poisoned.store(true, Ordering::Release);
+                }
+            }
+        }
+    }
+}
+
+/// One lockstep decode step against the live batched state.
+async fn op_step(
+    cfg: &EngineConfig,
+    session: &mut ActiveSession,
+    tokens: Vec<u32>,
+    prefix_lens: Vec<usize>,
+    padded_len: usize,
+    step: usize,
+) -> Result<Vec<Vec<f32>>> {
+    match (&cfg.backend, session) {
+        (BackendConfig::Single { worker, handle, .. }, ActiveSession::Single { .. }) => worker
+            .forward_logits_batch(*handle, tokens, prefix_lens, padded_len, step)
+            .await
+            .map_err(|e| anyhow::anyhow!("batched decode step {step}: {e}")),
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => pool
+            .generate_step_batch(
+                &tp.model_id,
+                tp.leader_handle,
+                tokens,
+                prefix_lens,
+                padded_len,
+                step,
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("TP batched decode step {step}: {e}")),
+        #[cfg(feature = "cuda")]
+        _ => anyhow::bail!("engine backend/session mismatch"),
+    }
+}
+
+/// Assemble per-sequence snapshots (by id) into the live batched
+/// state; returns the padded uniform KV length.
+async fn op_assemble(
+    cfg: &EngineConfig,
+    session: &mut ActiveSession,
+    seqs: Vec<(u64, usize)>,
+) -> Result<usize> {
+    match (&cfg.backend, session) {
+        (BackendConfig::Single { worker, handle, .. }, ActiveSession::Single { .. }) => {
+            let seqs = seqs
+                .into_iter()
+                .map(|(id, len)| (jobs::KvSnapshotId(id), len))
+                .collect();
+            worker
+                .assemble_kv_batch(*handle, seqs)
+                .await
+                .map_err(|e| anyhow::anyhow!("assemble_kv_batch: {e}"))
+        }
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => pool
+            .assemble_kv_batch(&tp.model_id, tp.leader_handle, seqs)
+            .await
+            .map_err(|e| anyhow::anyhow!("TP assemble_kv_batch: {e}")),
+        #[cfg(feature = "cuda")]
+        _ => anyhow::bail!("engine backend/session mismatch"),
+    }
+}
+
+/// Extract live batch rows into per-sequence snapshots; returns their
+/// ids (backend-scoped).
+async fn op_extract(
+    cfg: &EngineConfig,
+    session: &mut ActiveSession,
+    rows: Vec<(usize, usize)>,
+    padded_len: usize,
+    steps: usize,
+) -> Result<Vec<u64>> {
+    match (&cfg.backend, session) {
+        (BackendConfig::Single { worker, handle, .. }, ActiveSession::Single { .. }) => Ok(worker
+            .extract_kv_rows(*handle, rows, padded_len, steps)
+            .await
+            .map_err(|e| anyhow::anyhow!("extract_kv_rows: {e}"))?
+            .into_iter()
+            .map(|(id, _bytes)| id.0)
+            .collect()),
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
+            let ids: Vec<u64> = (0..rows.len())
+                .map(|_| tp.next_snapshot_id.fetch_add(1, Ordering::Relaxed))
+                .collect();
+            pool.extract_kv_rows(
+                &tp.model_id,
+                tp.leader_handle,
+                rows,
+                padded_len,
+                steps,
+                ids.clone(),
+            )
+            .await
+            .map_err(|e| anyhow::anyhow!("TP extract_kv_rows: {e}"))?;
+            Ok(ids)
+        }
+        #[cfg(feature = "cuda")]
+        _ => anyhow::bail!("engine backend/session mismatch"),
+    }
+}
+
+/// Snapshot the freshly prefilled single-sequence state for assembly.
+async fn op_snapshot_seq(cfg: &EngineConfig, session: &mut ActiveSession) -> Result<u64> {
+    match (&cfg.backend, session) {
+        (BackendConfig::Single { worker, handle, .. }, ActiveSession::Single { .. }) => worker
+            .snapshot_kv(*handle)
+            .await
+            .map(|(id, _bytes)| id.0)
+            .map_err(|e| anyhow::anyhow!("snapshot after prefill: {e}")),
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
+            let id = tp.next_snapshot_id.fetch_add(1, Ordering::Relaxed);
+            pool.snapshot_kv_cache(&tp.model_id, tp.leader_handle, id)
+                .await
+                .map_err(|e| anyhow::anyhow!("TP snapshot after prefill: {e}"))?;
+            Ok(id)
+        }
+        #[cfg(feature = "cuda")]
+        _ => anyhow::bail!("engine backend/session mismatch"),
+    }
+}
+
+/// Drop a temporary per-sequence snapshot (best-effort).
+async fn op_drop_snap(cfg: &EngineConfig, session: &mut ActiveSession, id: u64) {
+    match (&cfg.backend, session) {
+        (BackendConfig::Single { worker, handle, .. }, ActiveSession::Single { .. }) => {
+            let _ = worker
+                .drop_kv_snapshot(*handle, jobs::KvSnapshotId(id))
+                .await;
+        }
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
+            let _ = pool
+                .drop_kv_snapshot(&tp.model_id, tp.leader_handle, id)
+                .await;
+        }
+        #[cfg(feature = "cuda")]
+        _ => {}
+    }
+}
+
+/// Prefill one joining request at B=1 through the backend's existing
+/// chunked-prefill + prefix-cache paths. Returns the prefill logits.
+async fn op_prefill(
+    cfg: &EngineConfig,
+    session: &mut ActiveSession,
+    prompt_tokens: &[u32],
+) -> Result<Vec<f32>> {
+    let prompt_len = prompt_tokens.len();
+    let prefill_start = std::time::Instant::now();
+    let logits = match (&cfg.backend, session) {
+        (
+            BackendConfig::Single {
+                worker,
+                handle,
+                prefix_cache,
+                prefill_rate,
+                ..
+            },
+            ActiveSession::Single { .. },
+        ) => {
+            let prefix_cache = prefix_cache.as_deref();
+            let reused =
+                restore_or_clear_via_worker(worker, *handle, prefix_cache, prompt_tokens).await?;
+            let cut = if prefix_cache.is_some() {
+                stable_snapshot_cut(prompt_tokens, cfg.tokenizer.token_to_id("<|im_start|>"))
+                    .filter(|&c| c > reused)
+            } else {
+                None
+            };
+            let logits = match cut {
+                Some(c) => {
+                    chunked_prefill_via_worker(worker, *handle, &prompt_tokens[..c], reused)
+                        .await?;
+                    store_prefix_snapshot_via_worker(
+                        worker,
+                        *handle,
+                        prefix_cache,
+                        prompt_tokens[..c].to_vec(),
+                    )
+                    .await;
+                    chunked_prefill_via_worker(worker, *handle, prompt_tokens, c).await?
+                }
+                None => chunked_prefill_via_worker(worker, *handle, prompt_tokens, reused).await?,
+            };
+            prefill_rate.record(prompt_len, prefill_start.elapsed());
+            logits
+        }
+        #[cfg(feature = "cuda")]
+        (BackendConfig::Tp { .. }, ActiveSession::Tp { tp, pool }) => {
+            let reused = super::candle::restore_or_clear_tp(pool, tp, prompt_tokens).await?;
+            let cut = if tp.prefix_cache.is_some() {
+                stable_snapshot_cut(prompt_tokens, cfg.tokenizer.token_to_id("<|im_start|>"))
+                    .filter(|&c| c > reused)
+            } else {
+                None
+            };
+            let logits = match cut {
+                Some(c) => {
+                    super::candle::chunked_prefill_tp(
+                        pool,
+                        &tp.model_id,
+                        tp.leader_handle,
+                        &prompt_tokens[..c],
+                        reused,
+                    )
+                    .await?;
+                    super::candle::store_prefix_snapshot_tp(pool, tp, prompt_tokens[..c].to_vec())
+                        .await;
+                    super::candle::chunked_prefill_tp(
+                        pool,
+                        &tp.model_id,
+                        tp.leader_handle,
+                        prompt_tokens,
+                        c,
+                    )
+                    .await?
+                }
+                None => {
+                    super::candle::chunked_prefill_tp(
+                        pool,
+                        &tp.model_id,
+                        tp.leader_handle,
+                        prompt_tokens,
+                        reused,
+                    )
+                    .await?
+                }
+            };
+            tp.prefill_rate.record(prompt_len, prefill_start.elapsed());
+            logits
+        }
+        #[cfg(feature = "cuda")]
+        _ => anyhow::bail!("engine backend/session mismatch"),
+    };
+    Ok(logits)
 }
 
 /// One queued request. Admission has already been passed — the permit
@@ -180,8 +479,9 @@ async fn run_engine(cfg: EngineConfig, mut rx: mpsc::Receiver<EngineRequest>) {
     // `ExtractKvRows` key off.
     let mut padded_len = 0usize;
     let mut step = 0usize;
-    // Held while any slot is active — see `EngineConfig.inference_lock`.
-    let mut lock_guard: Option<tokio::sync::OwnedMutexGuard<()>> = None;
+    // Held while any slot is active — the model's serialization lock
+    // plus (TP) the upgraded model handle. See `ActiveSession`.
+    let mut session: Option<ActiveSession> = None;
 
     tracing::info!(
         model = %cfg.model_id,
@@ -207,38 +507,39 @@ async fn run_engine(cfg: EngineConfig, mut rx: mpsc::Receiver<EngineRequest>) {
             }
         }
 
-        // Take the model's inference lock before touching cache state;
-        // release it whenever the batch drains so vision/non-streaming
+        // Enter the active phase before touching cache state; release
+        // it whenever the batch drains so vision/non-streaming
         // requests get their turn.
-        if !joins.is_empty() && lock_guard.is_none() {
-            lock_guard = Some(Arc::clone(&cfg.inference_lock).lock_owned().await);
+        if !joins.is_empty() && session.is_none() {
+            match cfg.backend.acquire().await {
+                Some(s) => session = Some(s),
+                None => break 'main, // model unloaded (TP weak gone)
+            }
         }
+        let Some(sess) = session.as_mut() else {
+            continue;
+        };
 
         let needs_compaction = slots
             .iter()
             .any(|s| s.finished.is_some() || s.hangup.load(Ordering::Acquire));
         if (!joins.is_empty() || needs_compaction)
-            && let Err(e) = rebatch(&cfg, &mut slots, joins, &mut padded_len, &mut step).await
+            && let Err(e) = rebatch(&cfg, sess, &mut slots, joins, &mut padded_len, &mut step).await
         {
             fail_engine(&cfg, &mut slots, &mut rx, &e);
             break 'main;
         }
         if slots.is_empty() {
-            lock_guard = None; // every join finished during prefill
+            session = None; // every join finished during prefill
             continue;
         }
 
         // One lockstep decode step.
         let tokens: Vec<u32> = slots.iter().map(|s| s.next_token).collect();
         let prefix_lens: Vec<usize> = slots.iter().map(|s| s.prefix_len).collect();
-        let rows = match cfg
-            .worker
-            .forward_logits_batch(cfg.handle, tokens, prefix_lens, padded_len, step)
-            .await
-        {
+        let rows = match op_step(&cfg, sess, tokens, prefix_lens, padded_len, step).await {
             Ok(rows) => rows,
             Err(e) => {
-                let e = anyhow::anyhow!("batched decode step {step}: {e}");
                 fail_engine(&cfg, &mut slots, &mut rx, &e);
                 break 'main;
             }
@@ -330,7 +631,7 @@ fn fail_engine(
 ) {
     let chain = format!("{error:#}");
     if is_device_fault(&chain) {
-        cfg.poisoned.store(true, Ordering::Release);
+        cfg.backend.mark_poisoned();
         tracing::error!(
             model = %cfg.model_id,
             error = %chain,
@@ -356,6 +657,7 @@ fn fail_engine(
 /// geometry.
 async fn rebatch(
     cfg: &EngineConfig,
+    session: &mut ActiveSession,
     slots: &mut Vec<Slot>,
     joins: Vec<EngineRequest>,
     padded_len: &mut usize,
@@ -363,7 +665,7 @@ async fn rebatch(
 ) -> Result<()> {
     // 1. Extract survivors BEFORE any prefill clobbers the live state.
     let mut kept: Vec<Slot> = Vec::new();
-    let mut extracted: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = Vec::new();
+    let mut extracted: Vec<(u64, usize)> = Vec::new();
     let leavers_or_joiners = joins.len()
         + slots
             .iter()
@@ -380,12 +682,8 @@ async fn rebatch(
             .iter()
             .map(|&i| (i, slots[i].prefix_len))
             .collect();
-        let ids = cfg
-            .worker
-            .extract_kv_rows(cfg.handle, rows, *padded_len, *step)
-            .await
-            .map_err(|e| anyhow::anyhow!("extract_kv_rows: {e}"))?;
-        for (&i, (id, _bytes)) in survivors.iter().zip(ids) {
+        let ids = op_extract(cfg, session, rows, *padded_len, *step).await?;
+        for (&i, id) in survivors.iter().zip(ids) {
             let new_len = slots[i].prefix_len + *step;
             extracted.push((id, new_len));
         }
@@ -402,11 +700,14 @@ async fn rebatch(
 
     // 2. Prefill each join at B=1 (prefix cache + chunked prefill
     //    exactly as the per-request path).
-    let mut assemble: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = extracted.clone();
+    let mut assemble: Vec<(u64, usize)> = extracted.clone();
     for req in joins {
         let req_span = req.span.clone();
         // `None` = finished during prefill (EOS / hangup / max_new 0).
-        if let Some((slot, snap_id)) = prefill_join(cfg, req).instrument_in(req_span).await? {
+        if let Some((slot, snap_id)) = prefill_join(cfg, session, req)
+            .instrument_in(req_span)
+            .await?
+        {
             assemble.push((snap_id, slot.prompt_len));
             kept.push(slot);
         }
@@ -416,20 +717,15 @@ async fn rebatch(
     if kept.is_empty() {
         // Nothing active. Temp snapshots for extraction are dropped.
         for (id, _) in &assemble {
-            let _ = cfg.worker.drop_kv_snapshot(cfg.handle, *id).await;
+            op_drop_snap(cfg, session, *id).await;
         }
         *padded_len = 0;
         *step = 0;
         return Ok(());
     }
-    let seqs: Vec<(super::device_worker::jobs::KvSnapshotId, usize)> = assemble.clone();
-    let new_padded = cfg
-        .worker
-        .assemble_kv_batch(cfg.handle, seqs)
-        .await
-        .map_err(|e| anyhow::anyhow!("assemble_kv_batch: {e}"))?;
+    let new_padded = op_assemble(cfg, session, assemble.clone()).await?;
     for (id, _) in &assemble {
-        let _ = cfg.worker.drop_kv_snapshot(cfg.handle, *id).await;
+        op_drop_snap(cfg, session, *id).await;
     }
     *padded_len = new_padded;
     *step = 0;
@@ -443,8 +739,9 @@ async fn rebatch(
 /// up) — its Finish has been emitted and no slot joins the batch.
 async fn prefill_join(
     cfg: &EngineConfig,
+    session: &mut ActiveSession,
     req: EngineRequest,
-) -> Result<Option<(Slot, super::device_worker::jobs::KvSnapshotId)>> {
+) -> Result<Option<(Slot, u64)>> {
     use candle_transformers::generation::Sampling;
 
     let EngineRequest {
@@ -472,34 +769,10 @@ async fn prefill_join(
         LogitsProcessor::from_sampling(seed, sampling)
     };
 
-    let prefix_cache = cfg.prefix_cache.as_deref();
     let prompt_len = prompt_tokens.len();
     let prefill_start = std::time::Instant::now();
-    let reused =
-        restore_or_clear_via_worker(&cfg.worker, cfg.handle, prefix_cache, &prompt_tokens).await?;
-    let cut = if prefix_cache.is_some() {
-        stable_snapshot_cut(&prompt_tokens, cfg.tokenizer.token_to_id("<|im_start|>"))
-            .filter(|&c| c > reused)
-    } else {
-        None
-    };
-    let logits_vec = match cut {
-        Some(c) => {
-            chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens[..c], reused)
-                .await?;
-            store_prefix_snapshot_via_worker(
-                &cfg.worker,
-                cfg.handle,
-                prefix_cache,
-                prompt_tokens[..c].to_vec(),
-            )
-            .await;
-            chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens, c).await?
-        }
-        None => chunked_prefill_via_worker(&cfg.worker, cfg.handle, &prompt_tokens, reused).await?,
-    };
+    let logits_vec = op_prefill(cfg, session, &prompt_tokens).await?;
     let prefill_elapsed = prefill_start.elapsed();
-    cfg.prefill_rate.record(prompt_len, prefill_elapsed);
 
     // First token from the prefill logits.
     let generated: Vec<u32> = Vec::new();
@@ -572,11 +845,7 @@ async fn prefill_join(
     }
 
     // Snapshot the freshly prefilled state for assembly.
-    let (snap_id, _bytes) = cfg
-        .worker
-        .snapshot_kv(cfg.handle)
-        .await
-        .map_err(|e| anyhow::anyhow!("snapshot after prefill: {e}"))?;
+    let snap_id = op_snapshot_seq(cfg, session).await?;
     Ok(Some((slot, snap_id)))
 }
 
@@ -817,16 +1086,18 @@ mod tests {
         let admission = AdmissionController::new(&admission_cfg);
         let engine = EngineHandle::spawn(EngineConfig {
             model_id: "qwen3_next-tiny".into(),
-            worker: Arc::clone(&worker),
-            handle,
             tokenizer: tiny_tokenizer(512),
-            prefix_cache: None,
-            prefill_rate: Arc::new(PrefillRateEma::new()),
             reasoning_tokens: None,
             tool_call_tokens: None,
-            poisoned: Arc::new(AtomicBool::new(false)),
-            inference_lock: Arc::new(tokio::sync::Mutex::new(())),
             max_slots: 3,
+            backend: BackendConfig::Single {
+                worker: Arc::clone(&worker),
+                handle,
+                prefix_cache: None,
+                prefill_rate: Arc::new(PrefillRateEma::new()),
+                poisoned: Arc::new(AtomicBool::new(false)),
+                inference_lock: Arc::new(tokio::sync::Mutex::new(())),
+            },
         });
 
         let prompts: [&[u32]; 3] = [&[1, 2, 3], &[4, 5], &[7, 3, 2, 5, 6]];

crates/neuron/src/harness/tp/mod.rs

@@ -94,6 +94,37 @@ impl TpLeaderModel {
         }
     }
 
+    /// Lockstep batched decode step on rank 0 (#98). Only qwen3_5
+    /// batches (same gate as snapshots).
+    pub fn forward_batch_decode(
+        &mut self,
+        input: &candle_core::Tensor,
+        positions: &[usize],
+        attn_mask: Option<&candle_core::Tensor>,
+    ) -> candle_core::Result<candle_core::Tensor> {
+        match self {
+            TpLeaderModel::Qwen3(_) => {
+                candle_core::bail!("forward_batch_decode: qwen3 (dense) has no batched decode")
+            }
+            TpLeaderModel::Qwen3_5(m) => m.forward_batch_decode(input, positions, attn_mask),
+        }
+    }
+
+    /// Padding mask for a batched decode step (#98).
+    pub fn batch_decode_mask(
+        &self,
+        prefix_lens: &[usize],
+        padded_len: usize,
+        total_len: usize,
+    ) -> candle_core::Result<Option<candle_core::Tensor>> {
+        match self {
+            TpLeaderModel::Qwen3(_) => {
+                candle_core::bail!("batch_decode_mask: qwen3 (dense) has no batched decode")
+            }
+            TpLeaderModel::Qwen3_5(m) => m.batch_decode_mask(prefix_lens, padded_len, total_len),
+        }
+    }
+
     /// Whether this arch supports prefix snapshots (#11). Gates the
     /// pool fan-out so unsupported archs never even ask the ranks.
     pub fn supports_kv_snapshot(&self) -> bool {
@@ -837,6 +868,179 @@ impl WorkerPool {
         }
     }
 
+    /// One lockstep batched decode step across every rank (#98).
+    /// Same fan-out / leader-forward / always-drain shape as
+    /// [`Self::generate_step`]; every rank derives positions + mask
+    /// locally from the broadcast geometry. Returns one `[vocab]`
+    /// logits row per batch row from the leader's rank-0 shard.
+    #[cfg(feature = "cuda")]
+    pub async fn generate_step_batch(
+        &mut self,
+        model_id: &str,
+        leader_handle: super::device_worker::TpHandle,
+        tokens: Vec<u32>,
+        prefix_lens: Vec<usize>,
+        padded_len: usize,
+        step: usize,
+    ) -> Result<Vec<Vec<f32>>> {
+        for w in &mut self.workers {
+            w.send_only(&WorkerRequest::GenerateStepBatch {
+                model_id: model_id.to_string(),
+                tokens: tokens.clone(),
+                prefix_lens: prefix_lens.clone(),
+                padded_len,
+                step,
+            })
+            .await?;
+        }
+
+        let timeout = tp_step_timeout();
+        let leader_fut = self.leader_worker.tp_forward_logits_batch(
+            leader_handle,
+            tokens,
+            prefix_lens,
+            padded_len,
+            step,
+        );
+        let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
+            Ok(r) => r,
+            Err(_elapsed) => {
+                // Watchdog (#17 Stage 2) — same rationale as
+                // `generate_step`: abort the wedged comm, fail without
+                // draining, let auto-recovery restart the pool.
+                self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
+                anyhow::bail!(
+                    "tp watchdog: leader batched forward exceeded {}s deadline; aborted wedged \
+                     NCCL comm — model will auto-recover",
+                    timeout.as_secs()
+                );
+            }
+        };
+
+        let worker_errors = drain_workers(&mut self.workers, |r| match r {
+            WorkerResponse::GenerateStepOk => Ok(()),
+            WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
+            other => Err(format!("expected GenerateStepOk, got {other:?}")),
+        })
+        .await;
+
+        match leader_result {
+            Ok(rows) => {
+                if worker_errors.is_empty() {
+                    Ok(rows)
+                } else {
+                    anyhow::bail!(
+                        "GenerateStepBatch: leader succeeded but workers failed: {}",
+                        worker_errors.join("; ")
+                    )
+                }
+            }
+            Err(e) => Err(anyhow::Error::new(e).context(if worker_errors.is_empty() {
+                "GenerateStepBatch: leader forward failed".to_string()
+            } else {
+                format!(
+                    "GenerateStepBatch: leader forward failed and workers also failed: {}",
+                    worker_errors.join("; ")
+                )
+            })),
+        }
+    }
+
+    /// Assemble stored per-sequence snapshots into every rank's live
+    /// batched state (#98). Snapshot ids in `seqs` are pool-minted;
+    /// every rank (leader + subprocesses) assembles the same geometry
+    /// and the returned padded length is asserted identical.
+    #[cfg(feature = "cuda")]
+    pub async fn assemble_kv_batch(
+        &mut self,
+        model_id: &str,
+        leader_handle: super::device_worker::TpHandle,
+        seqs: Vec<(u64, usize)>,
+    ) -> Result<usize> {
+        for w in &mut self.workers {
+            w.send_only(&WorkerRequest::AssembleKvBatch {
+                model_id: model_id.to_string(),
+                seqs: seqs.clone(),
+            })
+            .await?;
+        }
+        let leader_result = self
+            .leader_worker
+            .tp_assemble_kv_batch(leader_handle, seqs)
+            .await;
+        let leader_padded = leader_result.as_ref().ok().copied();
+        let worker_errors = drain_workers(&mut self.workers, |r| match r {
+            WorkerResponse::KvBatchAssembled { padded_len } => {
+                if leader_padded.is_some_and(|lp| lp as u64 != padded_len) {
+                    Err(format!(
+                        "rank assembled padded_len {padded_len} != leader {}",
+                        leader_padded.unwrap_or(0)
+                    ))
+                } else {
+                    Ok(())
+                }
+            }
+            WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
+            other => Err(format!("expected KvBatchAssembled, got {other:?}")),
+        })
+        .await;
+        let padded = leader_result.map_err(|e| {
+            anyhow::Error::new(e).context("AssembleKvBatch: leader assembly failed")
+        })?;
+        if !worker_errors.is_empty() {
+            anyhow::bail!(
+                "AssembleKvBatch: leader succeeded but workers failed: {}",
+                worker_errors.join("; ")
+            );
+        }
+        Ok(padded)
+    }
+
+    /// Extract live batched-state rows into per-sequence snapshots on
+    /// every rank (#98), stored under the pre-minted `snapshot_ids`
+    /// (one per row, minted from the pool's snapshot counter).
+    #[cfg(feature = "cuda")]
+    pub async fn extract_kv_rows(
+        &mut self,
+        model_id: &str,
+        leader_handle: super::device_worker::TpHandle,
+        rows: Vec<(usize, usize)>,
+        padded_len: usize,
+        steps: usize,
+        snapshot_ids: Vec<u64>,
+    ) -> Result<()> {
+        for w in &mut self.workers {
+            w.send_only(&WorkerRequest::ExtractKvRows {
+                model_id: model_id.to_string(),
+                rows: rows.clone(),
+                padded_len,
+                steps,
+                snapshot_ids: snapshot_ids.clone(),
+            })
+            .await?;
+        }
+        let leader_result = self
+            .leader_worker
+            .tp_extract_kv_rows(leader_handle, rows, padded_len, steps, snapshot_ids)
+            .await;
+        let worker_errors = drain_workers(&mut self.workers, |r| match r {
+            WorkerResponse::KvRowsExtracted { .. } => Ok(()),
+            WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
+            other => Err(format!("expected KvRowsExtracted, got {other:?}")),
+        })
+        .await;
+        leader_result.map_err(|e| {
+            anyhow::Error::new(e).context("ExtractKvRows: leader extraction failed")
+        })?;
+        if !worker_errors.is_empty() {
+            anyhow::bail!(
+                "ExtractKvRows: leader succeeded but workers failed: {}",
+                worker_errors.join("; ")
+            );
+        }
+        Ok(())
+    }
+
     /// Image-bearing variant of [`Self::generate_step`] for the
     /// single-shot vision prefill. Identical fan-out / leader-forward /
     /// drain shape, but every rank runs the encode + splice path:

crates/neuron/src/harness/tp/rpc.rs

@@ -115,6 +115,48 @@ pub enum WorkerRequest {
         chunk_size: usize,
     },
 
+    /// One lockstep batched decode step (#98): `tokens[i]` is batch
+    /// row i's next token at position `prefix_lens[i] + step`.
+    /// Identical on every rank (the leader mirrors it through
+    /// `Job::TpForwardLogitsBatch`); each rank derives the per-row
+    /// positions and padding mask locally and discards its logits,
+    /// same as `GenerateStep`.
+    GenerateStepBatch {
+        model_id: String,
+        tokens: Vec<u32>,
+        prefix_lens: Vec<usize>,
+        /// Uniform padded KV length the batch was assembled to.
+        padded_len: usize,
+        /// Decode steps since the last rebatch.
+        step: usize,
+    },
+
+    /// Assemble stored per-sequence snapshots into one batched cache
+    /// state and install it as this rank's live state (#98). `seqs`
+    /// pairs pool-minted snapshot ids with true token lengths; every
+    /// rank holds symmetric shard snapshots under the same ids, so
+    /// the assembled geometry (padded length) is identical across
+    /// ranks. Source snapshots remain stored. Replies
+    /// `KvBatchAssembled { padded_len }`.
+    AssembleKvBatch {
+        model_id: String,
+        seqs: Vec<(u64, usize)>,
+    },
+
+    /// Extract rows of this rank's live batched state back into
+    /// contiguous single-sequence snapshots (#98) — the first half of
+    /// a rebatch. `rows` pairs batch-row indexes with prefix lengths;
+    /// `snapshot_ids` carries the pool-minted id to store each
+    /// extracted row under (one per entry of `rows`, so every rank
+    /// keys identically). Replies `KvRowsExtracted`.
+    ExtractKvRows {
+        model_id: String,
+        rows: Vec<(usize, usize)>,
+        padded_len: usize,
+        steps: usize,
+        snapshot_ids: Vec<u64>,
+    },
+
     /// Reset the KV cache for this model on this rank. Sent at the
     /// start of every inference so a fresh request doesn't accidentally
     /// attend over the previous one's tokens.
@@ -192,6 +234,14 @@ pub enum WorkerResponse {
     /// Reply to `ClearKvCache`. Empty payload.
     KvCacheCleared,
 
+    /// Reply to `AssembleKvBatch`. The uniform padded KV length the
+    /// batch was assembled to — the leader asserts every rank agrees.
+    KvBatchAssembled { padded_len: u64 },
+
+    /// Reply to `ExtractKvRows`. Total bytes of the extracted
+    /// snapshots on this rank (budget accounting only).
+    KvRowsExtracted { bytes: u64 },
+
     /// Reply to `QueryVram`. This rank's device VRAM in MiB.
     VramInfo { free_mb: u64, total_mb: u64 },
 

crates/neuron/src/harness/tp/tp_qwen3_5.rs

@@ -1739,6 +1739,63 @@ impl TpQwen3_5Model {
         self.forward_inner(input, offset, None, None, None)
     }
 
+    /// Lockstep batched decode step (#98) — the TP mirror of
+    /// `Qwen3_5Model::forward_batch_decode`: `(B, 1)` input, per-row
+    /// positions, padding mask from [`Self::batch_decode_mask`].
+    /// Every rank runs this with identical inputs, so the
+    /// replicated-hidden-state invariant holds per batch row.
+    /// Text-only: `rope_delta` is ignored — positions are explicit
+    /// and vision requests never enter the batch path.
+    pub fn forward_batch_decode(
+        &mut self,
+        input: &Tensor,
+        positions: &[usize],
+        attn_mask: Option<&Tensor>,
+    ) -> candle_core::Result<Tensor> {
+        let (b, l) = input.dims2()?;
+        if l != 1 {
+            candle_core::bail!("forward_batch_decode: expected (B, 1) input, got (B, {l})");
+        }
+        if positions.len() != b {
+            candle_core::bail!(
+                "forward_batch_decode: {} positions for batch of {b}",
+                positions.len()
+            );
+        }
+        let mut h = self.embed_tokens.forward(input)?;
+        let (cos, sin) = self.rotary.batch_cos_sin(positions)?;
+        for layer in &mut self.layers {
+            h = layer.forward(&h, attn_mask, &cos, &sin)?;
+        }
+        self.norm.forward(&h)
+    }
+
+    /// Additive padding mask for a batched decode step — the TP mirror
+    /// of `Qwen3_5Model::batch_decode_mask`: `(B, 1, 1, total_len)`,
+    /// `-inf` on each row's padding gap `[prefix_lens[i], padded_len)`,
+    /// `None` when no row is padded.
+    pub fn batch_decode_mask(
+        &self,
+        prefix_lens: &[usize],
+        padded_len: usize,
+        total_len: usize,
+    ) -> candle_core::Result<Option<Tensor>> {
+        if prefix_lens.iter().all(|&len| len == padded_len) {
+            return Ok(None);
+        }
+        let minf = f32::NEG_INFINITY;
+        let b = prefix_lens.len();
+        let mask: Vec<f32> = prefix_lens
+            .iter()
+            .flat_map(|&len| {
+                (0..total_len).map(move |j| if j >= len && j < padded_len { minf } else { 0. })
+            })
+            .collect();
+        Ok(Some(
+            Tensor::from_vec(mask, (b, 1, 1, total_len), &self.device)?.to_dtype(self.dtype)?,
+        ))
+    }
+
     /// Forward for a vision-prefill chunk: optional image-embedding
     /// splice plus explicit interleaved-M-RoPE `position_ids` (the
     /// chunk's slice of the full prompt's 3D positions). Used by
@@ -1948,6 +2005,33 @@ impl TpQwen3_5ForCausalLM {
         hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
     }
 
+    /// Lockstep batched decode step (#98): `(B, 1)` input, per-row
+    /// positions, padding mask from
+    /// [`TpQwen3_5Model::batch_decode_mask`]. Returns `(B, 1, vocab)`.
+    pub fn forward_batch_decode(
+        &mut self,
+        input: &Tensor,
+        positions: &[usize],
+        attn_mask: Option<&Tensor>,
+    ) -> candle_core::Result<Tensor> {
+        let hidden = self
+            .base
+            .forward_batch_decode(input, positions, attn_mask)?;
+        hidden.apply(&self.lm_head)
+    }
+
+    /// Padding mask for a batched decode step — see
+    /// [`TpQwen3_5Model::batch_decode_mask`].
+    pub fn batch_decode_mask(
+        &self,
+        prefix_lens: &[usize],
+        padded_len: usize,
+        total_len: usize,
+    ) -> candle_core::Result<Option<Tensor>> {
+        self.base
+            .batch_decode_mask(prefix_lens, padded_len, total_len)
+    }
+
     /// Forward for a vision-prefill chunk (optional image splice +
     /// explicit interleaved-M-RoPE `position_ids`). Mirrors `forward`
     /// but routes through `TpQwen3_5Model::forward_with_positions`.

crates/neuron/src/harness/tp/worker.rs

@@ -82,6 +82,37 @@ impl WorkerModel {
         }
     }
 
+    /// Lockstep batched decode step (#98). Only qwen3_5 batches — the
+    /// leader gates on its own `TpLeaderModel` support first.
+    fn forward_batch_decode(
+        &mut self,
+        input: &candle_core::Tensor,
+        positions: &[usize],
+        attn_mask: Option<&candle_core::Tensor>,
+    ) -> candle_core::Result<candle_core::Tensor> {
+        match self {
+            WorkerModel::Qwen3(_) => {
+                candle_core::bail!("forward_batch_decode: qwen3 (dense) has no batched decode")
+            }
+            WorkerModel::Qwen3_5(m) => m.forward_batch_decode(input, positions, attn_mask),
+        }
+    }
+
+    /// Padding mask for a batched decode step (#98).
+    fn batch_decode_mask(
+        &self,
+        prefix_lens: &[usize],
+        padded_len: usize,
+        total_len: usize,
+    ) -> candle_core::Result<Option<candle_core::Tensor>> {
+        match self {
+            WorkerModel::Qwen3(_) => {
+                candle_core::bail!("batch_decode_mask: qwen3 (dense) has no batched decode")
+            }
+            WorkerModel::Qwen3_5(m) => m.batch_decode_mask(prefix_lens, padded_len, total_len),
+        }
+    }
+
     /// Capture this rank's cache state for a prefix snapshot (#11).
     /// Only qwen3_5 exposes its state; the dense qwen3 arch errors —
     /// the leader never asks, because it gates on its own
@@ -248,6 +279,23 @@ impl WorkerState {
                 image_data_uris,
                 chunk_size,
             ),
+            WorkerRequest::GenerateStepBatch {
+                model_id,
+                tokens,
+                prefix_lens,
+                padded_len,
+                step,
+            } => self.handle_generate_step_batch(&model_id, tokens, prefix_lens, padded_len, step),
+            WorkerRequest::AssembleKvBatch { model_id, seqs } => {
+                self.handle_assemble_kv_batch(&model_id, &seqs)
+            }
+            WorkerRequest::ExtractKvRows {
+                model_id,
+                rows,
+                padded_len,
+                steps,
+                snapshot_ids,
+            } => self.handle_extract_kv_rows(&model_id, &rows, padded_len, steps, &snapshot_ids),
             WorkerRequest::ClearKvCache { model_id } => self.handle_clear_kv_cache(&model_id),
             WorkerRequest::SnapshotKvCache {
                 model_id,
@@ -541,6 +589,206 @@ impl WorkerState {
         }
     }
 
+    /// One lockstep batched decode step on this rank (#98). Mirrors
+    /// the leader's `forward_logits_batch`: derives per-row positions
+    /// and the padding mask locally from the broadcast geometry, runs
+    /// the batched forward, and discards the logits (the leader
+    /// samples from rank 0; the NCCL collectives are the point).
+    #[cfg(feature = "cuda")]
+    fn handle_generate_step_batch(
+        &mut self,
+        model_id: &str,
+        tokens: Vec<u32>,
+        prefix_lens: Vec<usize>,
+        padded_len: usize,
+        step: usize,
+    ) -> WorkerResponse {
+        use candle_core::Tensor;
+
+        let b = tokens.len();
+        if b == 0 || prefix_lens.len() != b {
+            return WorkerResponse::Error {
+                kind: "bad_request".into(),
+                message: format!(
+                    "GenerateStepBatch: {b} tokens vs {} prefix_lens",
+                    prefix_lens.len()
+                ),
+            };
+        }
+        let Some(model) = self.models.get_mut(model_id) else {
+            return WorkerResponse::Error {
+                kind: "model_not_loaded".into(),
+                message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
+            };
+        };
+        let device = model.device().clone();
+        let result = (|| -> candle_core::Result<()> {
+            let input = Tensor::from_vec(tokens, (b, 1), &device)?;
+            let positions: Vec<usize> = prefix_lens.iter().map(|&len| len + step).collect();
+            let total_len = padded_len + step + 1;
+            let mask = model.batch_decode_mask(&prefix_lens, padded_len, total_len)?;
+            model.forward_batch_decode(&input, &positions, mask.as_ref())?;
+            Ok(())
+        })();
+        match result {
+            Ok(()) => WorkerResponse::GenerateStepOk,
+            Err(e) => {
+                tracing::warn!(
+                    rank = self.config.rank,
+                    model = %model_id,
+                    error = %e,
+                    "worker GenerateStepBatch: forward failed"
+                );
+                WorkerResponse::Error {
+                    kind: "forward_failed".into(),
+                    message: format!("TP batched forward: {e}"),
+                }
+            }
+        }
+    }
+
+    #[cfg(not(feature = "cuda"))]
+    fn handle_generate_step_batch(
+        &mut self,
+        _model_id: &str,
+        _tokens: Vec<u32>,
+        _prefix_lens: Vec<usize>,
+        _padded_len: usize,
+        _step: usize,
+    ) -> WorkerResponse {
+        WorkerResponse::Error {
+            kind: "cuda_feature_not_enabled".into(),
+            message: "GenerateStepBatch requires --features cuda".into(),
+        }
+    }
+
+    /// Assemble stored per-sequence snapshots into this rank's live
+    /// batched state (#98). Mirrors the leader's `assemble_kv_batch`.
+    #[cfg(feature = "cuda")]
+    fn handle_assemble_kv_batch(
+        &mut self,
+        model_id: &str,
+        seqs: &[(u64, usize)],
+    ) -> WorkerResponse {
+        let Some(model) = self.models.get_mut(model_id) else {
+            return WorkerResponse::Error {
+                kind: "model_not_loaded".into(),
+                message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
+            };
+        };
+        let mut pairs = Vec::with_capacity(seqs.len());
+        for (id, len) in seqs {
+            let Some(snap) = self.kv_snapshots.get(&(model_id.to_string(), *id)) else {
+                return WorkerResponse::Error {
+                    kind: "snapshot_not_found".into(),
+                    message: format!(
+                        "AssembleKvBatch: no snapshot {id} for '{model_id}' on rank {}",
+                        self.config.rank
+                    ),
+                };
+            };
+            pairs.push((snap, *len));
+        }
+        let result =
+            crate::harness::arch::qwen3_5::snapshot::assemble_batch(&pairs).and_then(|batch| {
+                model.restore_kv_cache(&batch.snapshot)?;
+                Ok(batch.padded_len)
+            });
+        match result {
+            Ok(padded_len) => WorkerResponse::KvBatchAssembled {
+                padded_len: padded_len as u64,
+            },
+            Err(e) => WorkerResponse::Error {
+                kind: "assemble_failed".into(),
+                message: format!("AssembleKvBatch: {e}"),
+            },
+        }
+    }
+
+    #[cfg(not(feature = "cuda"))]
+    fn handle_assemble_kv_batch(
+        &mut self,
+        _model_id: &str,
+        _seqs: &[(u64, usize)],
+    ) -> WorkerResponse {
+        WorkerResponse::Error {
+            kind: "cuda_feature_not_enabled".into(),
+            message: "AssembleKvBatch requires --features cuda".into(),
+        }
+    }
+
+    /// Extract live batched-state rows into stored per-sequence
+    /// snapshots under pool-minted ids (#98). Mirrors the leader's
+    /// `extract_kv_rows`.
+    #[cfg(feature = "cuda")]
+    fn handle_extract_kv_rows(
+        &mut self,
+        model_id: &str,
+        rows: &[(usize, usize)],
+        padded_len: usize,
+        steps: usize,
+        snapshot_ids: &[u64],
+    ) -> WorkerResponse {
+        if snapshot_ids.len() != rows.len() {
+            return WorkerResponse::Error {
+                kind: "bad_request".into(),
+                message: format!(
+                    "ExtractKvRows: {} rows vs {} snapshot_ids",
+                    rows.len(),
+                    snapshot_ids.len()
+                ),
+            };
+        }
+        let Some(model) = self.models.get(model_id) else {
+            return WorkerResponse::Error {
+                kind: "model_not_loaded".into(),
+                message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
+            };
+        };
+        let live = match model.snapshot_kv_cache() {
+            Ok(s) => s,
+            Err(e) => {
+                return WorkerResponse::Error {
+                    kind: "extract_failed".into(),
+                    message: format!("ExtractKvRows: snapshot live state: {e}"),
+                };
+            }
+        };
+        let mut total_bytes = 0u64;
+        for (&(row, prefix_len), &id) in rows.iter().zip(snapshot_ids) {
+            match crate::harness::arch::qwen3_5::snapshot::extract_row(
+                &live, row, prefix_len, padded_len, steps,
+            ) {
+                Ok(snap) => {
+                    total_bytes += snap.size_bytes();
+                    self.kv_snapshots.insert((model_id.to_string(), id), snap);
+                }
+                Err(e) => {
+                    return WorkerResponse::Error {
+                        kind: "extract_failed".into(),
+                        message: format!("ExtractKvRows: row {row}: {e}"),
+                    };
+                }
+            }
+        }
+        WorkerResponse::KvRowsExtracted { bytes: total_bytes }
+    }
+
+    #[cfg(not(feature = "cuda"))]
+    fn handle_extract_kv_rows(
+        &mut self,
+        _model_id: &str,
+        _rows: &[(usize, usize)],
+        _padded_len: usize,
+        _steps: usize,
+        _snapshot_ids: &[u64],
+    ) -> WorkerResponse {
+        WorkerResponse::Error {
+            kind: "cuda_feature_not_enabled".into(),
+            message: "ExtractKvRows requires --features cuda".into(),
+        }
+    }
+
     /// Image-bearing prefill on this rank. Preprocesses each source data
     /// URI through the same deterministic `preprocess_data_uri` the
     /// leader runs, encodes through this rank's replicated tower, and