Working to take a 7B RAG service from 300 to 600 tok/s/GPU on 8×A100 while holding p95 <120 ms; speculative decoding gave about 18%, but attention kernels still stall at small dynamic batches. Has anyone nailed this with paged attention + KV cache sharding, or did prefix-caching with smarter batch bucketing scale better under bursty traffic?
But splitting prefill and decode into separate micro-batching queues and bucketing by ‘tokens_remaining’ (8/16/32) helped our small-batch stalls — think two barista lanes; with vLLM’s PagedAttention plus a per-request max_active_blocks we avoided page churn and held p95≈110 ms while nudging 600 tok/s/GPU (GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs). If your RAG has high prompt overlap, ‘prefix-caching’ with a warm prefix pool can beat KV sharding during bursts; are you on vLLM or TensorRT-LLM?
Switching paged attention to 16-token pages and capturing CUDA Graphs for prefill/decode killed most of our small-batch stalls; with a pinned KV allocator (no per-request malloc) we saw about 12% p95 drop and about 620 tok/s/GPU on 8×A100 with KV sharding. If you’re on vLLM, try page_size=16 and tensor-parallel KV shards (see GitHub - vllm-project/vllm: A high-throughput and memory-efficient inference and serving engine for LLMs); caveat: graphs can regress under bursty mixes unless capture shapes stay stable — have you tried that?
Building on @c_thompson91: the single change that steadied our p95 was putting decode on a higher‑priority CUDA stream (“decode > prefill”), which let us overlap shard all‑gathers without getting preempted; we saw about 10% and far fewer tail spikes… If you’re already using CUDA Graphs, stream priorities are cheap to try: CUDA C++ Programming Guide (Legacy) — CUDA C++ Programming Guide — have you tested that?
Given “p95 <120 ms”, the biggest win for us was running decode as a persistent kernel with a device‑side work queue; it kept tiny decodes from getting starved when batches thinned. Pair that with KV sharding across GPUs but keep hot heads local (GQA) and we saw fewer stalls than leaning solely on paged attention. @t_kimura21 did you also gate long prefills (>2k) during bursts so decode always has N active sequences?
One lever that steadied our small-batch decode was storing the K/V cache in int8 with per-head scales and dequantizing inside the attention kernel; on 8×A100 we stayed under “p95 <120 ms” while pushing a 7B RAG from about 300 toward about 600 by cutting HBM reads. If you try it, calibrate per-head ranges on a RAG‑like corpus and clamp outliers, since long contexts can drift without good scales.