CPU attention kernel boosts AI response speed with smarter scheduling
Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill
Machine Learning
Summary
When AI systems run tasks that need to remember long histories, they often rely on GPUs, but these can get overwhelmed. The authors found a way to speed up part of this work using CPUs by making better use of the computer’s memory and balancing the workload across processor cores. Their method, called CACPU, runs tasks faster than popular existing programs and helps large AI models deliver their first answer more quickly without losing accuracy.
What this means in practice
- •For data center operators: Run large AI model prefill computations efficiently on CPUs to reduce GPU workload and speed up response times in multi-agent scenarios.
- •For enterprise ai developers: Improve latency of AI systems processing long context sequences by integrating CACPU attention kernel in model-serving pipelines.
Authors
Kwun Fung Lau, Frost Mitchell, Kanya Mo, Emani Ashok
Abstract
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.