Table of Contents
Fetching ...

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.

PackInfer: Compute- and I/O-Efficient Attention for Batched LLM Inference

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 and throughput improvements around , 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.
Paper Structure (38 sections, 5 equations, 13 figures, 5 tables, 1 algorithm)

This paper contains 38 sections, 5 equations, 13 figures, 5 tables, 1 algorithm.

Figures (13)

  • Figure 1: Heterogeneity reduces GPU utilization.
  • Figure 2: PackInfer overview. It adaptively packs requests of heterogeneous input lengths into computation and I/O load-balanced groups.
  • Figure 3: Real-world inference workloads exhibit highly heterogeneous request lengths, with many shorter than 128 tokens.
  • Figure 4: Contiguous Memory Consolidation in PackInfer. PackInfer re-aligns scattered KV cache blocks into a unified, high-density buffer to maximize I/O throughput. A pre-allocated headroom is reserved for each request to accommodate future tokens.
  • Figure 5: PackInfer improves time-to-last-token (TTLT) latency.
  • ...and 8 more figures