FlashPrefill: 초고속 롱 컨텍스트 프리필링을 위한 즉각적 패턴 발견 및 임계값 설정

LLM이 긴 문서를 읽을 때 첫 응답이 나오기까지 걸리는 시간은 사용자 경험의 가장 큰 병목이다. FlashPrefill은 복잡한 정렬 과정 없이 중요한 데이터 패턴을 즉시 식별하여 연산량을 획기적으로 줄임으로써, 초거대 모델의 프리필링 속도를 비약적으로 높이고 운영 비용을 절감한다.

Transformer의 Self-Attention은 모든 토큰 간의 관계를 계산하므로 시퀀스 길이가 길어질수록 연산량이 제곱으로 증가한다. 특히 긴 문서를 처음 읽어들이는 프리필링 단계에서 이 연산량은 심각한 지연을 초래한다. 기존 방식은 중요한 토큰을 찾기 위해 전체 점수를 계산한 뒤 정렬하는 과정을 거치는데, 이 정렬 자체가 GPU 하드웨어에서 매우 무거운 작업이라는 한계가 있다.

FlashPrefill은 모든 토큰을 일일이 계산하는 대신, 블록 단위로 데이터를 묶어 평균값을 활용함으로써 중요한 정보의 위치를 대략적으로 먼저 파악한다. 이는 마치 두꺼운 책의 모든 글자를 읽는 대신 목차와 주요 키워드만 훑어보고 읽어야 할 페이지를 정하는 것과 같다. 블록 내 토큰들은 서로 유사한 특징을 가진다는 점을 이용해 계산량을 대폭 줄이면서도 정보의 손실을 최소화한다.

결과적으로 복잡한 정렬 단계 없이도 가장 중요한 정보(최댓값)를 기준으로 문턱값을 설정해 연산 대상을 압축한다. 이를 통해 256K와 같은 초장문 문맥에서도 모델의 정확도는 유지하면서 연산 속도만 수십 배 빠르게 끌어올리는 결과를 얻었다.

전체 접근 방식은 패턴 발견, 임계값 설정, 희소 어텐션 실행의 3단계로 구성된다. 핵심은 블록 근사(Block Approximation)를 통해 메모리 대역폭 낭비를 막고 연산을 융합하는 것이다.

패턴 발견 단계에서는 쿼리 타일과 풀링된 키 블록 간의 상호작용을 계산한다. [쿼리 타일 $Q_I$와 평균 풀링된 키 $k_J$를 입력으로] → [내적 연산 후 Max-Reduction 및 지수화를 수행해] → [블록별 에너지 점수 $S_{I,J}$를 얻고] → [이 값이 해당 블록의 상대적 중요도를 나타내는 지표로 활용된다]. 이 과정은 SRAM 내에서 1D Reduction으로 처리되어 메모리 트래픽을 획기적으로 줄인다.

동적 임계값 설정에서는 $thresh_I = \alpha \cdot \max_{J \le I}(Score_{I,J})$ 공식을 사용한다. [각 행의 최대 점수에 튜닝 가능한 인자 $\alpha$를 곱하는 연산을 수행해] → [최종 임계값을 결정하고] → [이 값보다 낮은 점수를 가진 블록은 연산에서 즉시 제외하여] → [연산 대상 블록의 밀도를 낮춘다].

최종 구현에서는 인덱스 기반 물리적 점프(Index-driven Physical Jumping) 메커니즘을 사용한다. [활성 블록의 좌표 인덱스를 입력으로] → [메모리 포인터를 해당 위치로 직접 리다이렉트하는 연산을 수행해] → [불필요한 조건문 분기 처리를 제거하고] → [하드웨어 파이프라인의 활용도를 극대화한다].

Qwen3-30B 모델 기준 256K 시퀀스에서 기존 Flash Attention 대비 27.78배의 연산 속도 향상을 기록했다. 특히 4K의 짧은 문맥에서도 1.71배의 속도 향상을 유지하여 문맥 길이에 상관없는 견고한 효율성을 입증했다.

vLLM 프레임워크에 통합하여 측정한 결과, 첫 토큰 생성 시간(TTFT)이 최대 7.22배 단축되었다. 이는 실제 서비스 환경에서 사용자가 체감하는 대기 시간을 획기적으로 줄일 수 있음을 의미한다.

정확도 측면에서는 "Needle In A Haystack" 테스트와 RULER 벤치마크에서 Full Attention과 거의 동일한 성능을 보였다. 희소 어텐션을 적용했음에도 불구하고 중요한 정보를 놓치지 않고 정확하게 추출해내는 능력을 유지했다.

FlashPrefill은 Softmax 연산의 Rank-order 불변성을 수학적 기반으로 삼는다. 블록 내 토큰들의 분산이 낮을 때 기하 평균(GM)이 산술 평균(AM)의 엄격한 단조 대리자(Strictly Monotonic Proxy) 역할을 수행할 수 있음을 활용하여 블록 단위 근사의 타당성을 확보했다.

아키텍처적으로는 Fused 2D-Reduction 커널을 통해 중간 행렬의 구체화(Materialization)를 방지했다. 기존 방식이 $O(L \cdot L/B)$의 중간 메모리 트래픽을 발생시키는 반면, FlashPrefill은 이를 $O((L/B)^2)$ 수준으로 억제하여 글로벌 메모리 대역폭 병목을 해결했다.

구현 측면에서 기존 Block-Sparse-Attention 코드베이스의 논리적 스킵(Logical Skipping) 방식이 가진 명령어 스트림 오버헤드를 지적했다. 이를 해결하기 위해 실제 메모리 주소를 직접 제어하는 물리적 점프 방식을 도입하여 GPU의 연산 유닛이 유휴 상태 없이 지속적으로 GEMM 연산을 수행하도록 최적화했다.