PaperlayerView exact revision on arXiv

Source-linked AI summary

UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification

Qihang Fan, Huaibo Huang, Zhiying Wu, Bingning Wang, Ran He

arXiv:2605.06221v1cs.CL

TL;DR

Long-context prefill acceleration is limited by methods specialized for full attention and difficult to integrate with continuous batching. UniPrefill performs token-level sparsification across model layers and integrates it into vLLM. It achieves up to 2.1× TTFT speedup with negligible accuracy loss, with gains scaling favorably with concurrency and context length.

  • Problem

    Long-context prefill methods largely target sparse attention, limiting their effectiveness on hybrid architectures and complicating integration with continuous batching.

  • Method

    UniPrefill estimates token importance at full-attention layers, drops redundant tokens, propagates sparsity across subsequent layers, and integrates continuous batching into vLLM.

  • Results

    Up to 2.1× TTFT speedup is achieved with negligible accuracy loss, with gains scaling favorably with request concurrency and context length.

  • Takeaways & Limitations

    UniPrefill provides an architecture-agnostic approach for efficient long-context serving across full-attention and hybrid models.

  • Takeaways & Limitations

    Existing acceleration methods remain difficult to integrate with continuous batching, while UniPrefill’s reported setup uses model-specific thresholds and retention settings.

Abstract

from arXiv · show

As large language models (LLMs) continue to advance rapidly, they are becoming increasingly capable while simultaneously demanding ever-longer context lengths. To improve the inference efficiency of long-context processing, several novel low-complexity hybrid architectures have recently been proposed, effectively alleviating the computational burden of long-context inference. However, existing research on long-context prefill acceleration remains predominantly focused on sparse attention mechanisms, which achieve their maximum speedup only on full-attention models. When transferred to emerging architectures--such as linear/full attention hybrids or sliding window/full attention hybrids--these prefill acceleration approaches suffer significant performance degradation. Furthermore, such methods are generally incompatible with continuous batching, making them difficult to integrate into modern inference engines such as vLLM. To this end, we propose UniPrefill, a prefill acceleration framework applicable to virtually any model architecture, which directly accelerates the model's computation at the token level. We further implement UniPrefill as a continuous batching operator and extend vLLM's scheduling strategy to natively support prefill-decode co-processing and tensor parallel for UniPrefill, enabling its seamless integration into vLLM. UniPrefill achieves up to 2.1x speedup in Time-To-First-Token (TTFT), with the acceleration becoming increasingly pronounced as the number of concurrent requests grows.

1. Introduction

Long-context inference faces prohibitive costs from quadratic attention, while existing prefill acceleration is limited on hybrid architectures and difficult to use with continuous batching. UniPrefill addresses these constraints through token-level sparsification and vLLM integration, achieving substantial TTFT and throughput gains with negligible accuracy loss.

  • Modern deployments increasingly require tens of thousands to million-token contexts, making quadratic Softmax attention computationally prohibitive.
  • Existing sparse-attention prefill methods lose marginal benefit on hybrid architectures because they accelerate only a fraction of layers.In a linear/full-attention hybrid with a 3:1 ratio, at most one in four layers is full attention.
  • Existing prefill acceleration methods are difficult to integrate with continuous batching because they assume isolated requests and static batch composition.This limits deployment in high-throughput inference engines such as vLLM.
  • UniPrefill scores token importance at full-attention layers, drops redundant tokens, and propagates sparsity through subsequent layers.The cascading mask reduces both attention and GEMM FLOPs across pure full-attention and hybrid architectures.
  • UniPrefill extends vLLM scheduling to support continuous batching and prefill-decode co-processing without changing model weights or serving infrastructure.
  • Up to 2.1× TTFT speedup is achieved with negligible accuracy degradation, and acceleration increases with concurrent requests.The evaluation uses RULER across multiple model architectures.

2. Related Works

Efficient and hybrid architectures address the cost of long-context modeling, but existing prefill acceleration remains largely specialized for full-attention models. Sparse-attention methods also leave non-attention computation untouched and are difficult to integrate with continuous batching.

  • Hybrid architectures interleave full attention with efficient alternatives to balance efficiency and expressiveness.
  • Existing prefill acceleration methods remain largely tailored to full-attention-only architectures, limiting effectiveness on newer hybrid models.
  • Sparse-attention methods accelerate prefill by skipping computations associated with static or dynamic sparse attention patterns.
  • These methods leave FFN and GEMM computations unaccelerated and are generally incompatible with continuous batching.UniPrefill addresses both limitations by operating at the token level and propagating sparsity across layers.

3. Method

UniPrefill estimates token importance at full-attention layers, retains enough blocks to reach a top-p mass threshold, and propagates dropped tokens across later computation. Its fused kernels and vLLM integration support continuous batching and tensor parallelism while reducing attention and GEMM work.

  • Token Importance Estimation: UniPrefill estimates token importance from the last n query positions, aggregates scores over non-overlapping token blocks, and selects blocks by cumulative attention mass.The partial attention computation uses online softmax before block reduction, reducing selection decisions from N tokens to ⌈N/G⌉ blocks.
  • Top-p Token Selection: The retained set is the smallest collection of descending-score blocks whose cumulative importance reaches p, while attention sinks and the query window are always retained.Top-p adapts the retained set to concentrated or diffuse attention and provides a uniform approximation-error bound; p = 0.99 discards at most 1% of attention mass.
  • Sparsity Propagation Across All Layers: After selection, dropped tokens are excluded from subsequent full-attention, linear-attention, sliding-window, and FFN sublayers, while their states are carried forward and scores are recomputed at later full-attention layers.This cascading reduction lowers token counts across all downstream layers rather than only within one attention sublayer.
  • FLOPs Analysis: A single drop at layer ℓ1 with retention ratio ρ saves (1 − ρ) · (L − ℓ1) · O(Nd^2) across remaining layers, whereas sparse attention leaves later GEMM costs intact.In the long-context regime N ≫ d, UniPrefill’s GEMM savings dominate.
  • Fused Kernel and vLLM Integration: Fused GPU kernels compute scores, block reductions, top-p masks, and token expansion directly on packed variable-length representations without padding or CPU round-trips.Tensor-parallel ranks synchronize partial block scores before applying the shared top-p decision.
  • Fused Kernel and vLLM Integration: vLLM integration updates layerwise metadata, KV-cache mappings, and per-request sequence lengths so continuous batching and tensor-parallel execution remain consistent after token dropping.The per-layer sequence correction is injected before each decode step without changing model weights or the PagedAttention allocator.

4. Experiments

UniPrefill is evaluated for accuracy and efficiency across full-attention and hybrid architectures using RULER and vLLM throughput measurements. The experiments also examine block size and query-window choices.

  • Evaluation setup: UniPrefill is evaluated across LLaMA-3.1-8B-Instruct, Qwen3-Next-80B-A3B, and Gemma-3-12B, representing full-attention, linear/full-attention, and sliding-window/full-attention architectures.RULER evaluates retrieval, multi-hop tracing, aggregation, and question answering at configurable lengths up to 128K tokens.
  • Accuracy and TTFT: At 128K, UniPrefill delivers up to 2.26×, 1.68×, and 1.49× TTFT speedup on LLaMA-3.1-8B, Qwen3-Next-80B-A3B, and Gemma-3-12B, respectively.The method retains accuracy close to the Baseline and provides the best reported accuracy-efficiency tradeoff among compared methods.
  • vLLM throughput: UniPrefill improves vLLM prefill throughput by up to +109%, +68%, and +42% on LLaMA-3.1-8B, Qwen3-Next-80B-A3B, and Gemma-3-12B, respectively.Throughput gains increase with context length and batch size.
  • Block-size ablation: At 128K, block size G=32 reaches up to +121% and +78% throughput gain on LLaMA-3.1-8B and Qwen3-Next-80B-A3B, respectively.The default G=64 balances selection overhead and token drop rate across context lengths.
  • Query-window ablation: n=32 causes noticeable accuracy loss, n=512 adds computational overhead, and n=128 provides the best balance for importance estimation.The default n=128 is adopted for the experiments.

5. Conclusion

UniPrefill accelerates long-context prefill across full-attention and hybrid architectures by propagating token sparsity beyond full-attention layers. Its vLLM integration supports continuous batching without model-weight changes, while experiments report up to 2.1× TTFT speedup with negligible accuracy loss.

  • Method: UniPrefill estimates token importance through block-wise top-p selection at full-attention layers and propagates the sparsity mask across subsequent sub-layers.This reduces attention and GEMM FLOPs across full-attention and hybrid architectures.
  • Deployment: UniPrefill integrates into vLLM’s continuous-batching scheduler without changing model weights.The implementation supports prefill-decode co-processing and tensor parallelism.
  • Results: Up to 2.1× TTFT speedup is reported with negligible accuracy loss, with gains scaling favorably with context length and batch size.The experiments are conducted on the RULER benchmark.

A. Implementation and deployment details

UniPrefill is deployed in vLLM with production-oriented batching and parallelism support. Its operator uses fused Triton kernels designed to be hardware-agnostic.

  • vLLM deployment: UniPrefill is implemented on vLLM v0.16.0 with support for prefill-decode co-processing and tensor parallelism.Throughput experiments use tensor parallelism degree TP = 8.
  • Kernel implementation: The continuous-batching operator is implemented as fused Triton kernels designed to be hardware-agnostic.

B. Experiment statistical significance.

UniPrefill’s performance remains stable across random-seed initializations. The statistical-significance experiment is reported as an ablation over different random seeds.

  • Random-seed robustness: Consistently stable performance across multiple random seeds demonstrates robustness to random-seed initialization.

C. Limitations and Broader Impacts

UniPrefill is scoped to accelerating long-context LLM prefill, while decoding acceleration and training-time efficiency remain future directions. Ethical deployment and safety alignment are outside the study's technical scope.

  • UniPrefill focuses on accelerating the prefill phase for long-context LLM inference.
  • Extending UniPrefill to decoding acceleration and training-time efficiency remains future work.
  • Ethical deployment and safety alignment lie outside this study's technical scope.
Loading 2605.06221v1…