arXiv:2511.22333cs.DCcs.CL2025-11中稿 · ASPLOS'26, code av…被引 5

通过感知前缀共享,显著降低大模型解码时的注意力延迟。

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

  • 按共享前缀打包查询,减少重复加载键值缓存
  • 多瓦片自适应内核使计算资源利用率提升53.5%
  • 适合高并发、长文本生成场景的推理部署

大模型服务中解码注意力操作因大量从全局内存加载键值缓存而成为内存瓶颈。实际负载中请求间存在显著的层级化前缀共享(如系统提示、工具模板、RAG)。现有注意力实现未能充分挖掘前缀共享:单查询每线程块执行导致共享前缀键值缓存反复加载,统一分块策略使片上资源闲置并加剧不均长度键值缓存的资源空泡。本文提出PAT,一种面向大模型解码的前缀感知注意力内核,采用打包-前向-合并范式。PAT按共享前缀打包查询以减少重复内存访问,使用定制化多瓦片内核实现高资源效率,进一步引入多流前向和键值拆分以减少资源空泡。最终合并阶段在线执行softmax,开销极低。我们在vLLM中实现PAT作为即插即用插件。在真实与合成负载下的评估显示,相比最先进注意力内核,PAT平均将注意力延迟降低53.5%,在相同配置下将端到端延迟(TPOT)降低17.0%-93.1%。PAT源码已公开于https://github.com/flashserve/PAT。

原文摘要 · Abstract (English)

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.

大模型推理注意力优化KV缓存vLLM

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。