Skip to content

Server: fused single-GPU decode for --batched-session (co-decode 2-4 sessions in one graph pass)#604

Open
iCreil wants to merge 10 commits into
antirez:mainfrom
iCreil:feat/rebase-on-batched-upstream
Open

Server: fused single-GPU decode for --batched-session (co-decode 2-4 sessions in one graph pass)#604
iCreil wants to merge 10 commits into
antirez:mainfrom
iCreil:feat/rebase-on-batched-upstream

Conversation

@iCreil

@iCreil iCreil commented Jul 25, 2026

Copy link
Copy Markdown

--batched-session N works beautifully on multi-GPU placements, where the
session pipeline fills different pipeline stages across devices. On a single
GPU I ran into two things while testing it on an RTX PRO 6000 (96 GB) with
the Flash q2 GGUF, and this PR is an attempt at addressing both. Happy to
rework any part of it — the branch is organized so the pieces can be taken
(or dropped) independently.

1. Resident session creation fails at normal context sizes

On 96 GB, --batched-session 4 --ctx 32768 (and even --batched-session 1 --ctx 32768) fails with failed to create cuda session: the lazy Q8→F16
weight caches warm up during engine open and grow until only the default
reserve (~1% of VRAM) stays free, and the resident sessions are created
after that point. Measured on q2: engine base lands at ~95.9 GiB while each
32k session only needs ~0.6 GiB of context buffers.

Fix (2 commits): the server passes the resident session count to the engine
(resident_sessions config field), and engine open computes a cache reserve
from the per-session ds4_context_memory estimate right before
accelerator_cache_model_tensors() — the estimate needs the model shape, so
it has to run inside open, not before. An explicit
DS4_CUDA_Q8_F16_CACHE_RESERVE_MB still wins. With this, --batched-session 4 --ctx 32768 --prefill-chunk 1024 boots at 97.2/97.9 GiB and single-stream
speed is unchanged (the hot part of the cache still fits).

2. Single-GPU decode batches don't scale

Without a placement, ds4_sessions_eval_batch_cuda encodes the one-token
graphs back to back. That preserves exactness, but each session re-reads the
full weight stream, so decode is memory-bound at 1x: measured 62 tok/s
aggregate at 2 and 4 concurrent clients vs 61 tok/s for one client.

This PR adds a fused single-tier path to the same function (the single-GPU
counterpart of metal_graph_encode_session_pipeline_batch): the dense
projections run once on the shared batch workspace with n_tokens = count,
so the weights cross memory once per step; KV append and compressor/indexer
bookkeeping run per sequence with the same single-token helpers the decode
layer uses; attention runs as one multi-sequence launch with per-sequence
fallbacks (indexed path, score-buffer overflow), so there is no position
limit.

Numbers (RTX PRO 6000, Flash q2, --batched-session 4 --ctx 32768 --prefill-chunk 1024, 200 generated tokens per request, aggregate =
total tokens / wall):

concurrent clients main (one-token pipelining) this PR (fused)
1 61 61 (same path)
2 62 68
4 could not boot / 62 at smaller ctx 85.6

Commit walkthrough

CUDA kernels (additive):

  • CUDA: add multi-sequence decode attention kernel — one query token per
    sequence, each attending to its own KV ring/compressed cache
    (ds4_gpu_attn_seqview, by-value in kernel params, max 4).
  • CUDA: add pos-vector RoPE tail variants — decode tokens of different
    sequences sit at arbitrary positions; per-element math identical to the
    single-token kernels (token t matches a pos0=pos[t] launch bitwise).
  • CUDA: narrow-batch q8_0 matmul for 2-4 token decode — same idea as your
    tok2 kernel, generalized to n_tok 2-4 with n_tok as a parameter: one
    warp per output row, each 32-weight block read once and dotted against
    all tokens. Dispatched ahead of cuBLAS/MMA for n_tok 2-4
    (DS4_CUDA_NO_Q8_NARROW=1 restores the previous dispatch; the GLM verify
    force_decode_warp path is untouched).

Graph driver:

  • metal_graph_decode_compressor_indexer — the per-token compressor/indexer
    segment of the decode layer phase, taken as a helper so the batched driver
    can run it per sequence on batch-scratch row views. It is a copy (the
    single-token path is untouched); if you'd rather have the phase function
    call the helper too to avoid the duplication, glad to do that instead.
  • metal_graph_encode_layer_attention_decode_multi +
    metal_graph_encode_decode_multi — the fused pass. FFN and output head
    reuse the existing batch functions; logits land one row per sequence in
    spec_logits.
  • Graph: run fused-decode routed experts per sequence — at batch 2-4 the
    selected experts of different sequences are almost always disjoint, so the
    batch expert path costs more than n single-token passes. This turned out
    to matter twice: without it the fused path was slower than pipelining (54
    vs 62) and one of two concurrent greedy streams diverged from the solo
    run; with it, both concurrent greedy streams match the solo stream
    token-for-token on the prompts I tested. DS4_FUSED_BATCH_MOE=1 keeps the
    batch expert path for comparison.

Server/glue:

  • The fused path plugs into ds4_sessions_eval_batch_cuda for single-tier,
    non-streaming, non-distributed decode batches of 2-4;
    DS4_NO_FUSED_SESSION_BATCH=1 restores the one-token pipelining (and was
    used for the A/B numbers above).

Correctness notes and limitations

  • The fused pass is not bitwise-identical to serial decode (batched matmul
    reduction order); greedy token streams matched the solo run in every check
    I ran (q2, 150-200 tokens, 2 and 4 concurrent). The kill switch makes the
    comparison easy to reproduce, and I kept it default-on only because the
    numbers looked convincing — happy to flip the default if you prefer
    opt-in.
  • CUDA only: the multi-sequence kernels have no Metal twins, so Metal keeps
    its current paths untouched.
  • SSD streaming sessions are excluded from the fused path (they fall back to
    the existing behavior).
  • spec_logits is now allocated unless the fused path is disabled (~8 MiB
    per graph), since the fused output head lands there.

Test setup details, exact configs and raw runs are in the commit messages;
I can share the bench harness if useful. And if you had a different
direction in mind for single-GPU batching, happy to adapt this or just take
whatever pieces are useful.

iCreil added 10 commits July 24, 2026 22:12
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.
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.
…ode driver)

Port of the fork's batched-decode driver onto the current main:
- metal_graph_decode_compressor_indexer: per-sequence compressor/indexer
  helper re-extracted from the current decode layer phase (kept as a copy,
  single-token path untouched).
- metal_graph_encode_layer_attention_decode_multi + encode_decode_multi:
  dense projections batched on the aliased batch workspace with
  n_tokens = n_seqs (narrow q8_0 kernel serves 2-4 rows), pos-vector RoPE,
  per-sequence KV store and compressor/indexer, multi-sequence attention
  with per-sequence fallbacks, FFN and output head via the existing batch
  functions.
- ds4_sessions_eval_batch_cuda: single-tier decode batches take the fused
  pass instead of one-token pipelining (DS4_NO_FUSED_SESSION_BATCH=1
  restores it); logits land one row per sequence in spec_logits, now
  allocated unless the fused path is disabled.
In the fused session-batch decode each FFN batch row is one token of a
distinct sequence, so the selected experts are almost always disjoint and
the generic batch expert path costs more than n_rows single-token passes.
Route the routed-MoE dispatch of metal_graph_encode_layer_ffn_batch through
ds4_gpu_routed_moe_one_tensor per row while the fused driver is on the
stack (file-scope flag; the graph worker serializes encodes).
DS4_FUSED_BATCH_MOE=1 keeps the batch expert path for comparison.
With --batched-session N the lazy Q8->F16 weight caches warm up during
engine creation and grow until only the default reserve (~1% of VRAM)
stays free, so the resident sessions allocated afterwards fail even at
modest context sizes (on a 96 GB card even one 32k session could not be
created).  Derive the reserve from the per-session context-buffer
estimate times N plus the shared prefill workspace, and export it before
engine creation; an explicit DS4_CUDA_Q8_F16_CACHE_RESERVE_MB wins.
Move the reserve computation into engine open, right before the startup
cache warm: the context-memory estimate needs the model shape, which is
only loaded here (the earlier server-side attempt underestimated it and
still failed session creation).  The server passes the resident session
count via the new resident_sessions config field.
iSevenDays added a commit to iSevenDays/ds4 that referenced this pull request Jul 25, 2026
…tirez#605 race fix

Cherry-picked onto feat/mtp-native (which has upstream main + MTP/DSpark + expert LRU):
- antirez#604: 10 commits (session-batch decode, fused routed experts, VRAM reserve for resident sessions)
- antirez#605: 1 commit (batched server session recovery race fix)
- antirez#605 streaming improvements (4 commits): deferred — conflict with per-tier LRU, adapt later

MTP/DSpark: present (from upstream fc9efd1 via feat/mtp-native).
--batched-session --prefill-chunk 1024 now available (antirez#604).
iSevenDays added a commit to iSevenDays/ds4 that referenced this pull request Jul 25, 2026
…r LRU

The integration commit 8577818 cherry-picked PR antirez#604 (session batching +
VRAM reserve) onto feat/mtp-native, but the cherry-pick left 33 unresolved
git conflict markers throughout ds4_cuda.cu, plus a duplicate definitions
block. Each conflict was HEAD=per-tier LRU (our f836788 port, uses ssd_current()
and g_ssd[t].expert_cache) vs 10ba298=upstream per-class (uses
g_stream_expert_caches[DS4_CUDA_STREAM_EXPERT_CLASSES]).

Resolution:
- Resolved all 33 conflict regions in favor of HEAD (per-tier), dropping
  upstream's per-class infrastructure (g_stream_expert_caches[],
  g_stream_expert_runtime_caps[], g_stream_expert_class_*_bytes[],
  cuda_stream_expert_cache_class_index/_configured_budget_class/
  _release_class, DS4_CUDA_STREAM_EXPERT_CLASSES define).
- Removed the duplicate definitions block (the second enum, duplicate
  struct cuda_stream_expert_cache_slot/_cache, and the second
  g_stream_expert_budget_override) introduced by the cherry-pick; the
  original per-tier definitions at the top of the file are authoritative.
- Dropped the leftover upstream direct-load block in
  cuda_stream_selected_cache_begin_load that referenced the undefined
  g_stream_selected_cache global and direct_loads/cache_hits/cache_misses
  counters (HEAD's per-tier flow already handles this via sc->gate_ptr and
  copied_from_global_cache).
- Removed the duplicate no-op stub of cuda_model_load_progress_finish;
  the real implementation at line 1550 stands.
- Kept cuda_stream_layer_expert_ranges_valid (pure bounds-check helper,
  no per-tier/per-class state) because it is still called from
  ds4_gpu_stream_expert_cache_seed_experts outside any conflict region.

Build passes clean with CUDA_ARCH=sm_89. The resulting binary exposes
all the expected flags: --ssd-streaming-cache-experts (LRU),
--batched-session (PR antirez#604), --mtp*/--dspark* (MTP/DSpark).

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant