UniPrefill: 블록 단위 동적 희소화를 통한 범용적 롱 컨텍스트 프리필 가속화

최신 LLM들이 긴 문맥을 처리하기 위해 다양한 하이브리드 구조를 채택하고 있지만, 기존 가속화 기법들은 특정 구조에만 국한되거나 실제 서비스 환경인 vLLM과의 통합이 어려웠습니다. UniPrefill은 모델 구조에 상관없이 적용 가능한 토큰 삭제 전략을 통해 연산량을 획기적으로 줄이면서도 정확도를 유지하며, 실제 추론 엔진에 즉시 배포 가능한 수준의 통합을 구현했습니다.

기존의 가속 기법들은 주로 Attention 연산 자체를 줄이는 데 집중했습니다. 하지만 최신 모델들은 Attention 외에도 FFN(Feed-Forward Network) 등 다른 연산 비중이 큽니다. UniPrefill은 '중요하지 않은 토큰은 이후 연산에서도 계속 중요하지 않다'는 직관에서 출발합니다.

먼저 Full Attention 레이어에서 쿼리와 키의 유사도를 계산하여 각 토큰의 기여도를 측정합니다. 이때 Softmax 함수를 통해 출력된 확률값이 0에 가까운 토큰들은 다음 토큰 예측에 미치는 영향이 미미하다고 판단합니다. UniPrefill은 이러한 토큰들을 아예 계산에서 제외(Drop)해 버립니다.

중요한 점은 이 삭제 결정이 해당 레이어에만 머물지 않고 이후의 모든 하위 레이어로 전파된다는 것입니다. 마치 필터에서 걸러진 찌꺼기가 다음 공정으로 넘어가지 않는 것과 같습니다. 이를 통해 Attention 연산뿐만 아니라 모델 연산의 큰 비중을 차지하는 GEMM(행렬 곱셈) 연산량까지 동시에 줄여 전체적인 속도를 높입니다.

UniPrefill은 블록 단위의 토큰 중요도 추정 및 선택 프로세스를 따릅니다. 입력 시퀀스를 고정된 크기 G의 블록으로 나누고, 마지막 n개의 쿼리 토큰을 사용하여 각 블록의 평균 Attention 점수를 계산합니다. [입력: 쿼리(n)와 키(N)의 내적] → [연산: 블록별 평균 점수 산출 및 Softmax 정규화] → [출력: 블록별 중요도 점수] → [의미: 각 블록이 다음 토큰 예측에 기여하는 정도를 수치화].

선택 단계에서는 Top-p 전략을 사용합니다. 중요도 점수가 높은 순서대로 블록을 정렬한 뒤, 누적 합계가 임계값 p(예: 0.99)에 도달할 때까지의 블록들만 유지합니다. [입력: 정렬된 블록 점수] → [연산: 누적 합 계산 및 p 기준 필터링] → [출력: 유지할 토큰 인덱스 집합] → [의미: 정보 손실을 최소화하면서 연산 대상을 압축].

시스템 통합 측면에서는 vLLM의 스케줄러를 확장하여 프리필과 디코딩이 공존하는 환경에서 토큰 삭제 이력을 관리합니다. 삭제된 토큰의 위치를 추적하는 메타데이터를 전파하여, 이후 디코딩 단계에서도 모델이 올바른 KV 캐시 위치를 참조할 수 있도록 설계했습니다.

Llama-3.1-8B(Full Attention), Qwen3-Next-80B(Linear 하이브리드), Gemma-3-12B(Sliding Window 하이브리드) 모델을 대상으로 실험을 진행했습니다. RULER 벤치마크에서 UniPrefill은 정확도 손실을 최소화하면서도 Llama-3.1-8B 기준 128K 문맥에서 최대 2.26배의 TTFT 가속을 달성했습니다.

처리량(Throughput) 측면에서도 뛰어난 성능을 보였습니다. vLLM 통합 환경에서 문맥 길이가 길어지고 배치가 커질수록 가속 효과가 뚜렷해졌으며, Llama-3.1-8B의 경우 표준 프리필 대비 최대 109%의 처리량 향상을 기록했습니다. 이는 기존 Sparse Attention 기법들이 하이브리드 모델에서 1.1배 미만의 성능 향상에 그친 것과 대조적입니다.

Ablation Study를 통해 블록 크기 G=64와 관찰 윈도우 n=128이 정확도와 효율성 사이의 최적의 균형점임을 확인했습니다. 또한 Top-p 선택 방식이 고정된 개수를 뽑는 Top-k 방식보다 문맥의 특성에 유연하게 대응하여 더 안정적인 성능을 유지함을 입증했습니다.

UniPrefill 아키텍처의 핵심은 'Sparsity Propagation'입니다. 특정 Full Attention 레이어에서 토큰 삭제가 결정되면, 해당 블록 내의 모든 하위 레이어(FFN, Linear Attention 등)는 오직 남겨진 토큰들만을 입력으로 받습니다. 이는 연산 복잡도를 O(N^2)에서 O(rho * N^2)로 줄일 뿐만 아니라, FFN의 O(N) 연산 또한 O(rho * N)으로 선형적으로 감소시킵니다 (rho는 토큰 유지 비율).

수학적으로는 각 서브레이어가 Lipschitz 연속성을 가진다고 가정할 때, 누적 오차의 상한이 (1-p)에 비례함을 증명하여 이론적 안정성을 확보했습니다. 실제 구현에서는 Layer Normalization과 Residual Connection이 오차 증폭을 억제하는 역할을 합니다.

시스템 구현에서는 Triton을 사용하여 4개의 융합 커널(Fused Kernel)을 개발했습니다. 이 커널들은 GPU 메모리(HBM)에서 데이터를 로드하고, 온라인 Softmax, 블록 리덕션, Top-p 선택을 단일 파이프라인으로 처리하여 CPU-GPU 간 통신 오버헤드를 제거했습니다. 텐서 병렬화 환경에서는 각 랭크가 계산한 부분 점수를 동기화한 후 동일한 삭제 결정을 내리도록 설계되었습니다.

본 연구는 주로 프리필 단계의 가속에 집중하고 있으며, 디코딩 단계의 가속이나 학습 효율성 개선은 직접적으로 다루지 않습니다. 또한 매우 짧은 문맥에서는 토큰 선택 오버헤드가 가속 이득보다 클 수 있습니다.