CACPU:CPU 端因果注意力内核,比 PyTorch SDPA 快 2.01 倍
Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill
多智能体把 GPU 打满时,这篇论文教你怎么用 CPU 跑预填充,比 PyTorch SDPA 快 2 倍,还接进了 SGLang。
CACPU 是一个面向因果全序列预填充的 CPU 注意力内核,采用分层缓存分块和内核内任务分发机制。在基准测试中,其速度最高达到 PyTorch SDPA、FlexAttention 和 SGLang attention 内核的 2.01 倍。集成进 SGLang 后,可将 Llama-3.1-8B 的首 token 生成时间(TTFT)缩短至最高 1.74 倍,且模型精度不受影响。该工作针对多智能体并发长上下文预填充挤占 GPU 算力的问题,把部分负载转到 CPU。论文已发布在 arXiv(编号 2610.38809)。
Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill
Agentic AI drives significant productivity gains across diverse domains, but concurrent long-context prefills from multiple agents can saturate GPU compute resources. Although CPUs offer a viable alternative, existing CPU attention kernels fall short in adapting to multi-tier cache hierarchies and lack effective thread scheduling. To address this, we introduce CACPU, a high-performance attention kernel featuring tiered cache blocking and an in-kernel task dispatching mechanism designed for causal full-sequence prefilling. CACPU achieves up to a 2.01x speedup over PyTorch SDPA, FlexAttention, and SGLang attention kernels. Integrated into SGLang, it reduces Time to First Token (TTFT) of Llama-3.1-8B by up to 1.74x while preserving model accuracy.