Batched decode: co-decode N live sessions in one GPU pass (opt-in) by iCreil · Pull Request #530 · antirez/ds4

iCreil · 2026-07-10T14:15:05Z

Draft, stacked on #516 (which is stacked on #515). Opened as draft so the
branch is visible while #515/#516 await review — please look at those first:
this branch builds on the --parallel N interleaved-session server from #516.
The diff below shows those commits too until #516 lands; once it merges I'll
rebase and only the 14 batched-decode commits remain, then mark this ready.

What

With --parallel N (from #516) the server round-robins several live sessions
through the single graph worker, one token at a time (interleaving). This PR adds
an opt-in batched decode: when two or more slots are simultaneously in plain
decode, their one-token-per-session steps are committed together in a single
ds4_session_eval_multi call, so the memory-bound dense stages (attention
projections, shared expert, output head, RoPE/norm) read each weight block from
VRAM once and reuse it across all the sessions' tokens, instead of re-reading
it per session.

Default behaviour is unchanged: batching is gated on DS4_SERVER_BATCHED_DECODE=1,
requires --parallel > 1 and an engine that reports
ds4_engine_supports_batched_decode() (CUDA today). With the flag off the server
uses the existing interleaving path bit-for-bit.

Why it's worth it — the narrow-batch matmul

The dense stages already had "batch" kernels (used for prefill, n_tokens
32–2048). Reusing them for n_tokens 2–4 was actually slower than
interleaving (B=2 ≈ 0.62×): the batch-warp q8 matmul puts n_tok on the grid Y
axis, so every token re-reads the weight rows from VRAM — no bandwidth saving in a
memory-bound decode.

The key commit here is a narrow-batch q8_0 matmul kernel
(matmul_q8_0_preq_narrow_warp8): one warp per output row, it reads each 32-weight
block once into registers and dots it against all n_tok activation blocks
(weights stay hot in L1, activations are tiny). That is exactly the reuse the
memory-bound decode wants.

Measured (RTX PRO 6000 Blackwell, DeepSeek V4 Flash q2, isolated decode step)

ms/step is the time to advance all the batched sessions by one token, so
the throughput comparison is ms/step against n_tok × single-token time:

batch	ms/step	speedup vs interleaving
1 (single)	21.1 (1 tok)	1.00×
2	28.7 (2 tok)	1.47×
3	—	1.67×
4	—	1.83×

(narrow kernel; the prior batch-warp path gave B=2 ≈ 67.5 ms/step = 0.62×.)
End-to-end (with prefill, two concurrent chats, temp 0): ~1.30× aggregate.

Design note — dense stages batched, routed experts per sequence

The HC pre/norm, router, shared expert and output head run batched
(n_tokens = n_seqs): same weights, every sequence. The routed experts run
one sequence at a time through the existing single-token decode MoE kernel on
batch-scratch row views: at batch 2-4 the selected experts of different
sequences are almost always distinct, so a generic batched expert path reads
more expert weights than n_seqs single passes, not fewer. This was measured,
not assumed — an earlier revision routing the MoE through the prefill expert
path was 0.62-0.80× (slower than interleaving); the per-sequence expert path
plus the narrow matmul is what flips it to 1.47×.

Correctness

Batched decode is not bitwise-identical to single decode by design: the MoE
down-projection sums the selected experts with atomicAdd (reduction order
depends on occupancy), and the hyper-connection fusion / matmul reductions differ
between the single and batch paths. It is token-identical in practice: greedy
runs match single decode across the test prompts. tests/batched_decode_smoke
checks the narrow path against the already-validated batch-warp path (first-step
logits bitwise-identical, token streams match), which is the right invariant.

Scheduler & observability

The scheduler batches only slots in plain decode. A slot whose step wants the
MTP speculative path (tool-call payloads force greedy mid-stream, so eligibility
flips per token) advances inline with its own burst; batched decode and
MTP-greedy are mutually exclusive per step.
DS4_SERVER_BATCH_DIAG=1 reports the batched fraction and a per-cause breakdown
of the lone-decode turns (peer_prefill / peer_finish / no_peer). Useful
finding: on sporadic interactive chat the batched fraction is arrival-bound,
not scheduler-bound — ~90–95% when two streams co-arrive, ~20–30% when arrivals
are staggered or lengths differ (one stream decodes before the other arrives or
after it finishes: no_peer). Batching pays most under sustained, synchronized
concurrency; it never hurts the interleaving path since it is opt-in.
On realistic mixed chat traffic ~47% of decode turns actually co-decode
(staggered arrivals, different lengths), rising to ~62% for near-identical
concurrent requests — so the end-to-end gain scales with real overlap, and
everything else falls back to the Server: add --parallel N interleaved sessions (default 1) #516 interleaving behaviour.

Known limitation

The multi-sequence decode attention kernel currently handles the raw-SWA and
f32-compressed paths; the compressed-f16 / large-context online variant
(where n_comp exceeds the score buffer, roughly past ~31k tokens) is not
implemented for the multi path yet. ds4_gpu_attention_decode_multi_supported()
gates on this, so those turns cleanly fall back to per-slot single decode —
correct, just without the batch speedup. Adding the f16/online multi attention
kernel is the natural follow-up.

--ssd-streaming is excluded by the ds4_engine_supports_batched_decode gate.
That combination is actually where batching should pay the most (amortizing the
expert-miss H2D traffic across sequences), but it needs the streaming cache to
accept a per-batch union of selected experts on the decode path — it deserves
its own PR.

Commit walk-through (14 commits on top of #516)

CUDA: add multi-sequence decode attention kernel — per-seq KV via seqview.
CUDA: add pos-vector RoPE tail variants — B sessions have arbitrary positions.
Refactor: extract per-token compressor/indexer step — reused per-seq (no behaviour change).
Graph: add batched-decode encoder (one token per sequence) — dense stages batched, per-seq KV/attention.
Session: add ds4_session_eval_multi — public API + ds4_engine_supports_batched_decode.
Server: batched decode turns — opt-in scheduler batching.
Graph: run batched-decode routed experts per sequence — MoE via decode selected-expert path.
tests: add batched decode step microbench.
CUDA: narrow-batch q8_0 matmul for 2-4 token decode — the performance kernel above.
Server: DS4_SERVER_BATCH_DIAG observability counters.
Server: attribute single-decode turns by cause in batch-diag.
test(server): cross-slot KV frontier guard + disconnect harness — unit test for the
per-slot KV frontier invariant (extracted slot_kv_swap_in/out helpers, no behaviour
change) + tests/disconnect_mid_decode_test.py integration harness.
fix(server-test): route per-session live state through a slot — repairs the ds4_test
server suite, which no longer compiled after the --parallel refactor moved
responses_live/session into server_slot (pre-existing on Server: add --parallel N interleaved sessions (default 1) #516; happy to fold this
into Server: add --parallel N interleaved sessions (default 1) #516 instead if you prefer).
test(disconnect): count reasoning_content as decode output — the harness must count
thinking tokens too, or the mid-decode abort never triggers on thinking models.

Flipping the default

With the per-sequence expert path and the narrow matmul in place, the batched
path dominates interleaving at every measured batch size. A reasonable adoption
path: keep it opt-in for one release, then default it on with
DS4_SERVER_BATCHED_DECODE=0 as the escape hatch. With the flag off (and with
--parallel 1) behaviour stays byte-identical to #516.

Note for maintainers

The *_multi_tensor kernels are implemented in ds4_cuda.cu only. The Metal
backend will need equivalents (or stubs) for the multi-sequence attention,
pos-vector RoPE and ds4_session_eval_multi to build on macOS — happy to add
stubs if you point me at the pattern you prefer.

…prefills Mixed-quant GGUFs (e.g. the official Flash q2-q4-imatrix) interleave layers whose routed experts have different byte sizes. The CUDA expert cache was a single global slab keyed on the last-seen size: every q2<->q4 layer transition reset the runtime cap and released the whole cache, so decode on mixed models was stuck at ~1 t/s with a permanent cap-notice flood. Additionally, prefill batch loads released the entire resident expert cache unconditionally before allocating their staging buffers, forcing a full cache re-warm on every request (this also affects uniform models). Changes: - keep up to 4 expert caches, one per (gate,down) byte-size class; runtime caps and cap notices are tracked per class - begin_compact_load: try the staging allocation first and only sacrifice the expert cache when VRAM is actually short (retry once) - let the prefill batch path read the global cache (hits are device-to- device copies instead of host uploads), appending on miss only while free capacity remains, never evicting: a long prompt cannot cycle out the decode working set On an RTX PRO 6000 Blackwell 96GB with the official Flash q2-q4-imatrix GGUF, a 200-word completion goes from not finishing within 400s (~0.9 t/s) to ~25 t/s sustained decode with --ssd-streaming-cache-experts 9472 and DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB=12 (97% cache hit rate per the VERBOSE stats). Uniform-quant GGUFs keep their existing behaviour.

Split ds4_gpu_graph into per-session state (KV caches, compressor and indexer frontiers, MTP verification state) and a new engine-owned ds4_gpu_scratch holding the transient work buffers: one-token decode tensors, per-layer stage tensors, MTP draft tensors and the batched prefill buffers. Everything in the scratch is written and consumed inside a single sync/eval call, so sessions sharing an engine can share one scratch as long as they are driven from one thread at a time -- which is exactly the serialization the server's single graph worker already provides. This is what makes a second live session affordable: the prefill batch buffers dominate a session's non-KV footprint (~1.8 MiB per prefill-chunk token), and they are now allocated once per engine instead of once per session. The scratch grows on demand to the union of the attached sessions' capacities and is freed with the engine. No behavior change with a single session: tensors, sizes and the execution order are unchanged, only their ownership moved. Add tests/two_sessions_smoke: drives two sessions of one engine interleaved (one greedy token each, alternating on one thread) and checks both streams match isolated single-session runs token by token. Verified on CUDA (RTX PRO 6000, DeepSeek V4 Flash q2): MATCH for both sessions, and isolated outputs identical to the pre-refactor build.

Move the request-scoped locals of generate_job() into a gen_state struct and split the body into job_begin / job_prefill / job_start_decode / job_decode_round_init / job_decode_step / job_finish. The decode_again tool-recovery label becomes a GEN_STEP_REDECODE result from job_finish. generate_job() is now a small driver that performs exactly the same calls in exactly the same order, so behavior is unchanged; this is groundwork for a scheduler that drives several jobs stepwise (one prefill chunk or a few decode tokens at a time) from the single graph worker instead of owning it until a job completes. Verified against the previous build on CUDA (DeepSeek V4 Flash q2): greedy chat completion is byte-identical in both stream and non-stream mode, and the SSE frame flow is intact.

Replace the single live session with an array of slots, each owning one ds4_session plus the protocol live state (Responses/Anthropic/thinking bindings) tied to that timeline. The graph worker becomes a scheduler: it hands queued jobs to idle slots (longest-common-prefix affinity, then LRU) and steps one slot at a time through the gen_state phases -- a whole begin/finish transition, one prefill slice, or a burst of decode tokens per turn. Prefill preemption reuses the cooperative cancel callback: syncs are interrupted at chunk boundaries once the slice deadline passes and another job is waiting, and resume from the session checkpoint on the slot's next turn. Decode interleaves round-robin between slots, so every connected client keeps streaming while another prompt prefills. The scheduler swaps the kvstore continued-checkpoint frontier in and out per slot; parser-side live call-id lookups scan every slot under the existing tool mutex. --parallel stays opt-in: the default is one slot, the yield callback never fires (no contention), and the phase sequence is exactly the old generate_job order -- verified byte-identical against the previous build on a greedy chat completion (stream and non-stream). N>1 requires resident experts and is refused with --ssd-streaming. Timeslices are tunable via DS4_SERVER_SCHED_PREFILL_MS (default 1500) and DS4_SERVER_SCHED_DECODE_TOKENS (default 6). Measured on CUDA (RTX PRO 6000 96GB, DeepSeek V4 Flash q2, --parallel 2 --ctx 32768 --prefill-chunk 1024, 94.9 GiB VRAM): with a 350-token generation in flight, a second client got its first token in 0.7 s and the first stream's worst inter-token gap during the newcomer's prefill was 0.15 s; a third concurrent request queued and completed once a slot freed.

Add ds4_gpu_attention_decode_heads_multi_tensor(): one query token per sequence, each attending to its own raw + compressed caches with decode (everything visible) semantics. Sequence b reads q row b and writes heads row b; per-sequence cache pointers, ring geometry and optional indexer mask row travel in a ds4_gpu_attn_seqview array (max 4, passed by value in the kernel parameters). The kernel is a standalone adaptation of the n_tokens==1 launch of attention_decode_mixed_kernel, so the single-session hot path is untouched. Sequences whose compressed row count exceeds the shared score buffer are rejected for now: the large-context online variant is a follow-up, and callers gate batched decode on this. tests/attention_multi_smoke drives two sequences with different shapes, ring offsets and masks and checks the multi launch against two independent single-sequence launches: bitwise identical on both the head_dim==512 fast path and the generic path (RTX PRO 6000, CUDA 13).

Batched decode evaluates one token per sequence, each at its own arbitrary position, while rope_tail_kernel and head_rms_norm_rope_tail_kernel derive the position as pos0 + t (* stride). Add _multi twins that take a small by-value position array (max DS4_GPU_DECODE_MULTI_MAX entries) indexed by token row, with per-element arithmetic kept identical so each row matches a single-token launch bitwise. tests/rope_multi_bench checks that equivalence (plain and YaRN paths) and measures why the variants exist: on an RTX PRO 6000, per-sequence loops of single-token launches over the 4 decode RoPE sites cost ~1.0 ms extra per 43-layer step at B=4 (~516 launches x ~2 us) versus batched launches.

…ayer Pure extraction, no behavior change: the compressed-KV update/emit, the ratio-4 indexer state update and the indexer score/top-k selection move from metal_graph_encode_decode_layer into metal_graph_decode_compressor_indexer(). The helper takes the token's attn_norm/qr_norm tensors and position explicitly so batched decode can call it once per sequence on batch-scratch row views.

metal_graph_encode_decode_multi evaluates one decode token for each of up to DS4_GPU_DECODE_MULTI_MAX sequences in a single pass over the layers. Dense projections run once on the batch scratch with n_tokens = n_seqs, reusing the prefill FFN and output-head batch functions unchanged (token ids go through scr->prefill_tokens for hash routing; logits land one row per sequence in scr->spec_logits). RoPE uses the pos-vector kernels, and KV append plus compressor/indexer bookkeeping run per sequence with the same single-token helpers the decode layer uses, so per-sequence state advances exactly as independent evals would. Attention runs as one multi-sequence launch when every sequence is served by the regular decode kernel; sequences on the indexed path enqueue their attention inside the per-sequence loop (scr->comp_selected is shared and must be consumed before the next sequence overwrites it), and anything the multi kernel cannot serve (compressed rows beyond the score buffer, F16 compressed cache) falls back to the existing per-sequence kernels, so there is no position limit. Additions only; nothing calls this yet — the session-level API and the server scheduler wire it up next.

One token for each of up to DS4_SESSION_EVAL_MULTI_MAX distinct GPU sessions of the same engine, in a single command buffer via metal_graph_eval_decode_multi. Positions come from each session's checkpoint, logits land in each session's own buffer, and checkpoints advance exactly as with independent evals; pending MTP drafts are dropped since they belong to the single-token flow. n_sessions == 1 degrades to ds4_session_eval. ds4_engine_supports_batched_decode gates the feature: GPU backend, no ssd-streaming, no directional steering.

Split job_decode_step into job_decode_sample and job_decode_post, then add a batched turn to the scheduler: when the picked slot is decoding, every slot in plain decode advances one token per round through ds4_session_eval_multi for the usual per-turn budget. A slot whose step wants the MTP speculative path advances inline with its own burst (tool-call payloads force greedy sampling mid-stream, so eligibility can flip per token); a slot that stops moves to SLOT_FINISH and completes on its next regular turn. spec_logits is now allocated unconditionally: it used to exist only with MTP, and batched decode reads its per-sequence logits rows from it. DS4_BATCHED_DECODE_FORCE keeps n==1 on the batched path for testing. tests/batched_decode_smoke drives two prompts through single evals and through eval_multi (n==1 forced and n==2) and compares greedy streams and logits: on CUDA (RTX PRO 6000, V4 Flash q2) all 24-token streams are token-identical. Off by default for now: the batched MoE stage goes through the prefill expert path, which is slower than interleaving at batch 2-4 (the decode selected-expert kernels are not used). End-to-end wall time for two concurrent 200-token generations was 0.72x sequential; flip the default once the routed experts run per-sequence through the decode path.

metal_graph_encode_layer_ffn_decode_multi replaces the prefill FFN batch in the batched-decode driver. The dense stages (HC pre/norm, router, shared expert, HC post) stay batched with n_tokens = n_seqs, where the same weights serve every sequence. The routed experts now run one sequence at a time through the single-token decode MoE kernel (ds4_gpu_routed_moe_one_tensor) on batch-scratch row views: at batch 2-4 the selected experts of different sequences are almost always distinct, so the generic batch expert path reads them all without sharing, costing more than n_seqs single-token passes. Token streams stay identical to single decode (batched_decode_smoke MATCH).

Isolates decode-step cost (prefill and contention excluded) for single vs batched B=2..4, reporting aggregate tok/s and the ratio against serving the same streams one-token-at-a-time on the shared GPU. On an RTX PRO 6000 with V4 Flash q2: single 21.1 ms/token; batched B=2 67.5 ms/step (0.62x), B=3 70.4 ms (0.90x), B=4 75.2 ms (1.12x). The near flat B=2..4 step time shows a large fixed per-layer cost in the batch dense kernels at tiny n_tokens, which only amortizes past B=4; at B=2 interleaved single-token decode wins.

The batch warp kernel puts n_tok on the grid Y axis, so every token re-reads the same weight row from VRAM; in decode, where weights are the whole cost, that scales the time with n_tok for no benefit. The generic per-row kernel (used when blocks>32) is worse still. The new narrow kernel keeps one warp per output row and, for each 32-weight block it reads once, dots it against all n_tok activation blocks, so the weight bytes cross memory once and are reused from L1. Output layout matches the batch kernel; the n_tok==1 and prefill paths are untouched. Dispatched for n_tok in [2, DS4_CUDA_NARROW_MAX] ahead of cuBLAS and the batch warp kernel; DS4_CUDA_NO_Q8_NARROW forces the old path. Modeled on the n_tok==1 warp8 kernel, it loops blocks with a lane stride and handles any in_dim. Decode-step bench (RTX PRO 6000, V4 Flash q2), aggregate vs interleaving: B=2 1.47x, B=3 1.67x, B=4 1.83x (was B=2 0.62x through the batch path). Correctness: batched_decode_smoke shows narrow and batch-warp produce identical token streams. Batched decode still differs slightly from single decode, but that is the existing non-deterministic MoE reduction (atomicAdd over experts), not the kernel: verified by first-step logits being bit-identical between narrow and batch-warp.

Optional (env-gated) counters that log, on each request completion, how many decode turns were batched vs run single-slot and the batched fraction. Used to measure that batched decode only fires while two sessions genuinely co-decode: ~47% of turns for mismatched concurrent requests, ~62% for identical ones — the rest run single because the two streams overlap only partially (staggered arrival, different lengths). No effect unless DS4_SERVER_BATCH_DIAG is set.

Extend the DS4_SERVER_BATCH_DIAG counters so a lone decode turn (one slot in SLOT_DECODE while no peer is also decoding, so no batch forms) is attributed to why it could not batch: - peer_prefill: a peer is still in begin/prefill/start-decode and will join decode shortly (recoverable overlap the scheduler could in principle win); - peer_finish: a peer is past decode (finishing); - no_peer: there is genuinely no other job (asynchronous arrival, or the only other stream already finished). This turns "batched fraction is low" into a measurable cause. On an interactive chat workload (parallel 4, ctx 32768, prod-realistic ~1700-token prompts) the lone-decode turns are overwhelmingly no_peer: ~95% batched when two streams co-arrive, but ~20-30% when arrivals are staggered or lengths differ, because one stream decodes before the other has arrived or after it has finished. peer_prefill stays at 0-2 turns even with long prefills, so the batched-fraction ceiling on sporadic chat is set by the arrival pattern, not by scheduler policy. Pure observability; gated on DS4_SERVER_BATCH_DIAG, no effect on the hot path.

- extract the per-slot continued-frontier swap into slot_kv_swap_in/out (identical assignments, golden-preserving) and add test_kv_cache_continued_frontier_is_per_slot to guard a slot from ever observing another slot's frontier under --parallel. - tests/disconnect_mid_decode_test.py: integration harness proving a slot is freed on client disconnect mid-decode and peers stay unaffected. - document the batched-decode f16-comp attention gate as a known limitation.

The --parallel refactor moved responses_live/anthropic_live/session from server to server_slot, but the tool-continuation and think-tool-recovery unit tests still set them on server directly, so ds4_test stopped compiling (latent: the suite was not rebuilt since the refactor). Wire a one-slot server (slots/n_slots/active) and set the live state on the slot, matching how the client-thread validators scan slots.

ds4/DeepSeek default to thinking, so a turn streams reasoning_content before any content. The harness only counted content deltas, so the abort never triggered (it drained the whole stream) and the serve-checks saw an empty but actually-served response. Count both content and reasoning as decode progress; bump default rounds past --parallel so a slot leak still exhausts.

iCreil · 2026-07-10T14:46:19Z

You mentioned session batching in your latest video — this stack is a working CUDA implementation of exactly that: opt-in batched decode across --parallel slots, 1.47–1.83× aggregate over interleaving at B=2–4, greedy token streams identical to single decode (numbers and design notes in the description).

It's marked draft only because it stacks on #515/#516 — happy to adapt it to whatever direction you had in mind.

…cache' into feat/batched-decode

iCreil · 2026-07-25T09:08:10Z

Superseded by #604, closing.

iCreil added 18 commits July 5, 2026 16:31

Merge remote-tracking branch 'origin/fix/cuda-streaming-mixed-expert-…

81e9c4f

…cache' into feat/batched-decode

iCreil closed this Jul 25, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Batched decode: co-decode N live sessions in one GPU pass (opt-in)#530

Batched decode: co-decode N live sessions in one GPU pass (opt-in)#530
iCreil wants to merge 19 commits into
antirez:mainfrom
iCreil:feat/batched-decode

iCreil commented Jul 10, 2026 •

edited

Loading

Uh oh!

iCreil commented Jul 10, 2026

Uh oh!

iCreil commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

iCreil commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What

Why it's worth it — the narrow-batch matmul

Measured (RTX PRO 6000 Blackwell, DeepSeek V4 Flash q2, isolated decode step)

Design note — dense stages batched, routed experts per sequence

Correctness

Scheduler & observability

Known limitation

Commit walk-through (14 commits on top of #516)

Flipping the default

Note for maintainers

Uh oh!

iCreil commented Jul 10, 2026

Uh oh!

iCreil commented Jul 25, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

iCreil commented Jul 10, 2026 •

edited

Loading