PIVOT은 희소 어텐션(sparse-attention) 모델을 위한 새로운 추론 시점 기법으로, 모델의 가중치를 전혀 수정하지 않고도 인덱서 비용을 최대 4배 절감하고 전체 지연 시간을 약 1.6배 줄여줍니다. 이 방식은 쿼리를 그룹화하고 단일 "프록시(proxy)" 쿼리가 주요 작업을 수행하도록 하는 원리로 작동합니다.

희소 어텐션이 100K 토큰에서 정체되는 이유

희소 어텐션은 연산 비용을 저렴하게 유지하면서 트랜스포머가 100K 토큰 이상의 매우 긴 시퀀스를 처리할 수 있도록 설계되었습니다. 하지만 실제로는 시퀀스가 수만 토큰을 넘어서면 약속되었던 속도 향상이 사라집니다. 그 주범은 인덱서(indexer)입니다. 인덱서는 어떤 토큰이 희소 패턴에 포함될지 결정하기 위해 모든 토큰을 모든 쿼리에 대해 점수화합니다. 인덱서의 작업량은 $O(L^2)$ ($L$ = 시퀀스 길이)로 증가하므로, 100K 토큰에서는 인덱서 작업이 실행 시간의 대부분을 차지하여 희소 어텐션이 주는 이점을 모두 상쇄해 버립니다.

PIVOT의 근간이 되는 두 가지 관찰 결과

연구진은 인접한 쿼리들이 거의 항상 동일한 top-k 토큰을 선택한다는 사실, 즉 약 90%의 중복이 발생한다는 점을 발견했습니다. 이는 단 하나의 대표 쿼리가 인접한 쿼리 묶음 전체를 대신하더라도 여전히 유용한 후보 세트를 찾아낼 수 있음을 의미합니다. PIVOT은 다음과 같은 방식으로 이를 활용합니다.

  1. 정해진 수의 연속된 쿼리를 그룹화(size $g$).
  2. 그룹의 평균을 내어 프록시 쿼리를 생성.
  3. 원래 쿼리마다 한 번씩 실행하는 대신, 프록시에 대해 인덱서를 단 한 번 실행.
  4. 그룹의 각 구성원에 대해 프록시의 후보 목록을 정교화(Refining).

이를 통해 연산량은 $O(L^2)$에서 $O(L^2/g)$로 줄어듭니다. 그룹 크기가 8인 경우, 인덱서의 전체 스캔 횟수가 8배 감소합니다.

두 가지 동작 모드

  • PIVOT-Refine은 밀집(dense) 인덱서의 정확도를 유지하면서 인덱서 단계에서 약 3배의 속도 향상을 제공합니다.
  • PIVOT-Reuse는 최대 처리량(throughput) 이득을 위해 약간의 정확도를 희생하는 대신 속도를 더욱 극대화합니다.

DeepSeek-V3.2 및 GLM-5.1 모델에 대한 벤치마크 결과, 추론 시 PIVOT을 적용했을 때 인덱서 속도는 일관되게 4배 빨라졌으며, 엔드 투 엔드(end-to-end) 지연 시간은 1.6배 감소했습니다.

플러그 앤 플레이(Plug-and-play) 구현

이 기술은 기존의 어떤 Dynamic Sparse Attention (DSA) 파이프라인에도 바로 적용할 수 있는 레퍼런스 구현체 형태로 제공됩니다. 가중치 변경이 필요 없으므로, 표준 희소 어텐션 방식으로 학습된 모델을 그대로 사용할 수 있습니다. 유일한 주의 사항은 레퍼런스 코드가 범용 GPU 커널에서 실행된다는 점입니다. 실제 프로덕션 환경에서 최대 성능을 끌어내려면 수동으로 최적화된 Triton 또는 CUDA 커널이 필요합니다.

트레이드오프(Trade-offs) 분석

PIVOT-Reuse의 속도 이점은 어텐션 품질의 약간의 저하를 동반하며, 이는 정확한 토큰 선택에 매우 민감한 작업에서는 문제가 될 수 있습니다. 따라서 팀은 이러한 품질 손실과 지연 시간 허용 범위 사이의 균형을 고려해야 합니다. 또한, 커스텀 커널이 필요하다는 점은 GPU 커널 전문 지식이 없는 조직에게는 엔지니어링 오버헤드로 작용할 수 있습니다.

향후 주목할 점

PIVOT은 쿼리를 그룹화하고 프록시 스캔을 공유하는 식의 영리한 작업 재배치를 통해, 재학습 비용 없이도 진정한 의미의 긴 컨텍스트를 위한 희소 어텐션의 가능성을 되살리고 실질적인 속도 향상을 가져올 수 있음을 보여줍니다.