Static batched decode (LLaMA + Qwen3): B independent sequences per step, up to 41x aggregate throughput#129
Conversation
…r-slot KV, 41x aggregate throughput
There was a problem hiding this comment.
Pull request overview
Implements a new static batched decode path for FP16 LLaMA that advances B independent sequences per step, using MMA/GEMM-heavy projections to amortize weight reads and increase aggregate decode throughput. This adds a new decode-layer task-graph pipeline, decode-specific KV-addressing kernels, and benchmark/doc support to reproduce and validate performance/correctness.
Changes:
- Add
LlamaFP16LayersBatchDecodeMMAto build per-layer TornadoVM TaskGraphs for batched decode with per-slot KV and per-slot positions. - Add decode-specific kernels in
TransformerBatchPrefillKernelsfor RoPE+KV write and FP16 flash attention over per-slot KV regions. - Add end-to-end and micro-benchmarks plus a detailed write-up (
BATCHED_DECODE.md) for validation and reproduction.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 7 comments.
Show a summary per file
| File | Description |
|---|---|
| src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/decode/LlamaFP16LayersBatchDecodeMMA.java | New per-layer TaskGraph builder for static batched decode with per-slot KV addressing. |
| src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerBatchPrefillKernels.java | Adds decode variants of RoPE+KV write and flash-attention kernels with per-slot KV regions/positions. |
| src/main/java/org/beehive/gpullama3/bench/BatchedProjectionBench.java | New synthetic benchmark illustrating projection batching benefits. |
| src/main/java/org/beehive/gpullama3/bench/BatchedDecodeEngine.java | New end-to-end batched decode benchmark/engine wiring activation + decode layers + logits. |
| src/main/java/org/beehive/gpullama3/bench/BatchedDecodeAttentionBench.java | New microbench validating/benchmarking decode attention with per-slot KV caches. |
| BATCHED_DECODE.md | Documentation/write-up describing design, correctness checks, and performance results. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, | ||
| context, state.wrapHbFP16Batch, | ||
| weights.w2Layered[layerIndex].asHalfFloatArray(), | ||
| state.w2Out, batchSize, dim, hidDim) |
| int B = args.length > 0 ? Integer.parseInt(args[0]) : 32; | ||
| int K = args.length > 1 ? Integer.parseInt(args[1]) : 2048; | ||
| int N = args.length > 2 ? Integer.parseInt(args[2]) : 2048; | ||
| int iters = 100; |
| int B = args.length > 0 ? Integer.parseInt(args[0]) : 32; | ||
| int seqLen = args.length > 1 ? Integer.parseInt(args[1]) : 256; // each slot attends 0..seqLen-1 | ||
| int iterations = 100; |
| if (P + nDecode >= decodeCtx) { | ||
| throw new IllegalStateException("prompt(" + P + ")+n(" + nDecode + ") >= ctx(" + decodeCtx + ")"); | ||
| } |
| long kvElems = (long) B * numLayers * decodeCtx * kvDim; | ||
| FloatArray keyCacheBatch = new FloatArray((int) kvElems); | ||
| FloatArray valueCacheBatch = new FloatArray((int) kvElems); |
| int nDecode = Integer.getInteger("batch.decode.n", 64); | ||
| boolean cudaGraphs = Boolean.parseBoolean(System.getProperty("batch.decode.cudaGraphs", "true")); |
| // ── Decode: static batch, greedy. Position advances in lock-step. | ||
| // First few steps capture the CUDA graphs; exclude them from timing so | ||
| // the reported latency is steady-state replay, not capture. | ||
| int warmup = cudaGraphs ? Math.min(4, nDecode - 1) : 0; |
…RoPE decode kernels); CUDA-backend docs + multi-model perf
mikepapadim
changed the title
| .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, | ||
| context, state.wrapHbFP16Batch, | ||
| weights.w2Layered[layerIndex].asHalfFloatArray(), | ||
| state.w2Out, batchSize, dim, hidDim) |
| long kvElems = (long) B * numLayers * decodeCtx * kvDim; | ||
| FloatArray keyCacheBatch = new FloatArray((int) kvElems); | ||
| FloatArray valueCacheBatch = new FloatArray((int) kvElems); | ||
| keyCacheBatch.init(0.0f); | ||
| valueCacheBatch.init(0.0f); | ||
| IntArray seqPositions = new IntArray(B); | ||
| FloatArray finalScaleBatch = new FloatArray(B); | ||
| HalfFloatArray normedFinalFP16 = new HalfFloatArray(paddedB * dim); | ||
| normedFinalFP16.init(new uk.ac.manchester.tornado.api.types.HalfFloat(0.0f)); | ||
| FloatArray logitsBatch = new FloatArray(paddedB * vocab); | ||
|
|
| // First few steps capture the CUDA graphs; exclude them from timing so | ||
| // the reported latency is steady-state replay, not capture. | ||
| int warmup = cudaGraphs ? Math.min(4, nDecode - 1) : 0; | ||
| long decodeNs = 0; | ||
| int steps = 0; |
| double decodeSec = decodeNs / 1e9; | ||
| int genTokens = steps * B; | ||
| System.out.printf("[perf] prefill %d pos in %.1f ms | decode %d steps × B=%d = %d tokens in %.1f ms%n", | ||
| P, prefillNs / 1e6, steps, B, genTokens, decodeNs / 1e6); | ||
| System.out.printf("[perf] per-step latency %.2f ms | aggregate throughput %.0f tok/s | per-stream %.1f tok/s%n", | ||
| decodeNs / 1e6 / steps, genTokens / decodeSec, steps / decodeSec); |
| int B = args.length > 0 ? Integer.parseInt(args[0]) : 32; | ||
| int K = args.length > 1 ? Integer.parseInt(args[1]) : 2048; | ||
| int N = args.length > 2 ? Integer.parseInt(args[2]) : 2048; | ||
| int iters = 100; |
…-queue, +20% throughput / +24% utilization vs static wave
| int B = args.length > 0 ? Integer.parseInt(args[0]) : 32; | ||
| int K = args.length > 1 ? Integer.parseInt(args[1]) : 2048; | ||
| int N = args.length > 2 ? Integer.parseInt(args[2]) : 2048; | ||
| int iters = 100; |
| long kvElems = (long) B * numLayers * decodeCtx * kvDim; | ||
| FloatArray keyCacheBatch = new FloatArray((int) kvElems); | ||
| FloatArray valueCacheBatch = new FloatArray((int) kvElems); |
| .task("w2Proj", TransformerBatchPrefillKernels::gemmMMA, | ||
| context, state.wrapHbFP16Batch, | ||
| weights.w2Layered[layerIndex].asHalfFloatArray(), | ||
| state.w2Out, batchSize, dim, hidDim) | ||
| .task("w2Resid", TransformerBatchPrefillKernels::batchedResidualAddFP32, |
…KV memory at ~1% overhead (LLaMA)
…locks, +85% throughput / 85% fewer prefill tokens
…d paged flash kernel); +104% throughput with prefix cache
GPU utilization (nsys) on larger modelsProfiled the continuous+paged+prefix engine with Nsight Systems 2024.5
(steady GPU util = Σ kernel time / (steps × wall-per-step); nsys adds ~10–20% wall overhead.) Kernel mix is the same on both — >95% of GPU time is the tensor-core projection GEMMs: So when the GPU is working, it is compute-bound on the tensor cores (the intended batched-decode Where the idle time goes (the actionable finding). Per step the host must: copy the full
Re the other vLLM opts: the profile re-prioritizes them. Chunked prefill — the engine |
Exact model output (greedy, verified)Prompt (
Continuous + paged + prefix-cache, greedy sampling. All completed requests are mutually Llama-3.2-1B (B=128, 128 requests): Qwen3-4B (B=32, 96 requests): Both coherent, on-prompt, not gibberish. (Qwen3 emits its |
…ids not the 65MB logits tensor, +30% throughput
On-device sampling — the nsys-identified fix, benchmarkedPer the GPU-utilization profile above, the engine was host-bound: every step copied the full Effect (nsys, Llama-1B B=128, same workload):
End-to-end (non-nsys, 512-request continuous+paged+prefix, identical greedy output):
+30% throughput, −23% wall. GPU kernel work per step is unchanged (18.8 ms) — the entire |
Static batched decode — B independent LLaMA sequences per step
Decode B independent sequences at once, one token per step, each attending its
own KV cache. Turns the bandwidth-bound single-token matvecs of decode into
compute-bound tensor-core GEMMs (one weight read amortized across B tokens), so
aggregate throughput scales ~linearly with B until the 128-row MMA tile fills.
41× aggregate throughput vs single-stream decode on an RTX 4090 (Llama-3.2-1B
FP16), output verified coherent + bit-exact against the single-stream greedy reference.
Full write-up in
BATCHED_DECODE.md.Why it wins
Single-token decode is memory-bound: each projection reads the whole weight matrix to
make one output vector (intensity ~1). Batch B tokens sharing the weights → each weight
row is read once and applied to all B → intensity ~B, a GEMM. Past B≈16 the
projections go compute-bound and the tensor cores engage. Attention doesn't batch this
way (per-sequence KV) but is a small fraction of decode FLOPs.
Design
One step routine serves both phases. Prompt tokens are fed with logits discarded —
fills every slot's KV region with the same RoPE the decode step uses (no prefill/decode
mismatch). Then generated tokens feed back.
The 12-task layer is identical to the existing batched-prefill MMA layer
(
LlamaFP16LayersBatchPrefillMMA). Only two kernels change, and only in KV addressing:batchedRopeWithKVCachePacked—pos=start+b, shared cachebatchedDecodeRopeWithKVCachePacked—pos=seqPositions[b], per-slot basebatchedFlashAttentionFP16Out— shared causal KVbatchedDecodeAttentionFP16Out— own KV region per slotRoPE math, online-softmax flash attention, register-partitioned P·V — all untouched; the
forks only change the KV index.
Qwen3 reuses the pattern (
Qwen3FP16LayersBatchDecodeMMA): keeps the per-head Q/KRMS-norm task, swaps in split-half
batchedDecodeRopeWithKVCacheQwen3Packed+ the sameper-slot flash kernel (parameterized by qDim/nEmbdHead/gqa). Engine dispatches on model
type; embeddings + batched-logits GEMM are model-agnostic. KV cache is B-sized (
batchIdxstride =numLayers·ctx·kvDim), engine-owned soctxcan be capped for VRAM.Correctness
Greedy + B copies of one prompt → all B streams bit-exact identical AND equal to the
single-stream greedy reference (each slot reproduces the reference independently):
Temperature + per-slot RNG → B distinct, individually coherent continuations (real
concurrent requests, not replication):
Performance (RTX 4090, Llama-3.2-1B FP16, ctx=512, 64 steps, CUDA graphs, steady-state)
Across models (largest B that fits 24 GB; deeper/wider → higher per-step):
All bit-exact vs single-stream greedy + coherent.
Per-step is ~flat (29–31 ms) from B=8 to B=128: every step runs at
paddedBatch=128(MMA tile is 128 rows) and the logits GEMM spans the full vocab regardless of B. So
aggregate scales ~linearly to B=128 (no pad waste) — B=128 is the efficient operating
point. Residual cost is real GEMM compute (full-vocab logits ~33 GFLOP + 16 layer
projections at 128 rows), not launch overhead — hence CUDA graphs help only ~6%.
Reproduce
TornadoVM: needs the CUDA backend + MMA
KernelContext(
mmaFragment/mmaLoadA/mmaMultiply,HalfFloatArray) — the same TornadoVM thisfeat/mma_cudabranch already requires. Tested against TornadoVM 5.0.1-jdk21-dev,CUDA backend, JDK 21 (
make BACKEND=cuda). All numbers here are from the CUDA backend.Build:
mvn -Pjdk21 -Dtornadovm.base.version=5.0.1 -Djdk.version.suffix=-jdk21-dev clean package -DskipTestsRun (
-Dllama.prefillBatchSizeMUST equal-Dbatch.decode.B; takellama-tornado --show-command, swap main class, prepend flags):Two micro-benchmarks isolate the regimes:
BatchedProjectionBench(compute-boundprojection crossover, ~20×) and
BatchedDecodeAttentionBench(memory-bound per-slotattention, ~2×, bit-exact vs CPU ref).
Scope
Additive — no existing path touched (new decode layer graph + 2 kernels + engine/benches).
FP16 LLaMA + Qwen3; same-length prompts (static). Qwen3-8B is garbage in the stock decode path on this build (not this feature). Ragged / continuous batching is the
follow-up — per-slot
seqPositions+ per-slot KV base already allow per-slot positions.Update: continuous batching + PagedAttention (vLLM-style)
Continuous (iteration-level) batching (
-Dbatch.decode.continuous=true): each slot isindependently prefilling or decoding; a slot that hits stop/max is evicted and immediately
refilled from a pending queue — new requests join a partly-decoded batch mid-flight. Zero
kernel changes (the per-step forward already feeds one token per slot at its own
seqPositions[b]). Same 512-request workload (Llama-1B, B=128, max-gen∈[8,64], greedy):+20% throughput / +24% relative utilization; grows with length variance.
PagedAttention (
-Dbatch.decode.paged=true, LLaMA): KV in a shared pool of fixed-sizeblocks + per-slot block table; the two decode KV kernels index
blockTable[b][pos/blockSize].Pool sized to actual concurrent demand, not
B·ctx:~10.7× less KV memory at ~1% overhead, still bit-exact. Undersizing the pool throws a
clear
KV block pool exhausted(where admission control takes over). Prerequisite forprefix caching. Details in
BATCHED_DECODE.md.Update: prefix caching + Qwen3 paging
Prefix caching (
-Dbatch.decode.prefixCache=true, needs paging): a shared prompt prefix(system prompt / preamble) is prefilled once into pinned blocks; every request points its
block-table prefix rows at them and starts decoding at
pos = sharedPrefixLen, skipping theprefix's prefill. The attention kernel walks the block table, so it reads shared-then-private
blocks with no kernel change. Same 512-request continuous+paged workload, 48-tok shared
prompt (Llama-1B, B=128):
~2× fewer steps, +85% throughput; bit-exact.
Qwen3 now runs the full stack (batched decode → continuous → paging → prefix caching) —
adds a split-half paged RoPE kernel; the paged flash kernel is shared (
qDim==nHeads·nEmbdHead).Qwen3-1.7B, 256-request continuous+paged, 48-tok shared prompt:
+104% throughput; prefix sharing also halves peak KV blocks. All bit-exact.
vLLM-style feature set now on both LLaMA and Qwen3: continuous batching + PagedAttention +
prefix caching. Details + reproduction in
BATCHED_DECODE.md.How to test — exact prompts and flags
Launch. The engine is a
mainon the standard TornadoVM module path. Easiest: capture theJVM command from the normal runner, swap the main class to the engine, and prepend the
-Dflags:
-Dllama.prefillBatchSizemust equal-Dbatch.decode.B(it sizes the batch buffers).Larger models need more device memory — add
-Dtornado.device.memory=20GBto<ENGINE_FLAGS>.Prompts used (verbatim):
P_SHORT=Tell me a short story about a robot learning to paint.P_LONG(48-token shared prefix) =You are a helpful and knowledgeable assistant. Please write a detailed, vivid and engaging short story about a small curious robot named Zeta who lives in an old cluttered workshop and slowly learns the delicate art of painting with oil colors.P_QA=What is the capital of France?Exact
<ENGINE_FLAGS>per result (model, prompt):-Dllama.prefillBatchSize=128 -Dbatch.decode.B=128 -Dbatch.decode.ctx=512 -Dbatch.decode.n=64…B=128 … -Dbatch.decode.temp=0.7…B=128 … -Dbatch.decode.continuous=true -Dbatch.decode.requests=512 -Dbatch.decode.minN=8 -Dbatch.decode.refill=true(baseline:refill=false)…continuous=true -Dbatch.decode.paged=true -Dbatch.decode.blocks=384 -Dbatch.decode.requests=512 -Dbatch.decode.minN=8…continuous=true -Dbatch.decode.paged=true -Dbatch.decode.blocks=1024 -Dbatch.decode.prefixCache=true -Dbatch.decode.requests=512 -Dbatch.decode.minN=8(off:prefixCache=false)-Dtornado.device.memory=20GB -Dllama.prefillBatchSize=64 -Dbatch.decode.B=64 -Dbatch.decode.ctx=512 -Dbatch.decode.n=64 -Dbatch.decode.continuous=true -Dbatch.decode.paged=true -Dbatch.decode.blocks=768 -Dbatch.decode.prefixCache=true -Dbatch.decode.requests=256 -Dbatch.decode.minN=8(off:prefixCache=false)-Dtornado.device.memory=20GB -Dllama.prefillBatchSize=32 -Dbatch.decode.B=32 -Dbatch.decode.ctx=512 -Dbatch.decode.n=48Single-stream baselines (stock decode, for the speedup ratios): normal runner with CUDA
graphs, e.g.
JAVA_TOOL_OPTIONS=-Dllama.cudaGraphs=true python3 llama-tornado --gpu --cuda --model <m.gguf> --prompt "<P_SHORT>" --instruct -n 64→ readsachieved tok/s.All flags (defaults in parens):
batch.decode.B,.ctx(512),.n(64),.temp(0=greedy),.cudaGraphs(true),.continuous(false),.refill(true),.requests(4·B),.minN(n/2),.paged(false),.blockSize(16),.blocks(B·ctx/blockSize),.prefixCache(false).Each run prints
[verify] … : true(correctness) and[perf] …(throughput).