UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification
Summary
UniPrefill is a new prefill acceleration framework proposed in a research paper that enables block-wise dynamic sparsification for universal long-context processing in LLMs. It integrates with vLLM to achieve up to 2.1x speedup in Time-To-First-Token across various model architectures.
View Cached Full Text
Cached at: 05/08/26, 07:25 AM
# UniPrefill: Universal Long-Context Prefill Acceleration via Block-wise Dynamic Sparsification
Source: [https://arxiv.org/html/2605.06221](https://arxiv.org/html/2605.06221)
\\reportnumber
001\\githubhttps://github\.com/qhfan/UniPrefill\.git
Qihang Fan1,2,3,∗, Huaibo Huang1,2,†, Zhiying Wu3, Bingning Wang3,‡, Ran He1,2 1MAIS&NLPR, CASIA2UCAS3WeChat, Tencent
###### Abstract
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 proposeUniPrefill, 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 to2\.1xspeedup in Time\-To\-First\-Token \(TTFT\), with the acceleration becoming increasingly pronounced as the number of concurrent requests grows\.
†††Corresponding Author\.††‡Project Leader\.††∗Work done during internship at WeChat\.## 1Introduction
The rapid advancement of large language models \(LLMs\) has driven their deployment across an increasingly diverse range of real\-world applications, from document understanding and code generation to multi\-turn dialogue and retrieval\-augmented generation\[llama,llama2,qwen2\.5\-1m,qwen25technicalreport,qwen3technicalreport,qwentechnicalreport,glm2024chatglm\]\. Alongside this expansion in capability, the context lengths that LLMs are expected to process have grown dramatically — modern deployments routinely involve sequences of tens of thousands of tokens, and the demand for hundred\-thousand\-token or even million\-token contexts is becoming commonplace\. This trend places enormous pressure on inference efficiency, as the canonical Softmax Self\-Attention\[attention\]mechanism scales quadratically with sequence length, incurring prohibitive computational costs when processing long contexts\.
To address the quadratic complexity bottleneck, a new generation of hybrid architectures has emerged that interleave computationally efficient layers with full attention layers\. Two representative families have gained particular traction: linear/full attention hybrids, which replace a subset of attention layers with linear recurrent mechanisms\[mamba,mamba2,yang2024gla,fan2024rect,fan2024breaking\]to reduce per\-layer complexity fromO\(N2\)O\(N^\{2\}\)toO\(N\)O\(N\); and sliding window/full attention hybrids, which restrict most attention layers to a fixed local context window while retaining a small number of global full\-attention layers for long\-range dependencies\[gemmateam2025gemma3technicalreport,jiang2023mistral7b\]\. These hybrid designs substantially reduce the theoretical complexity of long\-context inference and have been widely adopted in recently released production\-grade models\.
Figure 1:Prefill throughput comparison between Standard Prefill and UniPrefill across three model architectures and varying batch sizes \(tensor parallel size is set to 8\)\. All experiments are conducted within vLLM, with UniPrefill deeply integrated into vLLM’s continuous batching scheduler\. We evaluate prefill throughput \(K tokens/s\) on LLaMA\-3\.1\-Instruct\-8B\[llama3\]\(full attention\), Qwen3\-Next\-80B\-A3B\[qwen3next\_blog\_2025\]\(linear/full attention hybrid\), and Gemma\-3\-12B\[gemmateam2025gemma3technicalreport\]\(sliding window/full attention hybrid\) across context lengths from 4K to 128K and batch sizes of 1, 4, 16, and 64\. Solid bars denote Standard Prefill and hatched bars denote UniPrefill\. UniPrefill consistently achieves higher throughput across all three architectures, with gains becoming more pronounced at longer context lengths and larger batch sizes\.Despite the proliferation of hybrid architectures, the research community’s efforts on prefill acceleration have remained heavily concentrated on sparse attention\[minference,mobamixtureblockattention,fan2026flashprefill\]\. Representative works such as MInference\[minference\]have demonstrated impressive prefill speedups, achieving up to 10× acceleration on long sequences under the full\-attention\-only setting\. However, this focus on sparse attention comes with a fundamental limitation: the acceleration is tightly coupled to the full attention operation itself\. In hybrid architectures where full attention constitutes only a fraction of all layers, the marginal benefit of accelerating solely those attention layers diminishes considerably\. For instance, in a linear/full attention hybrid with a 3:1 ratio, at most one out of every four layers can be accelerated by existing sparse attention methods, leaving the dominant computational budget entirely untouched\. This architectural mismatch renders existing prefill acceleration approaches far less effective on the new generation of hybrid models\.
A second, equally critical limitation of existing prefill acceleration methods is their incompatibility with continuous batching, the scheduling paradigm that underpins modern high\-throughput inference engines such as vLLM\[vLLM,zheng2024sglang\]\. Methods such as FlexPrefill\[flexprefill\]operate on individual requests in isolation and assume static batch composition, making them fundamentally difficult to integrate into a continuous batching scheduler where requests enter and exit the batch dynamically\. As a result, these methods have largely remained research prototypes and have not been successfully embedded into production inference systems\.
To overcome both limitations, we propose UniPrefill, a prefill acceleration framework that achieves architecture\-agnostic speedups by exploiting a key insight: token importance can be estimated at full attention layers and propagated across all subsequent layers\. Specifically, UniPrefill applies a lightweight block\-wise scoring criterion at each full attention layer to identify and drop computationally redundant tokens\. Once a token is dropped, it is excluded from all downstream computation in the remaining layers of the block\. This cascading effect means that a single token\-dropping decision at the attention layer translates into a proportional reduction in computation across the entire layer stack, not merely the attention sublayer\. As a result, UniPrefill achieves substantial reductions in both attention FLOPs and GEMM FLOPs simultaneously, making it effective regardless of whether the model is a pure full\-attention Transformer or hybrid architecture\.
Beyond the algorithmic design, we address the systems integration challenge by implementing UniPrefill as a continuous batching operator\[yu2022orca\]and extending vLLM\[vLLM\]’s scheduler to natively support prefill\-decode co\-processing under UniPrefill’s token\-dropping regime\. This tight integration allows UniPrefill to function as a transparent acceleration layer within production inference engines, without requiring changes to model weights or serving infrastructure\.
We evaluate UniPrefill on RULER\[hsieh2024ruler\]with multiple model architectures\. Results demonstrate that UniPrefill introduces no significant accuracy degradation while achieving up to2\.1×2\.1\\timesspeedup in Time\-To\-First\-Token \(TTFT\), as illustrated in Fig\.[1](https://arxiv.org/html/2605.06221#S1.F1)\. Notably, the speedup scales favorably with the number of concurrent requests \(see Fig\.[1](https://arxiv.org/html/2605.06221#S1.F1)\), making UniPrefill particularly well\-suited for high\-concurrency production serving scenarios where prefill cost is the dominant bottleneck\.
Our main contributions are summarized as follows:
- •We proposeUniPrefill, a token\-level prefill acceleration framework that drops tokens at full attention layers and propagates sparsity across all subsequent layers, reducing both attention and GEMM FLOPs simultaneously, which enables consistent speedups across heterogeneous hybrid architectures\.
- •We implement UniPrefill as a continuous batching operator and integrate it into vLLM\[vLLM\]via extended scheduling strategies that support prefill\-decode co\-processing and tensor parallel, enabling seamless production\-ready deployment\.
- •Extensive experiments on the long context benchmark RULER demonstrate that UniPrefill achieves up to2\.1×2\.1\\timesTTFT speedup with negligible accuracy loss, with acceleration gains scaling with request concurrency\.
## 2Related Works
#### Hybrid LLM Architectures\.
To overcome the quadratic complexity of Softmax attention, a rich body of work has proposed efficient sequence modeling alternatives, including state space models, linear attention variants, and recurrent architectures\[mamba,mamba2,sun2023retentivenetworksuccessortransformer,yang2024gla,yang2024deltanet,fan2025sec,fan2024rect,minimax01scalingfoundationmodels,yang2024gdn,zhang2025kda\]\. To balance efficiency and expressiveness, hybrid architectures have emerged that interleave full attention with these efficient alternatives\[qwen3next\_blog\_2025,lenz2025jamba,gemmateam2025gemma3technicalreport,xiao2026mimov2flash,jiang2023mistral7b\], and have been widely adopted in recently released production models\. However, existing prefill acceleration methods remain largely tailored to full\-attention\-only architectures, limiting their effectiveness on this new generation of models\.
#### Sparse Attention for Prefill Acceleration\.
Exploiting the inherent sparsity in attention score matrices is a well\-established strategy for accelerating the prefill stage\. A body of work identifies static or dynamic sparse patterns — such as vertical, slash, and block\-sparse structures — and skips the corresponding attention computations\[minference,native\-sparse\-attention,mobamixtureblockattention,optimizingmixtureblockattention,flexprefill,chen2026vsprefill\]\. These methods have demonstrated substantial speedups on full attention models\[minference,flexprefill,xattention,wang2025proxyattn\]\. However, they share two fundamental limitations: their acceleration is tightly coupled to the attention operation itself, leaving FFN and GEMM computations entirely unaccelerated, and they are generally incompatible with continuous batching\[yu2022orca\], making integration into production inference engines such as vLLM\[vLLM\]non\-trivial\. UniPrefill addresses both limitations by operating at the token level and propagating sparsity across all layers\.
## 3Method
In this section, we present UniPrefill, an architecture\-agnostic prefill acceleration framework\. The overall pipeline is illustrated in Fig\.[2](https://arxiv.org/html/2605.06221#S3.F2)\.
### 3\.1Preliminaries
Consider an input sequence𝐱=\[x1,…,xN\]\\mathbf\{x\}=\[x\_\{1\},\\ldots,x\_\{N\}\]processed by a hybrid LLM consisting ofBBblocks\. Each blockbbcontains a full attention layer followed byMbM\_\{b\}sublayers \(linear attention, sliding window attention, FFN, etc\.\)\. Let𝐇\(b,0\)∈ℝN×d\\mathbf\{H\}^\{\(b,0\)\}\\in\\mathbb\{R\}^\{N\\times d\}denote the block input\. The goal of prefill is to compute the final hidden state𝐡N\(L\)\\mathbf\{h\}\_\{N\}^\{\(L\)\}for next\-token prediction:
P\(xN\+1∣x1:N\)=LMHead\(𝐡N\(L\)\)P\(x\_\{N\+1\}\\mid x\_\{1:N\}\)=\\text\{LMHead\}\\\!\\left\(\\mathbf\{h\}\_\{N\}^\{\(L\)\}\\right\)\(1\)Standard prefill incurs𝒪\(N2dk\)\\mathcal\{O\}\(N^\{2\}d\_\{k\}\)per full attention layer and𝒪\(Nd2\)\\mathcal\{O\}\(Nd^\{2\}\)per GEMM sublayer, totaling𝒪\(N2dk\+MbNd2\)\\mathcal\{O\}\(N^\{2\}d\_\{k\}\+M\_\{b\}Nd^\{2\}\)per block\.
### 3\.2Token Importance Estimation
Since next\-token prediction depends solely on𝐡N\(L\)\\mathbf\{h\}\_\{N\}^\{\(L\)\}, the contribution of tokeniito the final hidden state at blockbbis:
𝐡N\(b,1\)=∑i=1N𝐀N,i\(b\)⋅𝐯i\(b\)\+𝐡N\(b,0\),\\mathbf\{h\}\_\{N\}^\{\(b,1\)\}=\\sum\_\{i=1\}^\{N\}\\mathbf\{A\}^\{\(b\)\}\_\{N,i\}\\cdot\\mathbf\{v\}\_\{i\}^\{\(b\)\}\+\\mathbf\{h\}\_\{N\}^\{\(b,0\)\},\(2\)where𝐀N,i\(b\)=softmaxi\(𝐪N\(b\)𝐊\(b\)⊤/dk\)\\mathbf\{A\}^\{\(b\)\}\_\{N,i\}=\\operatorname\{softmax\}\_\{i\}\\\!\\left\(\\mathbf\{q\}\_\{N\}^\{\(b\)\}\{\\mathbf\{K\}^\{\(b\)\}\}^\{\\top\}/\\sqrt\{d\_\{k\}\}\\right\)is the full\-sequence attention weight\. A tokeniiis negligible to next\-token prediction when𝐀N,i\(b\)≈0\\mathbf\{A\}^\{\(b\)\}\_\{N,i\}\\approx 0\.
To reduce estimation variance, we aggregate over the lastnnquery positions instead of a single position:
si\(b\)=1n∑j=N−n\+1N𝐀j,i\(b\),s\_\{i\}^\{\(b\)\}=\\frac\{1\}\{n\}\\sum\_\{j=N\-n\+1\}^\{N\}\\mathbf\{A\}^\{\(b\)\}\_\{j,i\},\(3\)requiring ann×Nn\\times Nattention computation at cost𝒪\(nNdk\)\\mathcal\{O\}\(nNd\_\{k\}\), negligible forn≪Nn\\ll N\.
In practice, importance estimation and token selection operate at*block granularity*\. We partition the input sequence into non\-overlapping blocks of sizeGG:ℬg=\{\(g−1\)G\+1,…,min\(gG,N\)\}\\mathcal\{B\}\_\{g\}=\\\{\(g\-1\)G\+1,\\ldots,\\min\(gG,N\)\\\},g=1,…,⌈N/G⌉g=1,\\ldots,\\lceil N/G\\rceil\. For efficiency, the partial GEMM𝐒=𝐐\[N−n:N\]𝐊⊤∈ℝn×N\\mathbf\{S\}=\\mathbf\{Q\}\_\{\[N\-n:N\]\}\\mathbf\{K\}^\{\\top\}\\in\\mathbb\{R\}^\{n\\times N\}is computed first; an online softmax is then applied*across the full sequence dimension*to obtain properly normalised attention weights, after which scores are reduced within each block:
s¯g\(b\)=1G∑i∈ℬg1n∑j=N−n\+1N𝐀j,i\(b\),\\bar\{s\}\_\{g\}^\{\(b\)\}=\\frac\{1\}\{G\}\\sum\_\{i\\in\\mathcal\{B\}\_\{g\}\}\\frac\{1\}\{n\}\\sum\_\{j=N\-n\+1\}^\{N\}\\mathbf\{A\}^\{\(b\)\}\_\{j,i\},\(4\)where the softmax normalisation is performed over the complete key sequence before the block reduction, ensurings¯g\(b\)\\bar\{s\}\_\{g\}^\{\(b\)\}reflects the true attention mass captured by blockgg\. This reduces the number of selection decisions fromNNto⌈N/G⌉\\lceil N/G\\rceilwhile preserving the accuracy of importance estimation\.
#### Relationship to SnapKV\.
Our importance estimation shares a surface\-level similarity with SnapKV\[li2024snapkv\], which also uses an observation window to identify important tokens\. However, the two methods differ fundamentally in objective and scope\. SnapKV completes a fullN×NN\\times Nprefill across all layers before applying its selection to compress the KV cache for decode—the prefill FLOPs are entirely unaffected\. UniPrefill applies selection*during*prefill, propagating the drop decision forward through all subsequent layers\. Formally, whereas SnapKV saves at most𝒪\(N⋅dkv\)\\mathcal\{O\}\(N\\cdot d\_\{kv\}\)in decode\-time memory per layer, UniPrefill saves\(1−ρ\(b\)\)⋅Mb⋅𝒪\(Nd2\)\(1\-\\rho^\{\(b\)\}\)\\cdot M\_\{b\}\\cdot\\mathcal\{O\}\(Nd^\{2\}\)in prefill\-time FLOPs per block, whereρ\(b\)=\|𝒮\(b\)\|/N\\rho^\{\(b\)\}=\|\\mathcal\{S\}^\{\(b\)\}\|/Nis the token retention ratio—a quantity that grows linearly withMbM\_\{b\}and is entirely absent in SnapKV\.
### 3\.3Top\-ppToken Selection
Letπ\\pibe the permutation sorting block\-level scores\{s¯g\(b\)\}\\\{\\bar\{s\}\_\{g\}^\{\(b\)\}\\\}in descending order\. We retain the minimal set of blocks:
𝒮\(b\)=\{π\(1\),…,π\(k∗\)\},k∗=minks\.t\.∑j=1ks¯π\(j\)\(b\)∑gs¯g\(b\)⩾p\\mathcal\{S\}^\{\(b\)\}=\\left\\\{\\pi\(1\),\\ldots,\\pi\(k^\{\*\}\)\\right\\\},\\qquad k^\{\*\}=\\min k\\ \\text\{ s\.t\. \}\\ \\frac\{\\sum\_\{j=1\}^\{k\}\\bar\{s\}\_\{\\pi\(j\)\}^\{\(b\)\}\}\{\\sum\_\{g\}\\bar\{s\}\_\{g\}^\{\(b\)\}\}\\geqslant p\(5\)The dropped set is𝒮¯\(b\)=\[N\]∖𝒮\(b\)\\bar\{\\mathcal\{S\}\}^\{\(b\)\}=\[N\]\\setminus\\mathcal\{S\}^\{\(b\)\}\. Two structural elements are always retained regardless of their scores: the firstAAtokens \(attention sinks\[xiao2023streamingllm\]\) and the lastnntokens \(the query window itself\), ensuring causal consistency and numerical stability\.
#### Error bound\.
The perturbation to any retained positionjjdue to dropping𝒮¯\(b\)\\bar\{\\mathcal\{S\}\}^\{\(b\)\}satisfies:
‖Δ𝐡j\(b,1\)‖⩽\(∑i∈𝒮¯\(b\)𝐀j,i\(b\)\)⋅Vmax\(b\)⩽\(1−p\)⋅Vmax\(b\)\\left\\\|\\Delta\\mathbf\{h\}\_\{j\}^\{\(b,1\)\}\\right\\\|\\;\\leqslant\\;\\left\(\\sum\_\{i\\in\\bar\{\\mathcal\{S\}\}^\{\(b\)\}\}\\mathbf\{A\}^\{\(b\)\}\_\{j,i\}\\right\)\\cdot V\_\{\\max\}^\{\(b\)\}\\;\\leqslant\\;\(1\-p\)\\cdot V\_\{\\max\}^\{\(b\)\}\(6\)whereVmax\(b\)=maxi‖𝐯i\(b\)‖V\_\{\\max\}^\{\(b\)\}=\\max\_\{i\}\\\|\\mathbf\{v\}\_\{i\}^\{\(b\)\}\\\|\. Settingp=0\.99p=0\.99guarantees that at most1%1\\%of the total attention mass is discarded, providing a direct information\-theoretic bound on the approximation error at the attention layer\.
#### Top\-ppvs\. top\-kk\.
A fixed top\-kkis insensitive to the actual distribution of attention: when attention is highly concentrated, top\-kkretains many unnecessary tokens; when diffuse, it may drop tokens with non\-trivial contributions\. Top\-ppadapts automatically—the retained set is small when attention is concentrated and large when it is diffuse—providing a uniform bound on approximation error regardless of sequence length or content, which top\-kkcannot guarantee\.
Figure 2:Overview of UniPrefill\.Left:UniPrefill estimates token importance via block\-level attention scores from the lastnnqueries\(1\), retains the smallest set of token blocks whose cumulative importance reachespp\(2\), and propagates the resulting sparsity across all subsequent sub\-layers within each repeating layer pattern\(3\)\.Right:UniPrefill is deeply integrated into vLLM via a fused kernel pipeline\(4\), with KV cache block tables\(5\), per\-layer sequence length tracking\(6\), and tensor\-parallel metadata synchronisation\(7\)updated accordingly\.
### 3\.4Sparsity Propagation Across All Layers
After token selection at the full attention layer of blockbb, dropped tokens𝒮¯\(b\)\\bar\{\\mathcal\{S\}\}^\{\(b\)\}are excluded from all subsequent sublayers within and beyond the block—every full attention, linear attention, sliding window attention, and FFN layer processes only the retained set𝒮\(b\)\\mathcal\{S\}^\{\(b\)\}:
𝐇𝒮\(b,m\+1\)=fm\(𝐇𝒮\(b,m\)\),m=1,…,Mb\\mathbf\{H\}\_\{\\mathcal\{S\}\}^\{\(b,m\+1\)\}=f\_\{m\}\\\!\\left\(\\mathbf\{H\}\_\{\\mathcal\{S\}\}^\{\(b,m\)\}\\right\),\\qquad m=1,\\ldots,M\_\{b\}\(7\)At blockb\+1b\+1, the full sequence is reconstituted by carrying dropped token states forward without update:
𝐇i\(b\+1,0\)=\{𝐇i\(b,Mb\+1\)i∈𝒮\(b\)𝐇i\(b,0\)i∈𝒮¯\(b\)\\mathbf\{H\}\_\{i\}^\{\(b\+1,0\)\}=\\begin\{cases\}\\mathbf\{H\}\_\{i\}^\{\(b,M\_\{b\}\+1\)\}&i\\in\\mathcal\{S\}^\{\(b\)\}\\\\\[4\.0pt\] \\mathbf\{H\}\_\{i\}^\{\(b,0\)\}&i\\in\\bar\{\\mathcal\{S\}\}^\{\(b\)\}\\end\{cases\}\(8\)and importance scores are recomputed fresh at each block’s full attention layer\. This means a single drop decision at layerℓ\\ellimmediately reduces the token count for all layersℓ′\>ℓ\\ell^\{\\prime\}\>\\ell, including subsequent full attention layers, linear attention layers, sliding window layers, and all FFN projections\.
#### FLOPs analysis\.
Letℒdrop=\{ℓ1,ℓ2,…\}\\mathcal\{L\}\_\{\\text\{drop\}\}=\\\{\\ell\_\{1\},\\ell\_\{2\},\\ldots\\\}denote the set of layers at which dropping is applied, and letρk\\rho\_\{k\}denote the retention ratio after thekk\-th drop\. The total FLOPs saved across allLLlayers is:
ΔFLOPs=∑k\(1−ρk\)⋅∑ℓ\>ℓkFLOPsℓ\(N\)\\Delta\\mathrm\{FLOPs\}=\\sum\_\{k\}\(1\-\\rho\_\{k\}\)\\cdot\\sum\_\{\\ell\>\\ell\_\{k\}\}\\mathrm\{FLOPs\}\_\{\\ell\}\(N\)\(9\)For a model withLLtotal layers each of cost𝒪\(Nd2\)\\mathcal\{O\}\(Nd^\{2\}\), a single drop at layerℓ1\\ell\_\{1\}with retention ratioρ\\rhosaves:
ΔFLOPs\(ℓ1\)=\(1−ρ\)⋅\(L−ℓ1\)⋅𝒪\(Nd2\)\\Delta\\mathrm\{FLOPs\}^\{\(\\ell\_\{1\}\)\}=\(1\-\\rho\)\\cdot\(L\-\\ell\_\{1\}\)\\cdot\\mathcal\{O\}\(Nd^\{2\}\)\(10\)This saving scales linearly with\(L−ℓ1\)\(L\-\\ell\_\{1\}\), the number of layers remaining after the drop point\. Sparse attention methods operating only within the attention sublayer save at most\(1−ρ\)⋅𝒪\(N2dk\)\(1\-\\rho\)\\cdot\\mathcal\{O\}\(N^\{2\}d\_\{k\}\)at that layer alone, leaving all subsequent GEMM costs intact\. The ratio of savings is:
ΔFLOPsUniPrefillΔFLOPsSparseAttn=\(L−ℓ1\)⋅Nd2N2dk→N→∞∞\\frac\{\\Delta\\mathrm\{FLOPs\}\_\{\\text\{UniPrefill\}\}\}\{\\Delta\\mathrm\{FLOPs\}\_\{\\text\{SparseAttn\}\}\}=\\frac\{\(L\-\\ell\_\{1\}\)\\cdot Nd^\{2\}\}\{N^\{2\}d\_\{k\}\}\\xrightarrow\{N\\to\\infty\}\\infty\(11\)In the long\-context regime whereN≫dN\\gg d, UniPrefill’s GEMM savings dominate, making it particularly effective precisely at the sequence lengths where prefill acceleration matters most\.
#### Error propagation\.
Assuming each sublayerfmf\_\{m\}isLmL\_\{m\}\-Lipschitz, the accumulated error at block end satisfies:
‖Δ𝐡j\(b,Mb\+1\)‖⩽\(1−p\)⋅Vmax\(b\)⋅∏m=1MbLm\\left\\\|\\Delta\\mathbf\{h\}\_\{j\}^\{\(b,M\_\{b\}\+1\)\}\\right\\\|\\leqslant\(1\-p\)\\cdot V\_\{\\max\}^\{\(b\)\}\\cdot\\prod\_\{m=1\}^\{M\_\{b\}\}L\_\{m\}\(12\)Layer normalization and residual connections constrain∏mLm\\prod\_\{m\}L\_\{m\}in practice, preventing unbounded error amplification across layers\.
### 3\.5Fused Kernel and vLLM Integration
#### Kernel design\.
We implement the importance estimation and top\-ppselection pipeline as a sequence of four fused kernels operating directly on the variable\-length packed token representation indexed bycu\_seqlens, without materializing per\-request tensors or padding\. The pipeline proceeds as follows:
𝐒=𝐐\[N−n:N\]𝐊⊤∈ℝn×N→online softmax𝐨∈ℝN→block reduce𝐛∈ℝ⌈N/G⌉→top\-pℳ∈\{0,1\}N\\mathbf\{S\}\\;=\\;\\mathbf\{Q\}\_\{\[N\-n:N\]\}\\mathbf\{K\}^\{\\top\}\\in\\mathbb\{R\}^\{n\\times N\}\\;\\xrightarrow\{\\text\{online softmax\}\}\\;\\mathbf\{o\}\\in\\mathbb\{R\}^\{N\}\\;\\xrightarrow\{\\text\{block reduce\}\}\\;\\mathbf\{b\}\\in\\mathbb\{R\}^\{\\lceil N/G\\rceil\}\\;\\xrightarrow\{\\text\{top\-\}p\}\\;\\mathcal\{M\}\\in\\\{0,1\\\}^\{N\}\(13\)The partial GEMM kernel computes𝐒\\mathbf\{S\}with tiledQQ\-KKblocking and inline causal masking\. The softmax kernel aggregatessoftmax\(𝐒\)\\operatorname\{softmax\}\(\\mathbf\{S\}\)over thennquery rows via a numerically stable two\-pass online algorithm, yielding per\-token importance scores𝐨\\mathbf\{o\}\. The block\-reduce kernel contracts𝐨\\mathbf\{o\}across both the head and spatial dimensions within each block of sizeGG, producing the block\-level score vector𝐛\\mathbf\{b\}\.
The top\-ppkernel performs sort\-and\-threshold entirely on\-GPU without CPU round\-trips\. We encode each \(score, index\) pair into a singleint64word via a monotone IEEE\-754 bitcast mapping:
φ\(x\)=\{bits\(x\)⊕0x80000000x⩾0bits\(x\)⊕0xFFFFFFFFx<0⇒packed=\(φ\(bg\)≪32\)\|g\\varphi\(x\)=\\begin\{cases\}\\operatorname\{bits\}\(x\)\\oplus\\texttt\{0x80000000\}&x\\geqslant 0\\\\ \\operatorname\{bits\}\(x\)\\oplus\\texttt\{0xFFFFFFFF\}&x<0\\end\{cases\}\\quad\\Rightarrow\\quad\\texttt\{packed\}=\\bigl\(\\varphi\(b\_\{g\}\)\\ll 32\\bigr\)\\;\\big\|\\;g\(14\)Sorting packed words descending, computing a cumulative sum of scores, and thresholding atppyields the keep maskℳ\\mathcal\{M\}, which is scattered back to original positions\. A final expansion kernel liftsℳ\\mathcal\{M\}from block to token granularity, unconditionally settingℳi=1\\mathcal\{M\}\_\{i\}=1for attention\-sink tokensi<Ai<Aand query\-window tokensi⩾N−ni\\geqslant N\-n\.
#### Tensor parallelism\.
Under tensor parallelism of degreeTT, each rank observes only1/T1/Tof the attention heads, yielding a partial block score𝐛\(t\)\\mathbf\{b\}^\{\(t\)\}\. We synchronize via:
𝐛=∑t=1T𝐛\(t\)\\mathbf\{b\}=\\sum\_\{t=1\}^\{T\}\\mathbf\{b\}^\{\(t\)\}\(15\)before the top\-ppkernel, ensuring a consistent drop decision across all TP ranks\.
#### vLLM scheduler integration\.
Integrating token dropping into vLLM’s continuous batching scheduler\[yu2022orca,vLLM\]requires maintaining correctness across three coupled state structures: layer\-wise attention metadata, KV cache slot mappings, and per\-request KV length tracking across decode steps\.
Upon a drop event at layerℓ\\ell, we propagate updated metadata to all downstream layersℓ′\>ℓ\\ell^\{\\prime\}\>\\ellby patchingquery\_start\_loc,seq\_lens, andnum\_actual\_tokensto reflect the compacted token stream\|𝒮\(ℓ\)\|\|\\mathcal\{S\}^\{\(\\ell\)\}\|\. Physical KV cache slot mappings for each layerℓ′\\ell^\{\\prime\}are recomputed as:
sloti\(ℓ′\)=block\_table\(ℓ′\)\[ri,⌊pi/B⌋\]⋅B\+\(pimodB\)\\texttt\{slot\}^\{\(\\ell^\{\\prime\}\)\}\_\{i\}=\\texttt\{block\\\_table\}^\{\(\\ell^\{\\prime\}\)\}\\\!\\left\[r\_\{i\},\\,\\lfloor p\_\{i\}/B\\rfloor\\right\]\\cdot B\+\(p\_\{i\}\\bmod B\)\(16\)wherepip\_\{i\}is the logical position of theii\-th retained token,BBis the KV block size, andblock\_table\(ℓ′\)\\texttt\{block\\\_table\}^\{\(\\ell^\{\\prime\}\)\}is the physical block table of layerℓ′\\ell^\{\\prime\}—which may differ between global and sliding\-window attention layers\[gemmateam2025gemma3technicalreport\]\.
During decode, each layerℓ′\\ell^\{\\prime\}must attend over only the tokens that were physically written to its KV cache during prefill\. We maintain a per\-request drop history\{\(ℓk,skr\)\}\\\{\(\\ell\_\{k\},s\_\{k\}^\{r\}\)\\\}recording the retained sequence lengthskrs\_\{k\}^\{r\}after each drop event at layerℓk\\ell\_\{k\}\. The effective KV length visible to layerℓ′\\ell^\{\\prime\}during decode is then:
sequsedr\(ℓ′\)=sr\(ℓ−\)\+Δr,Δr=kv\_lenr−orig\_lenr\\mathrm\{seqused\}^\{\(\\ell^\{\\prime\}\)\}\_\{r\}\\;=\\;s^\{\(\\ell^\{\-\}\)\}\_\{r\}\+\\Delta\_\{r\},\\qquad\\Delta\_\{r\}=\\mathrm\{kv\\\_len\}\_\{r\}\-\\mathrm\{orig\\\_len\}\_\{r\}\(17\)whereℓ−=max\{ℓk∈ℒdrop:ℓk<ℓ′\}\\ell^\{\-\}=\\max\\\{\\ell\_\{k\}\\in\\mathcal\{L\}\_\{\\mathrm\{drop\}\}:\\ell\_\{k\}<\\ell^\{\\prime\}\\\}is the last drop layer precedingℓ′\\ell^\{\\prime\}, andΔr\\Delta\_\{r\}counts autoregressive tokens appended since prefill\. This per\-layersequsedcorrection is injected into the forward context before each decode step, ensuring every attention layer observes a KV sequence length precisely consistent with its written cache entries—without any modification to model weights or the PagedAttention memory allocator\.
Method4K8K16K32K64K128KAvg4K8K16K32K64K128KLlama\-3\.1\-8B\-Instruct \(Full Attention\)Baseline97\.3695\.9894\.6291\.0286\.2976\.8990\.361\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\timesLazyLLM\[lazyllm\]89\.1681\.1270\.3264\.3956\.2849\.7168\.501\.09×1\.09\\times1\.19×1\.19\\times1\.28×1\.28\\times1\.74×1\.74\\times2\.21×2\.21\\times2\.51×2\.51\\timesSlimInfer\[sliminfer\]90\.2382\.0471\.3967\.1257\.1045\.3668\.871\.07×1\.07\\times1\.16×1\.16\\times1\.25×1\.25\\times1\.66×1\.66\\times1\.98×1\.98\\times2\.07×2\.07\\timesMInference\[minference\]96\.7195\.7895\.5190\.7687\.1278\.2190\.680\.82×0\.82\\times0\.86×0\.86\\times0\.98×0\.98\\times1\.03×1\.03\\times1\.15×1\.15\\times1\.34×1\.34\\timesFlexPrefill\[flexprefill\]96\.3495\.1294\.8388\.9684\.3178\.1389\.620\.83×0\.83\\times0\.89×0\.89\\times1\.02×1\.02\\times1\.08×1\.08\\times1\.24×1\.24\\times1\.46×1\.46\\timesXAttention\[xattention\]95\.9895\.2394\.6888\.0683\.9278\.1689\.340\.92×0\.92\\times0\.96×0\.96\\times1\.03×1\.03\\times1\.08×1\.08\\times1\.21×1\.21\\times1\.38×1\.38\\timesProxyAttn\[wang2025proxyattn\]96\.7895\.4695\.4989\.2885\.3178\.4990\.140\.82×0\.82\\times0\.89×0\.89\\times1\.03×1\.03\\times1\.11×1\.11\\times1\.36×1\.36\\times1\.79×1\.79\\timesUniPrefill96\.5395\.8395\.4189\.7785\.2879\.8790\.451\.21×1\.21\\times1\.34×1\.34\\times1\.37×1\.37\\times1\.62×1\.62\\times2\.01×2\.01\\times2\.26×2\.26\\timesQwen3\-Next\-80B\-A3B \(Linear/Full Attention Hybrid\)Baseline96\.8395\.6795\.0794\.3894\.5192\.0994\.761\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\timesLazyLLM\[lazyllm\]89\.1382\.3770\.6964\.3758\.1255\.1769\.981\.11×1\.11\\times1\.18×1\.18\\times1\.29×1\.29\\times1\.40×1\.40\\times1\.55×1\.55\\times1\.74×1\.74\\timesSlimInfer\[sliminfer\]88\.1180\.3667\.1363\.2257\.1355\.3668\.551\.14×1\.14\\times1\.22×1\.22\\times1\.27×1\.27\\times1\.34×1\.34\\times1\.42×1\.42\\times1\.56×1\.56\\timesMInference\[minference\]96\.6294\.3894\.4994\.2794\.2891\.8194\.310\.96×0\.96\\times0\.98×0\.98\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.02×1\.02\\times1\.05×1\.05\\timesFlexPrefill\[flexprefill\]96\.3695\.0394\.1793\.9192\.8991\.4493\.970\.96×0\.96\\times0\.98×0\.98\\times1\.00×1\.00\\times1\.01×1\.01\\times1\.04×1\.04\\times1\.08×1\.08\\timesXAttention\[xattention\]96\.0394\.8194\.0393\.0193\.0690\.2393\.530\.97×0\.97\\times0\.99×0\.99\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.02×1\.02\\times1\.05×1\.05\\timesProxyAttn\[wang2025proxyattn\]96\.3194\.1394\.2393\.4793\.5191\.6293\.880\.96×0\.96\\times0\.98×0\.98\\times1\.00×1\.00\\times1\.02×1\.02\\times1\.05×1\.05\\times1\.11×1\.11\\timesUniPrefill96\.6794\.4994\.2993\.6393\.1391\.4193\.941\.08×1\.08\\times1\.21×1\.21\\times1\.24×1\.24\\times1\.39×1\.39\\times1\.42×1\.42\\times1\.68×1\.68\\timesGemma\-3\-12B \(Sliding Window/Full Attention Hybrid\)Baseline94\.0189\.1285\.9880\.7668\.8961\.2279\.991\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.00×1\.00\\timesLazyLLM\[lazyllm\]86\.1180\.3475\.2368\.4254\.1243\.3867\.931\.23×1\.23\\times1\.32×1\.32\\times1\.37×1\.37\\times1\.42×1\.42\\times1\.53×1\.53\\times1\.64×1\.64\\timesSlimInfer\[sliminfer\]88\.2383\.1479\.1269\.3353\.0240\.1168\.831\.15×1\.15\\times1\.24×1\.24\\times1\.32×1\.32\\times1\.39×1\.39\\times1\.45×1\.45\\times1\.52×1\.52\\timesMInference\[minference\]93\.5689\.5186\.0180\.0467\.0959\.3179\.250\.98×0\.98\\times0\.99×0\.99\\times1\.00×1\.00\\times1\.00×1\.00\\times1\.01×1\.01\\times1\.03×1\.03\\timesFlexPrefill\[flexprefill\]93\.6389\.1685\.4979\.2365\.6958\.6378\.640\.98×0\.98\\times0\.99×0\.99\\times1\.00×1\.00\\times1\.01×1\.01\\times1\.02×1\.02\\times1\.04×1\.04\\timesXAttention\[xattention\]93\.0689\.2485\.6779\.1466\.1856\.2478\.260\.99×0\.99\\times1\.00×1\.00\\times1\.01×1\.01\\times1\.01×1\.01\\times1\.02×1\.02\\times1\.02×1\.02\\timesProxyAttn\[wang2025proxyattn\]93\.6789\.5285\.3178\.9765\.3159\.9378\.790\.98×0\.98\\times0\.99×0\.99\\times1\.01×1\.01\\times1\.01×1\.01\\times1\.03×1\.03\\times1\.06×1\.06\\timesUniPrefill93\.1889\.7686\.4779\.0866\.3258\.3878\.871\.15×1\.15\\times1\.21×1\.21\\times1\.22×1\.22\\times1\.26×1\.26\\times1\.31×1\.31\\times1\.49×1\.49\\times
Table 1:Performance vs\. efficiency across different models and methods\. Evaluation scores \(left\) and TTFT speedup relative to baselines \(right\) are reported\. For a fair comparison, all models are evaluated using HuggingFace Transformers with batch size equals to 1\.BSZ4K8K16K32K64K128K4K8K16K32K64K128KLlama\-3\.1\-8B\-Instruct \(Full Attention\)136984435864369739027303242101338522\(\+4%\+4\\%\)48932\(\+12%\+12\\%\)52314\(\+20%\+20\\%\)59984\(\+54%\+54\\%\)54786\(\+81%\+81\\%\)43672\(\+107%\+107\\%\)446372486324543640210308122105449336\(\+6%\+6\\%\)59372\(\+22%\+22\\%\)61148\(\+35%\+35\\%\)64113\(\+59%\+59\\%\)56108\(\+82%\+82\\%\)43698\(\+108%\+108\\%\)1653764526434822140671308342106263762\(\+19%\+19\\%\)70213\(\+33%\+33\\%\)66762\(\+38%\+38\\%\)66512\(\+64%\+64\\%\)57678\(\+87%\+87\\%\)44042\(\+109%\+109\\%\)6455431535414860340869——68721\(\+24%\+24\\%\)72139\(\+35%\+35\\%\)67361\(\+39%\+39\\%\)67618\(\+65%\+65\\%\)——Qwen3\-Next\-80B\-A3B \(Linear/Full Attention Hybrid\)115314265343571240312398073351214621\(−5%\-5\\%\)25364\(−4%\-4\\%\)36912\(\+3%\+3\\%\)48972\(\+21%\+21\\%\)50324\(\+26%\+26\\%\)49732\(\+48%\+48\\%\)426334384424332646534403223336425498\(−3%\-3\\%\)38872\(\+1%\+1\\%\)50442\(\+16%\+16\\%\)60894\(\+31%\+31\\%\)55136\(\+37%\+37\\%\)52442\(\+57%\+57\\%\)1647432518535154448303416933348948446\(\+2%\+2\\%\)62468\(\+20%\+20\\%\)62983\(\+22%\+22\\%\)66798\(\+38%\+38\\%\)57938\(\+39%\+39\\%\)56398\(\+68%\+68\\%\)6452334539365393647932——58936\(\+13%\+13\\%\)67848\(\+26%\+26\\%\)68123\(\+26%\+26\\%\)68631\(\+43%\+43\\%\)——Gemma\-3\-12B \(Sliding Window/Full Attention Hybrid\)119013212032473324763223671810318673\(−2%\-2\\%\)23132\(\+9%\+9\\%\)29436\(\+19%\+19\\%\)31023\(\+25%\+25\\%\)27384\(\+22%\+22\\%\)25673\(\+42%\+42\\%\)423531243612701326432225361840326783\(\+14%\+14\\%\)31432\(\+29%\+29\\%\)33014\(\+22%\+22\\%\)33432\(\+26%\+26\\%\)29232\(\+30%\+30\\%\)25932\(\+41%\+41\\%\)1629468298672903127012226131851334512\(\+17%\+17\\%\)37136\(\+24%\+24\\%\)35674\(\+23%\+23\\%\)34419\(\+27%\+27\\%\)29732\(\+31%\+31\\%\)26231\(\+42%\+42\\%\)6429471303242947127236——35016\(\+19%\+19\\%\)37365\(\+23%\+23\\%\)35912\(\+22%\+22\\%\)34578\(\+27%\+27\\%\)——
Table 2:Prefill throughput \(tokens/s\) of Standard Prefill and UniPrefill measured within vLLM \(TP=8=8\) across three model architectures, six context lengths \(4K–128K\), and four batch sizes \(BSZ\)\. The left half of each group reports Standard Prefill throughput; the right half reports UniPrefill throughput\.
## 4Experiments
We evaluate UniPrefill across two dimensions: accuracy and efficiency\. For accuracy, we compare UniPrefill against existing prefill acceleration methods on the RULER\[hsieh2024ruler\]long\-context benchmark across multiple model architectures\. For efficiency, we measure prefill throughput under varying context lengths and batch sizes within our vLLM deployment\. Finally, we conduct ablation studies to analyze the contribution of each design choice in UniPrefill\. Implementation and deployment details can be found inappendix\.
### 4\.1Experimental Setup
We select three model architectures to validate the effectiveness of UniPrefill: LLaMA\-3\.1\-8B\-Instruct\[llama3\], which consists entirely of full\-attention layers; Qwen3\-Next\-80B\-A3B\[qwen3next\_blog\_2025\], a linear/full\-attention hybrid with a 3:1 ratio; and Gemma\-3\-12B\[gemmateam2025gemma3technicalreport\], a sliding\-window/full\-attention hybrid with a 5:1 ratio\. We set the top\-ppthreshold to0\.990\.99,0\.990\.99, and0\.980\.98for the three models, respectively\. The minimum dropping granularity is set to a block size ofG=64G=64tokens, and importance scores are estimated using the lastn=128n=128query tokens\. To preserve attention sinks\[xiao2023streamingllm\], the first 128 tokens are always retained\.
### 4\.2Results on RULER
RULER\[hsieh2024ruler\]is a comprehensive long\-context benchmark that evaluates LLMs across diverse task categories including retrieval, multi\-hop tracing, aggregation, and question answering, with configurable context lengths up to 128K tokens\. Unlike prior benchmarks that rely on simple needle\-in\-a\-haystack tests, RULER provides a more rigorous and systematic assessment of true long\-context understanding, making it a widely adopted standard for evaluating long\-context LLM performance\.
Tab\.[1](https://arxiv.org/html/2605.06221#S3.T1)presents RULER scores and TTFT speedups across three model architectures\. UniPrefill achieves the best accuracy\-efficiency tradeoff among all acceleration methods\. LazyLLM and SlimInfer suffer notable accuracy degradation across all three architectures, while sparse attention methods preserve accuracy but yield diminishing speedups on hybrid architectures, with gains often below1\.1×1\.1\\timesat 128K\. UniPrefill strikes the optimal balance: it retains accuracy close to the Baseline while delivering up to2\.26×2\.26\\times,1\.68×1\.68\\times, and1\.49×1\.49\\timesTTFT speedup at 128K context length on LLaMA\-3\.1\-8B, Qwen3\-Next\-80B\-A3B, and Gemma\-3\-12B, respectively, demonstrating consistent effectiveness across full\-attention and hybrid architectures\.
### 4\.3vLLM Intergration
Tab\.[2](https://arxiv.org/html/2605.06221#S3.T2)reports prefill throughput within vLLM across three architectures\. UniPrefill consistently improves throughput as context length and batch size increase, achieving up to\+109%\+109\\%,\+68%\+68\\%, and\+42%\+42\\%gains on LLaMA\-3\.1\-8B, Qwen3\-Next\-80B\-A3B, and Gemma\-3\-12B, respectively\. The speedup scales favorably with both context length and batch size, demonstrating that UniPrefill is particularly effective in the high\-concurrency, long\-context regime that dominates production serving workloads\.
### 4\.4Ablation Study
#### Block Size\.
Tab\.[3](https://arxiv.org/html/2605.06221#S4.T3)presents the ablation results for block sizeG∈\{32,64,128\}G\\in\\\{32,64,128\\\}\. At short context lengths,G=128G=128yields the highest speedup due to lower selection overhead per drop decision\. As context length grows,G=32G=32surpassesG=128G=128in speedup, since finer granularity allows more tokens to be dropped, achieving up to\+121%\+121\\%and\+78%\+78\\%throughput gain on LLaMA\-3\.1\-8B and Qwen3\-Next\-80B\-A3B at 128K, respectively\. We adoptG=64G=64as the default, which balances selection overhead and drop rate across all context lengths\.
block size4K8K16K32K64K128KAvg4K8K16K32K64K128KLlama\-3\.1\-8B\-Instruct6496\.5395\.8395\.4189\.7785\.2879\.8790\.45\+19%\+19\\%\+33%\+33\\%\+38%\+38\\%\+64%\+64\\%\+87%\+87\\%\+109%\+109\\%12894\.3293\.6293\.0788\.1283\.3878\.9088\.57\+26%\+26\\%\+38%\+38\\%\+45%\+45\\%\+62%\+62\\%\+81%\+81\\%\+96%\+96\\%3293\.4294\.6395\.4690\.2385\.6779\.8889\.88\+19%\+19\\%\+32%\+32\\%\+36%\+36\\%\+72%\+72\\%\+96%\+96\\%\+121%\+121\\%Qwen3\-Next\-80B\-A3B6496\.6794\.4994\.2993\.6393\.1391\.4193\.94\+2%\+2\\%\+20%\+20\\%\+22%\+22\\%\+38%\+38\\%\+39%\+39\\%\+68%\+68\\%12896\.5294\.6994\.1393\.4192\.6792\.0693\.91\+5%\+5\\%\+22%\+22\\%\+23%\+23\\%\+34%\+34\\%\+37%\+37\\%\+56%\+56\\%3292\.1794\.8894\.7293\.8993\.6692\.6893\.67\+0%\+0\\%\+19%\+19\\%\+22%\+22\\%\+44%\+44\\%\+51%\+51\\%\+78%\+78\\%
Table 3:Ablation study of block sizeGG\. The left panel reports RULER scores under different values ofGG, and the right panel reports the corresponding TTFT speedup of UniPrefill relative to the Baseline\.lastnn4K8K16K32K64K128KAvgLlama\-3\.1\-8B\-Instruct12896\.5395\.8395\.4189\.7785\.2879\.8790\.453295\.3294\.1393\.1886\.2282\.6375\.1387\.7751296\.7296\.0495\.6390\.2384\.9679\.3890\.49
Table 4:Ablation study of lastnn\. RULER scores under different values of thennon LLaMA\-3\.1\-8B\-Instruct\.n=128n=128is adopted as the default\.
#### Lastnn\.
Tab\.[4](https://arxiv.org/html/2605.06221#S4.T4)reports RULER scores under different values of lastn∈\{32,128,512\}n\\in\\\{32,128,512\\\}\.n=32n=32leads to a noticeable accuracy drop, as too few query tokens introduce high variance in importance estimation\.n=512n=512recovers accuracy but incurs higher computational overhead\.n=128n=128achieves the best balance and is adopted as the default\.
## 5Conclusion
We present UniPrefill, an architecture\-agnostic framework for long\-context LLM prefill acceleration\. By estimating token importance via block\-wise top\-ppselection at full\-attention layers and propagating the sparsity mask across all subsequent sub\-layers, UniPrefill simultaneously reduces attention and GEMM FLOPs, making it effective across full\-attention and hybrid architectures alike\. We implement UniPrefill as a fused kernel pipeline and integrate it into vLLM’s continuous\-batching scheduler without any model weight changes\. Experiments on the RULER benchmark show that UniPrefill achieves the best accuracy\-efficiency tradeoff among all compared methods, delivering up to2\.1×2\.1\\timesTTFT speedup with negligible accuracy loss, with gains scaling favorably with context length and batch size\. We hope UniPrefill provides a practical and general solution for efficient long\-context LLM serving\.
## References
## Appendix AImplementation and deployment details
UniPrefill is implemented and evaluated on top of vLLM v0\.16\.0\[vLLM\], with full support for prefill\-decode co\-processing and tensor parallelism, making it compatible with standard production deployment configurations\. All throughput experiments are conducted under tensor parallelism degreeTP=8\\text\{TP\}=8, reflecting a typical multi\-GPU serving setup in real\-world scenarios\. The continuous\-batching operator is implemented as a set of fused Triton kernels, which are hardware\-agnostic by design and theoretically portable across different accelerator platforms\. All experiments are conducted under CUDA 12\.8\.
## Appendix BExperiment statistical significance\.
Tab\.[5](https://arxiv.org/html/2605.06221#A2.T5)reports results across multiple random seeds, and the consistently stable performance demonstrates that UniPrefill is robust to random seed initialization\.
random seed4K8K16K32K64K128K096\.5395\.8395\.4189\.7785\.2879\.8732196\.5395\.8395\.4189\.7785\.2879\.87346796\.5395\.8395\.4189\.7785\.2879\.87Table 5:Ablation study on different random seed\.
## Appendix CLimitations and Broader Impacts
This work focuses on accelerating the prefill phase for long\-context LLM inference\. While UniPrefill delivers consistent speedups across diverse architectures and sequence lengths without observable accuracy degradation, extending the framework to broader inference optimization dimensions—such as decoding acceleration or training\-time efficiency—remains an interesting direction for future work\. Broader societal considerations, including ethical deployment and safety alignment, are important but lie outside the technical scope of this study\.Similar Articles
Prefilling-dLLM: Predictive Prefilling for Long-Context Inference in Diffusion Language Models
This paper proposes Prefilling-dLLM, a training-free framework that partitions the prefix into chunks and caches KV representations, achieving state-of-the-art quality and up to 28x speedup for long-context inference in diffusion language models.
CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection
CompactAttention introduces Block-Union KV Selection to accelerate chunked prefill for long-context LLMs, achieving up to 2.72x attention speedup on LLaMA-3.1-8B at 128K context while maintaining accuracy close to dense attention.
PreFT: Prefill-only finetuning for efficient inference
PreFT proposes applying adapters only to prefill tokens, discarding them during decode, which increases throughput for multi-adapter serving with minimal performance loss.
@_avichawla: Prefill & decode in LLM inference. Have you ever noticed that the first token from an LLM always takes a moment to appe…
Explains the two phases of LLM inference - prefill and decode - detailing how GPU bottlenecks shift from compute-bound during prefill to memory-bound during decode, and the importance of KV caching.
Prefill vs. decoding and local LLM ROI: is prefill underrated?
An analysis comparing prefill and decoding phases in LLM inference, questioning whether prefill is underappreciated in terms of ROI for local LLM deployments.