FlashPrefill V2: 장문맥 LLM 서빙을 위한 블록-희소 프리필 어텐션

초록

장문맥 모델링은 대규모 언어 모델의 핵심 역량이지만, 어텐션의 이차 복잡도는 특히 계산 집약적인 프리필링 단계에서 여전히 중요한 병목 지점으로 남아 있다. 본 연구진의 이전 연구인 FlashPrefill은 즉각적인 패턴 발견과 최대값 기반 동적 임계값 설정을 통해 이러한 비용을 완화했으나, 실제 배포와는 거리가 먼 알고리즘적 프로토타입에 그쳤다. 본 논문에서는 FlashPrefill을 프로토타입 단계에서 실용적인 장문맥 서비스를 위한 도구로 발전시키는 FlashPrefill V2를 세 가지 측면에서 제시한다. 첫째, 평균 보정 항을 도입하여 근사 오차를 효과적으로 억제함으로써 극단적인 희소성 수준에서도 성능 저하를 관리 가능한 범위로 유지한다. 둘째, PackGQA 메모리 접근, 워프 전문화, 핑퐁 파이프라이닝을 적용하여 희소 어텐션 연산자를 재설계함으로써 최신 FlashAttention-3/4 구현과 완전히 정렬시키고, 실제 양자화 요구사항을 충족하기 위해 FP8 추론을 지원한다. 셋째, FlashPrefill V2는 페이지드 KV 캐시와 연속 배칭을 기본적으로 지원하여 SGLang과 같은 최신 추론 프레임워크에서 어텐션 백엔드로 통합될 수 있다. 가장 널리 배포된 추론 가속기 중 하나인 NVIDIA H20 GPU에서의 광범위한 평가를 통해, FlashPrefill V2는 FP8 및 BF16 정밀도에서 각각 128K 컨텍스트 길이 기준 FlashAttention-2 대비 최대 47.26배 및 27.19배의 속도 향상을 달성하며, FP8에서는 FA3/4 정렬 밀집 기준선 대비 여전히 30.49배의 속도 향상을 달성함을 입증한다.

English

Long-context modeling is a pivotal capability for Large Language Models, yet the quadratic complexity of attention remains a critical bottleneck, particularly during the compute-intensive prefilling phase. Our previous work, FlashPrefill, mitigates this cost through instantaneous pattern discovery and max-based dynamic thresholding; however, it remains an algorithmic prototype that is still distant from production deployment. In this paper, we present FlashPrefill V2, which evolves FlashPrefill from a prototype toward practical long-context serving along three dimensions. First, we introduce a mean correction term that effectively suppresses the approximation error, keeping performance degradation manageable even at extreme sparsity levels. Second, we redesign the sparse attention operator with PackGQA memory access, warp specialization, and pingpong pipelining, fully aligning with the latest FlashAttention-3/4 implementations and supporting FP8 inference to meet practical quantization requirements. Third, FlashPrefill V2 natively supports paged KV cache and continuous batching, allowing integration as an attention backend in modern inference frameworks such as SGLang. Extensive evaluations on NVIDIA H20 GPUs---among the most widely deployed inference accelerators---demonstrate that FlashPrefill V2 delivers up to 47.26x and 27.19x speedups over FlashAttention-2 at 128K context length under FP8 and BF16 precision, respectively, and, in FP8, still achieves a 30.49x speedup against an FA3/4-aligned dense baseline.