PAT: Accelerating LLM Decoding via Prefix-Aware Attention with Resource Efficient Multi-Tile Kernel

When a language model generates text one token at a time, each new token must read the entire stored history of keys and values (the KV cache) from GPU memory, so decoding is bottlenecked by memory bandwidth, not compute — and because many requests share the same opening text (system prompts, tool descriptions, retrieved documents), the shared portion of that history gets loaded from memory over and over, once per request; this paper (Tianjin University and collaborators — Jinjun Yi, Zhixin Zhao, Yitao Hu, Ke Yan, Weiwei Sun, Hao Wang, Laiping Zhao, Yuhao Zhang, Wenxin Li, Keqiu Li) presents PAT, an attention kernel that groups the requests sharing a prefix so its KV cache is loaded once and reused, runs a multi-tile kernel whose tile size adapts per group to keep the GPU's on-chip resources busy, and merges the partial results with a cheap online softmax — cutting attention latency by 53.5 percent on average as a drop-in plugin for vLLM.

The argument in one glance

Problem

Generating each token means re-reading a huge KV cache, and decoding is bound by memory

A large language model generates text autoregressively: after processing the prompt, it emits one token at a time, and each new token must attend over the keys and values of every token that came before, stored in the KV cache. That cache is large, and the attention step for each token has very little arithmetic per byte, so it is memory-bound — its speed is set by how fast the GPU can load the KV cache from its global memory, not by how fast it can compute. As requests get longer and batches get bigger, this decode attention comes to dominate the cost of serving a model. Anything that reduces the memory traffic of decode attention, or keeps the GPU's resources busier while it waits on memory, directly speeds up serving.
What people did

Give each query its own thread block with a fixed tile size

The standard attention kernels for decoding, like FlashDecoding and FlashInfer, assign one query to one CTA — one cooperative thread array, the GPU's unit of scheduled work — and each CTA loads the KV cache it needs and computes that query's attention with a single, fixed tile size chosen once for all CTAs. This is simple to schedule and works well when every request is independent and of similar length.
Why it fell short

Shared prefixes get reloaded per query, and one tile size wastes the GPU

Real workloads violate both assumptions. Requests share substantial, hierarchical prefixes — the same system prompt, the same tool and template definitions, the same retrieved documents in a RAG pipeline — so many queries need the exact same leading KV cache blocks. But with one query per CTA, each CTA independently loads that shared prefix from slow global memory, so the same KV data is fetched over and over, once per query; the L2 cache absorbs only a fraction because of its limited size and bandwidth. Separately, a single one-size-fits-all tile leaves on-chip resources idle for some CTAs and, when KV lengths are uneven across requests, produces execution bubbles where parts of the GPU stall waiting for the longest tile to finish. Both problems amplify memory-bandwidth pressure on an operation that is already memory-bound, so decode attention stalls.
This paper's idea

Pack queries by shared prefix, adapt the tile, then merge the pieces

PAT is a prefix-aware attention kernel for decoding organized as a pack-forward-merge pipeline. Pack: a scheduler builds a prefix tree of the decode batch (queries that share leading KV blocks share a path in the tree), scores candidate groupings with a memory-centric profit model, and packs queries that share a prefix into the same CTA so the shared KV cache is loaded once and reused by all of them, instead of once per query. Forward: a customized multi-tile kernel runs the packed work with an online tile-size selector — a small decision tree that picks the per-CTA tile size to fit each pack's shape — and adds multi-stream forwarding and long-KV splitting to keep the GPU busy and cut the bubbles that uneven KV lengths cause. Merge: because packing and splitting compute partial attention outputs, a final step recombines them with an online softmax at negligible overhead. Built as an off-the-shelf plugin for vLLM, PAT reduces attention latency by 53.5 percent on average and lowers time-per-output-token by 17.0 to 93.1 percent against state-of-the-art kernels.

Why this problem is worth solving

Decode attention is where the time and money of serving a language model increasingly go, so shaving its cost is directly valuable.

A transformer LLM serves a request in two phases. Prefill reads the whole prompt at once, which is compute-heavy but happens once. Decode then generates the answer one token at a time, and this phase usually dominates, because a long answer means many decode steps. At each decode step the new token’s query must attend over the keys and values of all previous tokens, which are kept in the KV cache — a structure that grows with the sequence and, across a batch of requests, becomes very large. The crucial fact is the arithmetic intensity: decode attention does only a little computation for each element of KV it reads, so it is memory-bound. Its speed is dictated by how fast the GPU can stream the KV cache out of its global memory (HBM), and the compute units spend much of their time waiting. This is why serving throughput and latency hinge so heavily on the memory traffic of decode attention: with the arithmetic essentially free, every redundant or avoidable byte loaded from global memory is a direct tax on speed, and every cycle the memory system is not delivering useful KV is wasted. Reducing that traffic, and hiding the unavoidable part behind useful work, is therefore the highest-leverage thing you can do for LLM serving, which is exactly the target PAT aims at.

What makes the opportunity real is that the redundant traffic is not random noise but highly structured sharing.

In production, LLM requests are far from independent. They almost always begin with the same boilerplate: a long system prompt setting the assistant’s behavior, definitions of the tools or functions the model may call, few-shot examples, or, in retrieval-augmented generation, the same retrieved passages prepended to many questions. Because the model processes tokens left to right, two requests that start with identical text produce identical keys and values for those leading tokens — so their KV cache blocks for the shared prefix are literally the same content. This sharing is also hierarchical: a group of requests may share a system prompt, and subgroups within it may further share a tool definition or a document, forming a tree of common prefixes. That structure is what makes the redundancy exploitable rather than merely regrettable. If the shared prefix’s KV were loaded from global memory once and reused by every request that shares it, the memory traffic of the most expensive part of decode attention would fall dramatically, because the shared portion is often the bulk of the sequence. The challenge is that this reuse must be organized inside the attention kernel and on the GPU’s scheduling model, without wasting on-chip resources or unbalancing the work — which is precisely the gap between the opportunity being visible and it being captured, and where PAT does its engineering.

What people did before, and where it broke

Query-centric kernels (FlashDecoding, FlashInfer). Assign one query per CTA and load whatever KV each needs, with a fixed tile size. Where it broke. Every CTA independently reloads the shared prefix from global memory, so the same KV is fetched once per query; the L2 cache only partly helps, so redundant memory traffic dominates a memory-bound operation.
One-size-fits-all tiling. Choose a single tile shape for all CTAs to keep the kernel simple. Where it broke. A fixed tile leaves on-chip resources idle for some work and, when KV lengths vary across requests, creates execution bubbles as parts of the GPU stall waiting for the longest tile, underusing the hardware.
KV-centric packing (RelayAttention). Pack shared KV so a first-level shared prefix is loaded once into a CTA and shared across queries, cutting redundant loads. Where it broke. Organizing purely around the KV can cause memory waste and does not flexibly handle deep, hierarchical sharing or uneven lengths, so resources are still used inefficiently.
Compute-oriented packing (FastTree). Use a cost model to pack and reduce repeated KV loads, tuned toward compute. Where it broke. For a memory-bound operation the right objective is to minimize global memory accesses, not to optimize compute; a compute-oriented model leaves memory traffic on the table.
The gap. No kernel packed queries by shared prefix under a memory-centric objective while also adapting the tile size per group and rebalancing uneven work — capturing the hierarchical prefix reuse and keeping the GPU’s on-chip resources busy at the same time.

How PAT works

A pack-forward-merge pipeline for prefix-aware decode attention

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Load the shared history once and let every request that shares it reuse the load. The organizing move in PAT is to notice that the expensive, memory-bound part of decode attention — streaming the KV cache from global memory — is being paid for redundantly, because the standard one-query-per-CTA layout makes each request reload the shared prefix independently. PAT instead packs the queries that share a prefix into the same unit of GPU work, so the shared KV is fetched from global memory a single time and reused by all of them. On an operation whose speed is set by memory traffic, deduplicating the largest shared component of that traffic is the highest-value change available. Arranging data and work so that a single expensive memory access serves many consumers, rather than each consumer paying to fetch the same bytes, is the same principle the work on placing co-activated neurons together for on-device inference covered elsewhere on this site is built on, co-locating what is used together so one load does the work of many. PAT applies that principle to shared prefixes in the KV cache, at the granularity of GPU thread blocks.

Match the kernel’s tiling to the work instead of fixing it in advance. Packing alone is not enough, because packed groups differ in shape and KV length, and a single fixed tile size cannot serve them all without wasting on-chip resources or stalling on the longest piece. PAT therefore makes tiling adaptive: a multi-tile kernel with an online tile-size selector chooses the per-CTA tile to fit each pack, and multi-stream forwarding plus long-KV splitting rebalance uneven work so the GPU’s many units stay busy rather than idling in bubbles. Tuning how a computation is broken into tiles and mapped onto a GPU’s on-chip resources, so the hardware is kept fully occupied rather than run with one rigid schedule, is the same kind of GPU-aware kernel optimization the work on an intermediate representation for optimized tensor computations on GPUs covered elsewhere on this site pursues, shaping the mapping to the machine. PAT does it online and per pack, because the right tiling depends on the batch’s prefix structure, which is only known at run time.

Keep the whole thing balanced and cheap, so the savings survive scheduling. A memory win can evaporate if the packing unbalances the GPU or the recombination is expensive, so PAT treats scheduling and merging as first-class. Its memory-centric profit model chooses packings that minimize global accesses rather than optimizing compute, its multi-stream and splitting logic smooths the load across the GPU, and its final online-softmax merge recombines the partial outputs with negligible overhead. The result is that the raw memory savings from prefix reuse actually translate into end-to-end speed. Scheduling many pieces of work onto shared GPU hardware so resources are used efficiently and no part sits idle, turning a raw efficiency into real serving gains, is the same concern the work on QoS-aware, resource-efficient multi-stage GPU serving covered elsewhere on this site is built on, balancing work across the accelerator. PAT carries that balancing into the attention kernel itself, which is why it can ship as a drop-in plugin and still deliver its full benefit.

Following a decode batch through PAT

Serving one batch of decode queries with prefix-aware attention

A trace, drawn from the paper's description, of how PAT executes one decode step for a batch of requests that share prefixes.

Step 1 of 5

A batch of queries arrives with their KV block tables

Each request in the decode batch contributes one query and a block table listing the KV cache blocks it must attend over. Requests that began with the same system prompt, tool definitions, or retrieved documents have identical leading block IDs, marking a shared prefix.

The shared prefix is often the bulk of the sequence, so its KV traffic dominates.

Why it works

Deduplicating the shared prefix attacks the true bottleneck. Decode attention is memory-bound and the shared prefix is the largest common part of the traffic → but packing shared-prefix queries loads that KV once instead of once per query → so the dominant source of redundant memory traffic disappears.
A memory-centric objective aims at the right target. Compute-oriented packing optimizes the wrong resource for a memory-bound op → but PAT’s profit model minimizes global memory accesses directly → so the chosen packings actually reduce the traffic that limits speed.
Adaptive tiling keeps on-chip resources busy. A fixed tile leaves some CTAs’ on-chip memory and bandwidth idle → but an online selector picks the per-CTA tile to fit each pack → so more of the GPU’s resources do useful work per step.
Multi-stream forwarding and KV splitting erase bubbles. Uneven KV lengths make parts of the GPU stall waiting on the longest tile → but splitting long KV and forwarding on multiple streams rebalances the work → so the execution bubbles that would idle the hardware are filled.
A cheap online-softmax merge preserves the win. Packing and splitting yield partial results that must be recombined correctly → but online softmax merges them with negligible overhead → so the memory savings survive as end-to-end latency reductions rather than being eaten by the merge.

What it costs (trade-offs and scope)

The gains depend on there being shared prefixes. PAT’s central win is reusing shared-prefix KV loads, so its benefit scales with how much prefix sharing a workload actually has. Batches of wholly independent requests with little common prefix offer little to deduplicate, and the packing machinery then mostly adds scheduling work without the memory payoff.
Packing and tile selection add run-time scheduling cost. Building the prefix tree, scoring packings with the profit model, and choosing tile sizes online all happen per decode batch. The paper reports this overhead is small and the merge negligible, but it is real work that must stay cheap relative to the attention it accelerates, especially for small batches.
It optimizes decode attention, not the whole model. PAT speeds up the attention kernel during decoding; it does not change prefill, the feed-forward layers, or other serving costs, so end-to-end gains depend on how much of serving time is decode attention for a given workload.
The improvements are relative to specific baselines and settings. The 53.5 percent average attention-latency reduction and 17.0 to 93.1 percent time-per-output-token reduction are measured against particular state-of-the-art kernels under matched decode batches or request rates; other models, sequence lengths, batch sizes, and GPUs would shift the numbers.
It is a kernel-and-scheduler plugin, tied to the serving stack. PAT is implemented as an off-the-shelf plugin for vLLM and reasons about the GPU’s CTA and tiling model; benefiting elsewhere requires integrating the same prefix-aware packing and multi-tile execution into that stack, and the exact tuning is GPU-specific.

The evaluation

Running decode attention for a batch with shared prefixes

Property	Query-centric (fixed tile)	KV-centric packing	PAT (prefix-aware)
Shared-prefix KV loads
Packing objective
Tile size
Handling uneven KV lengths
Reported outcome

A qualitative comparison of decode-attention execution strategies. PAT reduces attention latency by 53.5 percent on average and lowers time-per-output-token by 17.0 to 93.1 percent versus state-of-the-art kernels (FlashDecoding/FlashInfer, RelayAttention, FastTree), as a drop-in plugin for vLLM.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s abstract and its arXiv version (arXiv:2511.22333), which describe the mechanism in full: decode attention is memory-bound because each token loads the massive KV cache from global memory; real workloads have substantial hierarchical shared prefixes (system prompts, tools/templates, RAG); existing kernels waste this because one-query-per-CTA execution reloads shared prefix KV and one-size-fits-all tiling leaves on-chip resources idle and causes bubbles for uneven KV lengths. PAT uses a pack-forward-merge paradigm: a pack scheduler that builds a prefix tree, scores packings with a memory-centric profit model, and packs queries by shared prefix into memory-optimized CTAs; a multi-tile kernel with an online tile-size selector (a decision tree) plus multi-stream forwarding and long-KV splitting to reduce bubbles; and an output merge using online softmax with negligible overhead. It is implemented as an off-the-shelf plugin for vLLM. The headline figures are the paper’s own: attention latency reduced by 53.5 percent on average under matched decode batches, and time-per-output-token lowered by 17.0 to 93.1 percent under matched request rates, versus state-of-the-art attention kernels. Exact profit-model details, tile-selection specifics, and per-workload results should be confirmed in the paper; source code is at github.com/flashserve/PAT.

Background

Why LLM serving keeps chasing the KV cache

It helps to see PAT as part of a wave of work aimed squarely at the KV cache, which has become the central object of LLM serving efficiency. The reason is structural: as models are used for longer conversations, larger retrieved contexts, and agentic tool use, the KV cache grows without bound, and decode attention — which must stream that cache for every generated token — turns into a memory-bandwidth problem rather than a compute one. So the field has attacked the KV cache from many angles: compressing it, quantizing it, evicting parts of it, paging it between memory tiers, and, as here, avoiding redundant loads of the parts that many requests share. What unites these efforts is the recognition that in modern serving the bottleneck has migrated from the arithmetic of the model to the movement of its state, so the wins come from managing that state cleverly. Prefix sharing is an especially clean instance because the redundancy is exact — identical leading tokens produce identical KV — and pervasive in production prompts. The broader push to serve LLMs at low latency and high throughput under real, dynamic workloads is the same goal the work on shift parallelism for low-latency, high-throughput LLM inference covered elsewhere on this site pursues, reorganizing execution to fit how requests actually arrive. PAT contributes the kernel-level piece: making the attention computation itself prefix-aware and resource-efficient.

If you remember only three things

1. Serving a language model is increasingly dominated by decode attention, which is memory-bound: generating each token requires streaming the massive KV cache from global GPU memory, so speed is limited by memory bandwidth, not compute. Real workloads share substantial, hierarchical prefixes across requests (system prompts, tool templates, RAG documents), but standard kernels waste this because one-query-per-CTA execution reloads the shared prefix once per query, and a fixed one-size-fits-all tile leaves on-chip resources idle and creates execution bubbles when KV lengths differ.
2. PAT is a prefix-aware attention kernel organized as pack-forward-merge. It packs queries that share a prefix into the same CTA (guided by a prefix tree and a memory-centric profit model that minimizes global memory accesses) so the shared KV is loaded once and reused; it runs a multi-tile kernel whose per-CTA tile size is chosen online, with multi-stream forwarding and long-KV splitting to keep the GPU busy and cut bubbles; and it merges the resulting partial outputs with a cheap online softmax.
3. Implemented as a drop-in plugin for vLLM, PAT reduces attention latency by 53.5 percent on average under matched decode batches and lowers time-per-output-token by 17.0 to 93.1 percent under matched request rates, versus state-of-the-art attention kernels. Its gains depend on workloads actually having shared prefixes, add modest run-time scheduling cost, and target decode attention specifically rather than the whole model.

Check yourself

Check your understanding

1. Why is decode attention the bottleneck in LLM serving, and what makes it wasteful today?
2. What is PAT's core idea?
3. How does PAT keep the GPU busy when requests have uneven KV lengths?
4. What did PAT achieve, and what are its main limitations?

Sources

The paper is “PAT: Accelerating LLM Decoding via Prefix-Aware Attention with Resource Efficient Multi-Tile Kernel” by Jinjun Yi, Zhixin Zhao, Yitao Hu, Ke Yan, Weiwei Sun, Hao Wang, Laiping Zhao, Yuhao Zhang, Wenxin Li, and Keqiu Li (Tianjin University and collaborators), published at ASPLOS 2026 (Volume 2), DOI 10.1145/3779212.3790200. It is licensed CC BY.

This explainer is written from the paper’s abstract and its arXiv preprint (arXiv:2511.22333), which describe the mechanism in full: LLM serving is dominated by decode attention, a memory-bound operation because each token loads the massive KV cache from global memory; real workloads have substantial, hierarchical shared prefixes (system prompts, tools/templates, RAG); existing implementations fail to exploit this, since one-query-per-CTA execution repeatedly loads shared prefix KV, while one-size-fits-all tiling leaves on-chip resources idle and exacerbates bubbles for uneven KV lengths. PAT organizes execution with a pack-forward-merge paradigm: a pack scheduler builds a prefix tree, scores candidate packings with a memory-centric profit model that minimizes global memory accesses, and packs queries by shared prefix into memory-optimized CTAs; a multi-tile kernel uses an online tile-size selector (a decision tree that adapts the per-CTA tile dimensions) together with multi-stream forwarding and long-KV splitting to reduce resource bubbles; and an output merge performs online softmax with negligible overhead. PAT is implemented as an off-the-shelf plugin for vLLM. The reported results are the paper’s own: attention latency reduced by 53.5 percent on average under the same decode batches, and time-per-output-token lowered by 17.0 to 93.1 percent under the same request rate, against state-of-the-art attention kernels. No numbers beyond the paper’s have been introduced; the source code is at github.com/flashserve/PAT, and exact profit-model, tile-selection, and per-workload details should be confirmed in the paper.