PackInfer: Compute- and I/O-Efficient Attention for Batched LLM Inference
Rui Ning, Wei Zhang, Fan Lai
TL;DR
PackInfer introduces a kernel-level packing framework to enable compute- and I/O-efficient attention for batched LLM inference with heterogeneous input lengths. By jointly optimizing packed computation and packed IO, it forms balanced execution groups and reorganizes KV caches to minimize wasted work and memory traffic while preserving lossless FlashAttention semantics. Empirical results show latency reductions of $13.0$–$20.1\%$ and throughput improvements around $20\%$, with consistent gains across models, traces, and GPUs. The approach is modular and drop-in compatible with existing serving stacks, offering a practical path to higher-throughput, lower-latency LLM inference in real-world deployments.
Abstract
Attention efficiency is critical to large language model (LLM) inference. While prior advances optimize attention execution for individual requests (e.g., FlashAttention), production LLM serving relies on batching requests with highly heterogeneous sequence lengths for high serving throughput. This mismatch induces severe computation and I/O imbalance, exacerbates stragglers, and underutilizes GPU resources. We present PackInfer, a kernel-level attention framework that enables compute- and I/O-aware execution for heterogeneous batched inference. PackInfer orchestrates batched requests into load-balanced execution groups, effectively saturating GPU utilization by packing multiple requests into unified kernel launches. By constructing attention kernels directly over packed query-key regions, PackInfer eliminates redundant computation and balances thread-block execution. It then incorporates I/O-aware grouping that co-locates shared-prefix requests and reorganizes KV caches into group-contiguous layouts, reducing memory fragmentation and redundant data movement as generation evolves. Evaluations on real-world workloads show that PackInfer reduces inference latency by 13.0-20.1%, and improves throughput by 20% compared to the state-of-the-art FlashAttention.
