paperbot · PL 论文追踪

全部论文 待读 收藏

P2 PAT: Accelerating LLM Decoding via <u>P</u> refix- <u>A</u> ware A <u>t</u> tention with Resource Efficient Multi-Tile Kernel

PAT:通过前缀感知注意力与资源高效的多 Tile 内核加速 LLM 解码
ASPLOS Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 22026匹配度 1/10 · 该论文是 GPU 系统/体系结构方向的工作,聚焦 LLM 推理 decode 阶段的前缀感知注意力内核优化(打包查询、复用 KV、多 tile 调度),与静态分析、程序分析、形式化验证以及浮点/贝叶斯程序分析均无实质交叉。它虽以 agentic 场景(system prompt、工具模板、RAG 的共享前缀)为动机,但并未对 LLM/Agent 程序做任何分析或验证,只是把这类负载当作性能优化目标,因此不属于用户「PL↔AI」的研究脉络。引用 5 inbox
Jinjun Yi, Zhixin Zhao, Yitao Hu, Ke Yan, Weiwei Sun, Hao Wang, Laiping Zhao, Yuhao Zhang, Wenxin Li, Keqiu Li
crossref: 2026-09-12dblp: 2026-09-12arxiv: 2026-09-12

AI 速览

一句话:提出前缀感知的 decode 注意力内核 PAT,按共享前缀打包查询并复用 KV,平均降低注意力延迟 53.5%。

LLM 服务中解码阶段注意力日益成为主导开销,其受限于从全局内存反复加载海量 KV cache。真实负载普遍存在跨请求的层次化共享前缀(系统提示、工具/模板、RAG),但现有注意力实现未能充分利用:单 query 单 CTA 的执行方式会重复加载共享前缀 KV,而一刀切的 tiling 又使片上资源闲置,并在 KV 长度不均时加剧资源气泡,进而放大访存带宽压力。本文提出 PAT,一种面向 LLM 解码的前缀感知注意力内核,采用 pack-forward-merge 范式组织执行:先按共享前缀打包 query 以减少重复访存,再用定制的多 Tile 内核提升资源效率,并通过多流转发与 KV 切分减少资源气泡,最后以在线 softmax 完成合并,开销可忽略。PAT 作为即插即用插件集成到 vLLM。在真实与合成负载上,相较最先进的注意力内核,注意力延迟平均降低 53.5%,TPOT 降低 17.0–93.1%。

问题
LLM 解码阶段的注意力是访存受限操作,需从全局内存反复加载大规模 KV cache;而真实请求间存在大量层次化共享前缀,现有实现既重复加载共享前缀 KV,又因统一 tiling 造成片上资源闲置与资源气泡。
方法
提出 PAT 前缀感知注意力内核,采用 pack-forward-merge 执行范式:按共享前缀打包 query 以减少重复访存,使用定制多 Tile 内核提高资源效率,并结合多流转发与 KV 切分消除资源气泡,最终以在线 softmax 低开销合并,并作为插件接入 vLLM。
评估
在真实与合成工作负载上,与当前最先进的注意力内核在相同配置下对比;结果显示注意力延迟平均降低 53.5%,TPOT 降低 17.0–93.1%。
原文摘要(Abstract)

LLM serving is increasingly dominated by decode attention, which is a memory-bound operation due to massive KV cache loading from global memory. Meanwhile, real-world workloads exhibit substantial, hierarchical shared prefixes across requests (e.g., system prompts, tools/templates, RAG). Existing attention implementations fail to fully exploit prefix sharing: one-query-per-CTA execution repeatedly loads shared prefix KV cache, while one-size-fits-all tiling leaves on-chip resources idle and exacerbates bubbles for uneven KV lengths. These choices amplify memory bandwidth pressure and stall memory-bound decode attention. This paper introduces PAT, a prefix-aware attention kernel implementation for LLM decoding that organizes execution with a pack-forward-merge paradigm. PAT packs queries by shared prefix to reduce repeated memory accesses, runs a customized multi-tile kernel to achieve high resource efficiency. It further applies practical multi-stream forwarding and KV splitting to reduce resource bubbles. The final merge performs online softmax with negligible overhead. We implement PAT as an off-the-shelf plugin for vLLM. Evaluation on both real-world and synthetic workloads shows that PAT reduces attention latency by 53.5% on average and TPOT by 17.0-93.1% under the same configurations against state-of-the-art attention kernels. PAT's source code is publicly available at https://github.com/flashserve/PAT.

链接与引用

DOI 原文 · arXiv · PDF(开放获取) · DBLP

BibTeX
@inproceedings{YiZHYSWZZLL26,
  title = {PAT: Accelerating LLM Decoding via
                    <u>P</u>
                    refix-
                    <u>A</u>
                    ware A
                    <u>t</u>
                    tention with Resource Efficient Multi-Tile Kernel},
  author = {Jinjun Yi and Zhixin Zhao and Yitao Hu and Ke Yan and Weiwei Sun and Hao Wang and Laiping Zhao and Yuhao Zhang and Wenxin Li and Keqiu Li},
  booktitle = {ASPLOS},
  year = {2026},
  doi = {10.1145/3779212.3790200}
}