CPU上因果注意力预填充调度内核
原文：Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill
智能体AI在各领域推动了显著的生产力提升，但来自多个智能体的并发长上下文预填充可能使GPU计算资源饱和。尽管CPU提供了可行的替代方案，但现有的CPU注意力内核在适应多级缓存层次结构和缺乏有效线程调度方面存在不足。为解决这一问题，我们提出CACPU，一种高性能注意力内核，具有分层缓存分块和内核内任务调度机制，专为因果全序列预填充设计。CACPU相比PyTorch SDPA、FlexAttention和SGLang注意力内核实现了最高2.01倍的加速。集成到SGLang中后，它将Llama-3.1-8B的首Token时间（TTFT）降低了最高1.74倍，同时保持模型精度。作者：Kwun Fung Lau、Frost Mitchell、Kanya Mo、Emani Ashok。类别：cs.LG