CPU Can Efficiently Run Long-Context Attention Too: A Deep Dive into the CACPU Kernel

When multiple AI agents process long-context tasks simultaneously, GPU resources can easily become a bottleneck. A new study proposes CACPU—a CPU attention kernel designed specifically for causal full-sequence prefilling—which, through tiered cache blocking and in-kernel task scheduling, achieves up to 2.01× speedup over PyTorch SDPA, FlexAttention, and the SGLang attention kernel on CPU.

Why We Need a CPU Attention Kernel

Agentic AI is delivering significant productivity gains across many domains, but concurrent long-context prefilling from multiple agents quickly saturates GPU compute resources. The CPU is a viable alternative, yet existing CPU attention kernels have two key shortcomings: they cannot adapt to multi-tier cache hierarchies, and they lack effective thread scheduling mechanisms. CACPU is designed precisely to address these two pain points.

Causal full-sequence prefilling is the first phase of large-model inference: before generating any token, the model must process the entire input context in one pass and compute Key/Value caches for all positions. Unlike token-by-token decoding, computation during prefilling consists of highly parallel matrix operations with high computational density, making it sensitive to both memory bandwidth and cache capacity. "Causal" means each position can only attend to itself and preceding positions; this lower-triangular mask structure causes the amount of computation to vary in a stepwise fashion across positions—exactly the root of the load imbalance described below. In agentic AI scenarios, multiple agents each carry long contexts and issue concurrent requests, so the prefilling phase concentrates a large amount of compute demand. Offloading it to the CPU to pipeline with GPU-based decoding is therefore an attractive heterogeneous scheduling approach.

CACPU's Two Core Designs

Tiered Cache Blocking

Modern CPUs have a complex multi-level cache system (L1/L2/L3), with significant differences in capacity and latency across levels. CACPU adopts a tiered cache blocking strategy that matches the size of attention computation data blocks to cache levels, reducing frequent data movement across the memory hierarchy and thereby improving cache hit rates and computational efficiency.

Multi-tier cache hierarchy refers to the pyramid in modern CPUs formed by registers, L1/L2/L3 caches, and main memory: L1 has the smallest capacity (tens of KB) but latency of only a few cycles; L2 ranges from hundreds of KB to several MB; L3 is tens of MB and shared across cores; and main memory reaches tens of GB but with latency two orders of magnitude higher. Attention computation involves heavy round-trip access to Q, K, V matrices and intermediate results. If data block sizes do not match a given cache level's capacity, cache evictions and refills occur frequently, and actual performance falls far below peak compute. So-called "cache blocking/tiling" means splitting large matrices into small blocks that can reside in a specific cache level, so that as much reuse as possible occurs within a block before its data is evicted—turning the memory access bottleneck into a compute bottleneck.

In-Kernel Task Scheduling

Cache optimization alone is not enough. CACPU implements an in-kernel task dispatching mechanism that dynamically allocates thread tasks based on computational load, avoiding the imbalance caused by the causal attention mask—in causal attention, positions early in the sequence participate in far less computation than later positions, which easily leaves threads idle. This load-balancing scheduler makes fuller use of the CPU's parallel capability.

The causal attention mask is a lower-triangular matrix where position i can only attend to positions 0 through i. If the sequence is split evenly among threads by query position, threads at the beginning of the sequence have almost no work, while threads at the end must process all key-values, producing a classic "long-tail" load. In-kernel task dispatching differs from OS-level thread scheduling in that it happens inside a single compute kernel: task granularity is divided by actual computational workload rather than a fixed number of blocks, so threads that finish early dynamically pull new blocks from a shared task queue, achieving work-stealing-style balance. This fine-grained scheduling significantly reduces idle thread waiting and is one advantage CPUs hold over GPUs in implementing flexible scheduling.

Measured Results

The data from the study is quite convincing: CACPU achieves up to 2.01× speedup over PyTorch SDPA, FlexAttention, and the SGLang attention kernel. More importantly, when CACPU is integrated into the SGLang inference framework, the time-to-first-token (TTFT) for Llama-3.1-8B is reduced by up to 1.74× while model accuracy remains unchanged.

This means that in agentic AI scenarios, the CPU is no longer merely a fallback for the GPU but an effective source of compute that can handle real prefilling workloads. For deployment environments that must serve multiple agents concurrently with dense long-context requests, CACPU offers an optimization path that does not depend on scaling up GPUs.

Time to First Token (TTFT) is the time from when a request is issued to when the model outputs its first token, including request queuing, prefilling, and first sampling. It is the metric that most directly affects experience in interactive and agent scenarios. For multi-agent concurrent deployments, TTFT also determines how long a single request occupies resources: the lower the TTFT, the more requests the same CPU/GPU resources can serve per unit time. Note that CACPU optimizes the prefilling phase, so its benefits show up mainly in TTFT and have limited impact on token-by-token decoding throughput—which also explains why heterogeneous approaches that pipeline with GPU decoding are more practical.

Implications for AI Infrastructure

CACPU's value lies not only in its specific speedup numbers but in the direction it reveals for CPU attention optimization: cache-aware data layout and fine-grained thread scheduling are the keys to unlocking CPU compute. As agentic AI applications proliferate, efficient use of heterogeneous compute resources will become increasingly important, and deeply architecture-optimized kernels like this may become an indispensable part of inference infrastructure.

The research team integrated CACPU into SGLang and validated its end-to-end effectiveness, showing that this work is ready for engineering deployment. For teams focused on optimizing inference costs, this is a direction worth following.