feat(neuron): continuous batching — batched decode + engine (#98, slices 1–3b) #121

Merged

grenade merged 5 commits from feat/98-batched-decode into main 2026-07-03 12:39:08 +00:00

Conversation 0 Commits 5 Files Changed 9 +2140 -43

grenade commented 2026-07-03 12:38:59 +00:00

Owner

Copy Link

First four slices of F4d continuous batching (#98, epic #84): gang batching with lockstep decode and cat-based rebatch — no scatter kernels, no paged attention, no CUDA kernel changes.

Arch groundwork (qwen3_5) — per-row RoPE positions (batch_cos_sin + rank-dispatched apply), assemble_batch (right-pad per-seq snapshot KV to the batch max, cat on dim 0; GDN states cat on dim 0), forward_batch_decode((B,1), positions, mask), batch_decode_mask (one padding gap per row). Keys are stored post-RoPE, so padding is position-inert.

Worker jobs — AssembleKvBatch, ForwardLogitsBatch (per-row CPU [vocab] logits), ExtractKvRows (slice a live batch row back to a contiguous B=1 snapshot — the join/leave rebatch primitive). Fix en route: dense safetensors loads pick f32 on non-CUDA devices (candle CPU has no bf16 matmul).

Engine (harness/engine.rs) — per-model task multiplexing concurrent text chat streams through one (B,1) forward per step. Joins prefill at B=1 through the existing chunked-prefill + prefix-cache paths; per-slot router tasks own the detokenizer + reasoning/tool-call state machine; the engine holds inference_lock while active so vision/non-streaming requests stay safe. Default-off: activates only when [admission] max_in_flight > 1 on a snapshot-capable worker-path model; NEURON_BATCHING=0 kill switch. No production config changes in this PR — behaviour is unchanged everywhere until an operator raises the knob.

Parity tests on the committed qwen3_next-tiny fixture at three levels: arch (ragged lockstep == sequential, hidden-for-hidden), worker (end-to-end jobs incl. a leave-path rebatch), engine (concurrent == sequential byte-identical, with mid-flight joins).

Remaining slices tracked in #98: TP batched GenerateStep (the payoff — 27B/Coder-Next are TP-2) and the O5 bench gate.

🤖 Generated with Claude Code

https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

First four slices of F4d continuous batching (#98, epic #84): gang batching with lockstep decode and cat-based rebatch — no scatter kernels, no paged attention, no CUDA kernel changes. **Arch groundwork (qwen3_5)** — per-row RoPE positions (`batch_cos_sin` + rank-dispatched apply), `assemble_batch` (right-pad per-seq snapshot KV to the batch max, cat on dim 0; GDN states cat on dim 0), `forward_batch_decode((B,1), positions, mask)`, `batch_decode_mask` (one padding gap per row). Keys are stored post-RoPE, so padding is position-inert. **Worker jobs** — `AssembleKvBatch`, `ForwardLogitsBatch` (per-row CPU `[vocab]` logits), `ExtractKvRows` (slice a live batch row back to a contiguous B=1 snapshot — the join/leave rebatch primitive). Fix en route: dense safetensors loads pick f32 on non-CUDA devices (candle CPU has no bf16 matmul). **Engine (`harness/engine.rs`)** — per-model task multiplexing concurrent text chat streams through one `(B,1)` forward per step. Joins prefill at B=1 through the existing chunked-prefill + prefix-cache paths; per-slot router tasks own the detokenizer + reasoning/tool-call state machine; the engine holds `inference_lock` while active so vision/non-streaming requests stay safe. **Default-off**: activates only when `[admission] max_in_flight > 1` on a snapshot-capable worker-path model; `NEURON_BATCHING=0` kill switch. No production config changes in this PR — behaviour is unchanged everywhere until an operator raises the knob. Parity tests on the committed qwen3_next-tiny fixture at three levels: arch (ragged lockstep == sequential, hidden-for-hidden), worker (end-to-end jobs incl. a leave-path rebatch), engine (concurrent == sequential byte-identical, with mid-flight joins). Remaining slices tracked in #98: TP batched GenerateStep (the payoff — 27B/Coder-Next are TP-2) and the O5 bench gate. 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01TczcGF7JSjJs8r15RSSGpx

grenade added 5 commits 2026-07-03 12:39:00 +00:00

feat(neuron): batched lockstep decode groundwork for qwen3_5 (#98) ... 05b6c70a3a

Slice 1 of continuous batching (F4d): arch-level batched decode with
per-row positions, assembled from per-sequence prefix snapshots.

- rope: batch_cos_sin gathers cos/sin at arbitrary per-row positions
  ((B,1,half)); apply_cos_sin dispatches on cos rank — rank-3 takes a
  local GLM rotate-half broadcast path (rope_slow only broadcasts one
  position table across the batch).
- snapshot: assemble_batch cats per-sequence KvCacheSnapshots into one
  (B,…) batched state — attention K/V right-padded along the sequence
  axis (keys are stored post-RoPE, so padding is position-inert), GDN
  conv/recurrent states cat on dim 0. Text-only: rope_delta must be 0.
- model: forward_batch_decode((B,1), positions, mask) lockstep decode
  step returning (B,1,vocab); batch_decode_mask builds the (B,1,1,S)
  additive mask covering each row's padding gap [prefix_len,
  padded_len), None when lengths are uniform.

Gold test: three ragged sequences decoded in one lockstep batch match
the same sequences decoded sequentially, hidden-for-hidden at every
step, on the tiny CPU fixture. Plus rope gather/uniform-equivalence
and mask-geometry units.

No serving-path changes yet — the engine loop (slice 3) and worker
job plumbing (slice 2) build on these primitives.

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

feat(neuron): device-worker jobs for batched decode (#98) ... 9ac1a1586c

Slice 2: AssembleKvBatch + ForwardLogitsBatch route the batched-decode
primitives through the per-device worker thread.

- AssembleKvBatch assembles stored per-sequence KV snapshots into one
  batched (B,…) live state via snapshot::assemble_batch and installs
  it through the existing restore path; replies the padded KV length.
- ForwardLogitsBatch runs one lockstep decode step: builds the (B,1)
  input on the worker's device, derives per-row positions and the
  padding mask from prefix_lens/padded_len/step, and replies one CPU
  [vocab] logits row per batch row — same tensors-never-escape
  contract as ForwardLogits.
- ModelArch::forward_batch_decode / batch_decode_mask dispatch (qwen3_5
  only; other archs error as defence in depth).
- Dense safetensors loads now pick f32 on non-CUDA devices: candle's
  CPU backend has no bf16 matmul, so the CPU fallback (and the CPU
  test worker) was broken at first forward under the hardcoded BF16.

End-to-end test: tiny qwen3_next fixture loaded through the worker,
three ragged sequences prefilled+snapshotted, assembled, and every
ForwardLogitsBatch row checked against the sequential ForwardLogits
reference at each step.

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

feat(neuron): batch-row extraction for rebatch on join/leave (#98) ... 7f2fcc3527

Slice 3a: extract_row slices one row of a batched cache state back
into a contiguous single-sequence snapshot — attention KV keeps the
prefix [0, prefix_len) plus the lockstep decode columns [padded_len,
padded_len + steps) and drops the padding gap; GDN rows are sliced
whole and deep-copied. A join/leave rebatches by extracting every
surviving row and re-running assemble_batch, preserving the
one-gap-per-row invariant without any scatter kernels.

Worker side: Job::ExtractKvRows captures the live state once and
stores each extracted row in the snapshot slab; composes with
AssembleKvBatch / DropKvSnapshot for the full rebatch.

Tests: arch-level round-trip (batched decode → extract → solo B=1
continuation matches pure-sequential) and a worker end-to-end leave
path (3-row batch → extract 2 survivors → re-assemble → continued
lockstep decode matches the sequential reference).

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

feat(neuron): lockstep batched decode engine (#98) ... ec7a5b750a

Slice 3b: per-model engine task multiplexing concurrent text chat
streams through one (B,1) forward per decode step, replacing the
per-request inference_lock serialization when [admission]
max_in_flight > 1 on a snapshot-capable worker-path model.

- harness/engine.rs: engine loop (join → B=1 prefill via the existing
  chunked-prefill + prefix-cache paths → snapshot → rebatch via
  ExtractKvRows/AssembleKvBatch; lockstep ForwardLogitsBatch steps;
  per-slot CPU sampling with per-slot LogitsProcessor + repeat-penalty
  history). Per-slot router tasks own the incremental detokenizer and
  the reasoning/tool-call state machine (same logic as route_token!)
  and emit InferenceEvents — DecodeStream's tokenizer borrow lives
  inside the router's async block, and slow consumers decouple from
  the lockstep loop.
- The engine holds the model's inference_lock while it has active
  slots, so vision and non-streaming requests (which keep the direct
  path) still serialize safely against the batch; the lock releases
  when the batch drains.
- Fatal worker errors fail all active slots, mark the model poisoned
  on device-fault classification, and stop the engine (later submits
  fail fast).
- LoadedModel: poisoned/inference_lock/prefix_cache/prefill_rate move
  behind Arc (shared with the engine without keeping the model alive);
  new engine: Option<EngineHandle> field. Engine activates only when
  max_in_flight > 1 (config default 1 = existing behaviour everywhere)
  and NEURON_BATCHING != 0 (kill switch).
- Prefill helpers (restore_or_clear/chunked_prefill/store_prefix/
  emit_delta) ungated from cfg(cuda) so the engine and its CPU tests
  compile everywhere.

Gold test: three greedy requests submitted concurrently (ragged
prompts, mid-flight joins, rebatching) produce byte-identical text,
token counts, and finish reasons to the same requests run
sequentially through the same engine, on the tiny fixture.

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

fix(neuron): as_deref for Arc'd prefix_cache in cuda-only worker path (#98) ... d3ef9513f8

candle.rs:2441 sits behind cfg(cuda) — the CPU build can't see it, and
the CUDA type-check caught the Option<&Arc<Mutex<…>>> vs
Option<&Mutex<…>> mismatch introduced by the LoadedModel Arc change.

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

grenade merged commit 7454ec7744 into main 2026-07-03 12:39:08 +00:00

grenade referenced this issue from a commit 2026-07-03 12:39:09 +00:00

Merge pull request 'feat(neuron): continuous batching — batched decode + engine (#98, slices 1–3b)' (#121) from feat/98-batched-decode into main

grenade referenced this pull request 2026-07-03 12:39:35 +00:00

F4d — perf(neuron): continuous batching for agentic fan-out #98

grenade referenced this pull request 2026-07-03 14:20:21 +00:00

feat(neuron): continuous batching — TP batched decode + engine backend seam (#98) #122

Sign in to join this conversation.

Reviewers

No Reviewers

Labels

Clear labels

out-of-scope

Outside helexa's constraining principles: near-frontier models on consumer hardware, served transparently and predictably

p1-now

Active pain or unblocking work — pick up first

p2-next

Engine levers — start once P1 baselines exist

p3-later

Dependent on earlier work, or demand-driven

No Label

Milestone

No items

No Milestone

Projects

Clear projects

No project

Assignees

Clear assignees

grenade

No Assignees

1 Participants

Notifications

Subscribe

Due Date

No due date set.

Dependencies

No dependencies set.

Reference: helexa/helexa#121