Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill

Kwun Fung Lau; Frost Mitchell; Kanya Mo; Emani Ashok

cs.LGOct 28, 2026

Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill

Authors: Kwun Fung Lau, Frost Mitchell, Kanya Mo, Emani Ashok

Organizations: Intel Corp., USA.

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.

Figures & tables

Figure 1: Illustration of K and V tensor packing from DRAM into the L3 cache

Figure 3: Workload imbalance across threads

Input ( $S$ )	CPU	GQA	SDPA (ms)	SGLANG (ms)	FLEX (ms)	Ours (ms)	Ratio
4K	6767P	32:8	8.43	8.26	11.85	5.00	1.65 x
8K	6767P	32:8	27.46	63.22	39.74	15.74	1.74 x
16K	6767P	32:8	102.21	204.67	151.52	54.72	1.87 x
32K	6767P	32:8	391.52	492.22	595.84	207.79	1.88 x
64K	6767P	32:8	1562.45	1966.87	2339.86	812.10	1.92 x
128K	6767P	32:8	6376.64	8147.56	9495.27	3170.23	2.01 x

Table 1: Performance of different kernels on different sequence lengths using 1 batch size, 128 head dimensions

Panel of Figure 5 — Figure 5: Prefill (TTFT) results on Xeon 6767P and 265K. Left: Mistral-7B-v0.3 Right: Qwen3-8B

Panel of Figure 6 — Figure 6: Comparative analysis of the L3-Chunk-Scoped Cost-Prioritized on Xeon 6767P (Left) and 265K (Right)

CPU	$S$	FIFO (ms)	Ours (ms)	Speedup
6767P	4K	6.60	5.12	1.29 $\times$
6767P	8K	19.21	15.53	1.24 $\times$
6767P	16K	59.18	53.02	1.12 $\times$
6767P	32K	212.98	206.06	1.03 $\times$
6767P	64K	822.41	783.67	1.05 $\times$
6767P	128K	3520.41	3119.72	1.13 $\times$

Table 2: Causal attention latency of the Cost-Prioritized and FIFO work-queue schedules across different sequence lengths

CPU	$S$	Index(ms)	Ours (ms)	Speedup
6767P	1024	0.87	0.80	1.09 $\times$
6767P	4928	6.74	6.57	1.03 $\times$
265K	1056	9.36	8.91	1.05 $\times$
265K	1824	21.19	20.86	1.02 $\times$

Table 3: Representative sequences demonstrating speedup under Cost-Prioritized Scheduling

CPU	$S$	SDPA	SGLANG	FLEX	Ours
6767P	4K	0.05	0.04	0.05	0.03
6767P	8K	0.15	0.27	0.12	0.09
6767P	16K	0.62	1.66	0.59	0.42
6767P	32K	2.58	9.79	2.34	1.34
6767P	64K	23.29	53.40	21.53	5.42
6767P	128K	150.50	397.77	210.53	21.13

Table 4: DRAM read (GB) of different kernels across various sequence lengths

$S$	SDPA(%)	FLEX(%)	Ours (%)
4K	84.42	94.30	98.68
8K	79.78	90.28	97.48
16K	78.18	69.37	96.50
32K	74.92	52.62	94.66
64K	75.02	39.75	90.50
128K	75.67	33.78	88.29

Table 5: L3 Cache Hit Rate under various sequence lengths on 265K

Figure 7: Comparative analysis of scheduling strategies in CACPU and vLLM. Top left: Results of Xeon 6767P and 265K with MHA and Head Dim. 128. Top right: Results of Xeon 6767P and 265K with GQA(32:8) and Head Dim. 128. Bottom left: Results of Xeon 6767P and 265K with MHA and Head Dim. 192. Bottom right: Results of Xeon 6767P and 265K with GQA(32:8) and Head Dim. 192

CPU	$S$	Technique	Latency(ms)	Speedup versus last step
6767P	128K	All Off	8512.65	-
6767P	128K	+ L3-Chunk-Scoped Cost-prioritized Scheduling	7883.39	1.08 $\times$
6767P	128K	+ L3 Blocking (Two-Level Cache-Aware Scheduling)	5884.57	1.34 $\times$
6767P	128K	+ GEMM and Softmax Overlap	5239.27	1.12 $\times$
6767P	128K	+ Q-Grouping (Two-Level Cache-Aware Scheduling)	3147.51	1.66 $\times$
6767P	256K	All Off	30372.22	-

Table 6: Forward ablation study of the proposed algorithm across long sequences

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Input ( $S$ , $D$ )	CPU	GQA	SDPA (ms)	SGLANG (ms)	FLEX (ms)	CaCPU (ms)	Ratio
4096,128	6767P	-	8.56	8.23	12.24	4.70	1.75 x
8192,128	6767P	-	28.13	56.41	39.91	17.12	1.64 x
16384,128	6767P	-	107.02	150.20	152.62	56.41	1.90 x
32768,128	6767P	-	407.19	499.66	601.17	219.62	1.85 x
65536,128	6767P	-	1630.19	1942.53	2381.47	838.37	1.94 x
131072,128	6767P	-	6596.03	7875.30	9711.11	3279.43	2.01 x

Appendix

Table 7: Performance of different kernels on different sequence lengths, head dimensions, and GQA settings

$G$	4K (ms)	8K (ms)	16K (ms)	32K (ms)	64K (ms)	128K (ms)
1	89.5	336.9	1322.1	5312.6	21726.2	87370.0
2	88.9	333.8	1306.3	5132.9	20446.6	81665.9
4 (default)	89.7	333.6	1289.4	5093.8	20305.1	80543.8
6	89.0	339.9	1284.4	5067.1	20214.9	80501.8
8	89.6	334.3	1277.8	5061.8	20175.5	80563.4
10	89.8	336.7	1288.4	5045.9	20393.8	80641.7

Appendix

Table 8: Mean latency (ms) of CACPU on the Intel Core Ultra 7 265K across different grouping values

$G$	4K (ms)	8K (ms)	16K (ms)	32K (ms)	64K (ms)	128K (ms)
1	6.44	19.85	74.40	271.60	1042.94	4947.31
2	5.45	17.57	59.38	219.10	846.81	3361.34
4	5.47	16.13	54.87	202.14	781.92	3105.14
6	5.06	16.03	53.84	200.77	771.36	3090.10
8	5.17	15.58	53.60	200.13	779.52	3135.92
10	5.04	15.47	53.93	200.70	781.33	3153.19

Appendix

Table 9: Mean latency (ms) of CACPU on the Xeon 6767P across different grouping values

$S$	Algorithm	TFLOP/s	TFLOP Util. (%)	DRAM read (GB)	FLOP/B
4K	SDPA	14.01	8.82	0.05	2,923.40
4K	SGLang	16.99	10.69	0.04	3,523.08
4K	Flex	11.98	7.54	0.04	3,053.33
4K	Ours	29.33	18.46	0.03	4,432.26
8K	SDPA	20.85	13.12	0.15	3,740.14
8K	SGLang	15.19	9.56	0.27	2,051.49

Appendix

Table 10: TFLOPS utilization of different kernels across varying sequence lengths on the Xeon 6767P

$S$	Algorithm	TFLOP/s	TFLOP Util. (%)	DRAM read (GB)	FLOP/B
4K	SDPA	1.03	34.49	1.21	113.28
4K	SGLang	N/A	N/A	N/A	N/A
4K	Flex	0.49	16.25	0.40	343.79
4K	Ours	1.57	52.17	0.12	1,161.15
8K	SDPA	1.09	36.40	5.73	95.91
8K	SGLang	N/A	N/A	N/A	N/A

Appendix

Table 11: TFLOPS utilization of different kernels across varying sequence lengths on the 265K CPU

	Intel Xeon 6767P	Intel Core Ultra 7 265K
Microarchitecture	Granite Rapids (server)	Arrow Lake-S (hybrid)
Sockets $\times$ cores	$2\times 64$ ; one socket used	$1\times 20$ (8 P-cores + 12 E-cores)
SMT	Enabled	Not supported
ISA	AVX-512, AMX-BF16	AVX2, AVX-VNNI
Max. core frequency	3.9 GHz	5.4 GHz (P), 4.6 GHz (E)
L1D per core	48 KiB	48 KiB (P), 32 KiB (E)

Appendix

Table 12: Hardware and runtime configurations throughout the evaluation

	Xeon 6767P	Core Ultra 7 265K
Q-block rows $q_{r}$	64	64
KV-block columns $k_{c}$ (tile $T_{kv}$ )	512	512
L3 capacity used by the model	336 MiB	30 MiB
L3 budget fraction $f_{\mathrm{L3}}$	0.8	0.8
L2 capacity used by the model	2 MiB (per-core L2)	1 MiB (E-core share of the 4 MiB cluster L2)
L2 budget fraction $f_{\mathrm{L2}}$	0.65	0.65

Appendix

Table 13: Kernel parameters adopted throughout the evaluation

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,101 papers

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
Sep 11, 2026Varun Kumar Dasoju, Tian ZhaoEfficient Attention GPU Kernel Optimization
Department of Computer Science University of Wisconsin–Milwaukee
CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection
May 16, 2026Jiwon Song, Dongwon Jo, Beomseok Kang +1Self-Attention KV-Cache Management
Seoul National University
Multi-Segment Attention: Enabling Efficient KV-Cache Management for Faster Large Language Model Serving
Jun 1, 2026Chunan Shi, Yilei Chen, Yilin Chen +2KV-Cache Offloading LLM Inference
Peking University Beijing, China

Sep 11, 2026J/K move · Enter open · S save

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels

Varun Kumar Dasoju, Tian Zhao

Department of Computer Science University of Wisconsin–Milwaukee

Efficient Attention GPU Kernel Optimization

Open paper AskPDF arXiv

Modern AI systems are built on the Transformer architecture, whose core operation, attention, accounts for the majority of computation and memory cost. Researchers continually propose new attention variants to improve quality, efficiency, or context length, but each variant currently requires expert-written GPU code to run at usable speeds. PyTorch's recent flex_attention lets researchers describe custom attention patterns in Python and compile them to fused kernels, but its design is limited to modifications applied after the central matrix multiplication, excluding Rotary Position Embedding (RoPE), the positional encoding used by every major LLM. We introduce AttnFuse, a small DSL for attention that makes pre-multiplication transformations like RoPE first-class operations. Researchers compose ten high-level building blocks to describe a variant, and AttnFuse's compiler emits a single fused GPU kernel for the entire computation. On an RTX 3090, AttnFuse achieves a 2.10

\times

speedup over flex_attention on the RoPE+causal pattern. On an H100, it runs a full Llama-3-8B training step within 5% of PyTorch's hand-tuned backend. Our investigation reveals the Rotation Calculus: whether to fuse RoPE or apply it separately depends on the GPU's compute-to-bandwidth ratio, with a derived crossover that matches measurement. AttnFuse demonstrates that a small, attention-specific compiler can close the gap between flexible research code and production kernels.

Abstract

LLM Inference Acceleration

Explore similar work

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

Multi-Segment Attention: Enabling Efficient KV-Cache Management for Faster Large Language Model Serving