CPU 캐시 계층과 작업 대기열로 롱컨텍스트 프리필을 가속한다
Causal Attention on CPU: A Cache-Aware and Load-Balanced Scheduling Kernel for Prefill
무엇인가
에이전트 AI가 여러 도구 호출과 추론 단계를 길게 이어가면서 동시에 들어오는 롱컨텍스트 프리필이 GPU 연산을 포화시킨다. CPU를 대안으로 쓰려 해도 기존 CPU 어텐션 커널은 L2/L3 캐시 계층을 충분히 활용하지 못하고, 코어별 작업량을 고르게 나누지 못한다. 이 논문은 표준 causal attention의 프리필을 CPU에서 빠르게 끝내는 CACPU 커널을 제안한다. 모델 재학습이나 구조 변경 없이 PyTorch와 SGLang 같은 서빙 프레임워크에 바로 넣는 것을 목표로 한다.
어떻게 동작하나
첫 번째 축은 2단계 캐시 인지 스케줄링이다. L3에서는 K/V 텐서가 L3 예산 안에 들어오도록 Q 헤드를 그룹으로 묶고, 그룹을 바깥 루프로 돌려 K/V 블록이 축출되기 전에 재사용한다. 한 헤드의 K/V가 예산을 넘으면 시퀀스도 추가로 쪼갠다. 같은 소켓의 코어가 L3를 공유하므로 한 스레드가 가져온 K/V 블록을 다른 코어도 즉시 쓸 수 있다. L2에서는 KV 스텝 루프를 바깥으로, 연속한 G개 q-블록 루프를 안쪽으로 뒤집어 KV 스텝을 L2에 한 번 올리고 여러 q-블록이 재사용하게 한다. G는 L2 용량에서 공유 항 U와 q-블록별 항 Qb를 빼서 계산하므로 CPU마다 수동 튜닝하지 않는다.
무엇과 다른가
두 번째 축은 L3 청크 단위의 비용 우선 스케줄링이다. causal attention은 아래삼각 마스크 때문에 앞쪽 q-그룹은 일이 적고 뒤쪽은 많아 스레드 간 불균형이 생긴다. CACPU는 각 q-그룹의 비용을 마스크되지 않은 타일 수로 정확히 계산해 내림차순으로 정렬하고, L3 청크마다 다시 만든 공유 큐에 넣는다. 스레드는 원자적 인덱스로 가장 비용이 큰 남은 작업을 동적으로 가져가므로, 먼저 끝난 스레드가 추가 작업을 맡는다. 이 방식은 P코어와 E코어가 섞인 CPU에서도 빠른 코어가 더 많은 일을 가져가 균형을 맞춘다. 추가로 Xeon에서는 QK^T GEMM을 AMX에 맡기면서 이전 타일의 온라인 소프트맥스를 AVX-512로 처리해 GEMM과 소프트맥스를 겹친다.
어떻게 쓰나
실험은 서버용 Intel Xeon 6767P와 데스크톱용 Core Ultra 7 265K에서 진행됐다. 베이스라인은 PyTorch SDPA, PyTorch FlexAttention, SGLang 어텐션 커널이며 265K는 AMX가 없어 SGLang 비교에서 제외됐다. 커널 지연시간은 Xeon에서 4K 토큰에서 1.65~1.75배, 128K에서 2.01배 빨랐고, 265K에서는 1.45~1.58배가 시퀀스 길이 전반에서 유지됐다. SGLang에 통합한 TTFT 프리필 실험은 Llama-3.1-8B, Qwen3-8B, Mistral-7B-v0.3에서 이뤄졌고 Xeon 128K에서 약 1.74배 개선을 보였다. 265K는 시스템 메모리 제약으로 128K 테스트를 생략했고, AMX가 없어 비어텐션 GEMM 시간이 커 상대 이득이 제한됐다.
전제와 한계
스케줄링만 먼저 켠 경우 Xeon 128K에서 1.08배에 그쳤지만, 완전 최적화에서 이를 마지막에 제거하면 1.45배 느려졌다. 265K에서는 로드밸런싱이 즉시 1.44배를 냈고, 캐시 인지 스케줄링만 단독으로 켜면 오히려 성능이 나빠졌다. 두 기법을 함께 켰을 때 2.45배까지 개선됐다는 것이 저자들의 핵심 상호작용 주장이다. 비용 우선 스케줄링은 static 대비 Xeon에서 최대 1.71배, 265K에서 짧은 시퀀스 1.72배·긴 시퀀스 2.45배를 냈다. FIFO 대비 평균 이득은 Xeon 14.33%, 265K 4.16%였다. 캐시 쪽에서는 모든 시퀀스 길이에서 DRAM 읽기가 줄었고 265K의 평균 L3 적중률은 94.35%였으며, 기본 G 값은 최적값 대비 265K에서 0.58%, Xeon에서 2.34%만 느렸다.
정확도 검증은 FP64 PyTorch eager SDPA와 비교해 최소 코사인 유사도 0.999, 최대 절대차 0.01이었고, SGLang에 넣은 Llama-3.1-8B는 MMLU 66.2%로 공식 베이스라인 66.7%와 사실상 같았으며 RULER도 일관됐다. 저자들이 밝힌 한계는 표준 causal attention을 전제한다는 점이다. attention sink나 사용자 정의 희소 마스크가 들어오면 비용 모델과 메모리 접근 패턴이 달라져 L2/L3 지역성이 나빠질 수 있다. 다만 스케줄링 기준을 확장해 프리필 변형, 사용자 마스크, 디코드까지 적용할 수 있다고 본다. 개발자 입장에서는 GPU가 긴 프리필로 포화될 때 CPU 서빙 경로를 검토할 근거가 되지만, AMX 지원 여부와 이종 코어 구성, 캐시 크기에 따라 이득이 달라진다는 점을 먼저 확인해야 한다.