Fat-station continuous batching — the live-path integration (precise spec)

Target confirmed (2026-06-30): the live mistral runs as ONE full-model fat-station (fat-station.rs, pipeline: Mutex<NativeLlamaPipeline>) because Google hasn't granted more L4 quota yet — meshing 2–3 mistrals is the growth path. decode_step_batch is the right primitive for both: it batches concurrent requests on one station now, and becomes the per-station batched step when sharded later. Build single-station first.

Why this is the integration point (not the TS scheduler)

Station.generateTokens (host-side station.ts:514) HTTP/UDS-POSTs /generate to the Rust fat-station; the decode loop runs in Rust under station.pipeline.lock(). No TS loads the .node addon in the serving path. So the live serialize-all is the Rust Mutex, and continuous batching must live here. The napi binding + TS ContinuousBatchScheduler (gnosis 3d63dd73, vitest-green) validated the scheduling design — this is that design, ported to Rust where the decode actually runs.

What /generate does today (the serialize-all)

handle_generate (fat-station.rs:5776) per request: parse tokens/max_new_tokens/ temperature/seed/resume/terminal_prosody → PREFILL the prompt in chunks (populating self.kv_cache at [0,prompt_len)) → DECODE loop calling generate_forward_one(token, pos) (:979) + sample per step → respond with the full generated: Vec<u32> (request/response, not SSE — the host streams). The pipeline lock is held for the whole prefill+decode of one request, so two concurrent /generates run strictly serially. Each leaves (B−1)·w of weight-read on the table (the law).

The capture — a scheduler thread batching concurrent requests

decode_step_batch(&mut [KVCache], &mut [usize] positions, &[u32] tokens) (model.rs, parity-verified decode_step_batch_parity) advances K lanes by one token, one weight read for the wave. Wrap the station's one pipeline in a scheduler so concurrent /generates share each decode step.

Components (all in fat-station.rs, flag-gated)

1. struct BatchScheduler owned by Station (behind the existing pipeline Mutex, or a dedicated Mutex/channel). Holds: pending: VecDeque<Job>, the K seq_kv: Vec<KVCache>

   • positions: Vec<usize> lane slabs (alloc once at capacity), and per-lane Option<Lane>. Job { tokens, max_new_tokens, temperature, seed, terminal_prosody, reply: Sender<Vec<u32>> }. Lane { job, feed_count, generated: Vec<u32>, last_token, rng }.

2. handle_generate (flag ON): build a Job with a std::sync::mpsc::channel, push to pending (under the scheduler lock), notify the decode thread (Condvar), then block on reply.recv() and respond with the returned tokens. Same wire contract as today (readGeneratedTokensResponse unchanged) — only the internal execution batches. Flag OFF → the existing lock-and-loop path, byte-identical.

3. The decode thread (one, spawned at startup, owns the pipeline): loop —

   • fill(): for each free lane k with pending non-empty → seq_kv[k].clear(); positions[k]=0; assign the job; (this is resetLane done in-thread).
   • gather one input token per live lane: tokens[k] = feed_count < prompt.len() ? prompt[feed_count] : last_token.
   • logits = pipeline.decode_step_batch(&mut seq_kv, &mut positions, &tokens)? — the Fanout: one weight read serves the wave.
   • per lane (fold): feed_count += 1; if still in prompt → continue (KV seed only); else sample logits[k*vocab..] with the lane's temp/seed (reuse the existing greedy/ temperature-multinomial sampler verbatim), push to lane.generated, set last_token; retire on stop-token / terminal_prosody close / max_new_tokens → reply.send(generated); free the lane (so fill() re-uses it — true continuous batching).
   • if no live lanes and no pending → wait on the Condvar.

4. Flag: MOONSHINE_CONTINUOUS_BATCH (env, default OFF). Capacity: MOONSHINE_BATCH_LANES (e.g. 8), bounded by kv_cache.max_seq_len arena.

Parity (CPU, before any deploy)

Mirror decode_step_batch_parity: N concurrent Jobs through the scheduler must yield, per job, token-identical output to N serial handle_generate (greedy/seeded) runs of the same prompt. The kernel parity already guarantees the per-step math; this certifies the scheduler/prefill/retire/lane-reuse logic. Add tests/fat_station_batch_parity.rs driving the scheduler struct directly (no socket).

Throughput gate (GPU, owner-gated)

On the L4 Cloud Run rev: K concurrent /generate → aggregate tok/s approaches K× single until KV/SM saturates (the realized analogue of prefill's measured 6.2×). Confirm coherence ("Paris…"). Canary one rev at 0% → burst-test (min-instances=0 cold-start aware) → shift only after parity + throughput green + owner OK. Rollback = flag off.

Mesh extension (when L4 quota grows → 2–3 mistral)

Same decode_step_batch becomes the per-station batched step; the sharded path (fat-station-flow.rs handle_forward/handle_lm_head per layer-range) gains a batched /forward (K residuals through one weight read per station) and the host (openai-server.ts) drives a batched orchestration loop. The single-station scheduler is the foundation; the mesh distributes the same wave across stations.

Files

• src/bin/fat-station.rs — handle_generate (:5776), generate_forward_one (:979), the sampler, Station struct → the scheduler + decode thread.
• src/model.rs — decode_step_batch (the kernel, done) + NativeLlamaPipeline.
• tests/fat_station_batch_parity.rs — new, the CPU parity gate.
• tests/decode_step_batch_parity.rs — the kernel parity (done, green).