beehive-lab / beehive-lab/GPULlama3.java
Roadmap: vLLM-style serving features for TornadoVM batched decode
- Dominant language
- Java
- Stars
- 274
- Forks
- 37
- Avg merge
- 5d 1h
- Merged PRs (30d)
- 2
Description
## vLLM-style serving features for TornadoVM batched decode
Tracking the vLLM-inspired throughput/serving features for the GPU batched-decode engine
(`bench/BatchedDecodeEngine`, CUDA backend). Landed features link to the commit that adds
them on branch `feat/static-batched-decode` (PR #129). Full design + numbers in
`BATCHED_DECODE.md`.
### Core batched decode
- [x] Static batched decode — B independent sequences/step, per-slot KV, 41× aggregate (LLaMA) — `0c501a6`
- [x] Qwen3 support (per-head QK-norm + split-half RoPE decode kernels) — `b7cf126`
- [x] CUDA-graph capture/replay per step (~6% per-step) — `0c501a6`
- [ ] Temperature sampling → B divergent independent streams — `0c501a6`
### Serving / scheduling (vLLM)
- [x] **Continuous (iteration-level) batching** — evict-on-stop + admit-from-queue, +20% throughput / +24% utilization vs static wave — `4843943`
- [x] **PagedAttention** — block-pool KV + per-slot block table, ~10.7× less KV memory at ~1% overhead — `37e7f61`
- [x] **Prefix caching** — shared prompt prefix prefilled once into pinned blocks, +85% throughput / 85% fewer prefill tokens — `a81268f`
- [x] Paging + prefix caching extended to Qwen3 (+104% throughput) — `50eca7d`
### Sampling / logits
- [x] **On-device sampling** (GPU argmax) — transfer B token ids not the 65 MB logits tensor, +30% throughput — `fc1c16e`
- [x] Logits pruning — skip the full-vocab GEMM on pure-prefill steps — `3dd54da`
- [ ] On-device temperature / top-p sampling (needs a GPU RNG) — _future_
- [ ] Quantized (Q8/FP8) KV cache — _future_
### Deferred / WIP (rationale in PR #129 GPU-utilization comment)
- [ ] **Speculative decoding** (prompt-lookup draft + batch-prefill verify) — scheduling proven (1.71 tokens/forward, +71% fewer forwards) but the standalone verify-forward has an output bug; not shipped — _WIP_
- [ ] **Chunked prefill** — engine already interleaves token-by-token prefill with decode (anti-stall); multi-token-per-step chunking conflicts with the one-token-per-slot layout — _deferred_
- [ ] Host/GPU pipeline overlap (double-buffer sampling vs next step) — _future_
### Models
- [x] LLaMA (FP16) batched decode — `0c501a6`
- [x] Qwen3 (FP16) batched decode — `b7cf126`
- [ ] Gemma 4 batched decode (per-layer dims + sliding-window + GeGLU + PLE) — _in progress on `feat/gemma4-batched-decode`_
### Documentation
- [x] `BATCHED_DECODE.md` — design (ASCII), per-model perf, exact prompts + flags — `19ce1b0`
- [x] nsys GPU-utilization analysis (host-bound finding → on-device sampling) — PR #129 comments
Reproduction, exact prompts, and per-result flags are in `BATCHED_DECODE.md`. Profiling shows
the engine is compute-bound on the projection GEMMs when the GPU is working; on-device sampling
removed the main host-side stall (65–78 MB/step D2H + CPU argmax).
Contributor guide
Assessment
This issue has not been assessed yet.