TileMix, 하드웨어 타일 단위로 LLM 어텐션 정밀도 라우팅

배경

LLM이 긴 입력을 처리하는 프리필 단계에서는 생성 전에 많은 토큰을 한꺼번에 계산해야 합니다. 밀집 자기 어텐션은 질의와 키의 모든 조합에 대해 스코어를 계산하므로, 컨텍스트가 길어질수록 계산량과 메모리 접근이 빠르게 늘어납니다. INT8 같은 저정밀 연산은 효율을 높일 수 있지만, 어텐션 전체를 일괄적으로 INT8로 바꾸면 긴 컨텍스트 작업의 품질이 떨어질 수 있습니다.

TileMix는 상호작용 자체를 삭제하는 대신, 어텐션 스코어 행렬의 어느 영역에 더 높은 수치 정밀도가 필요한지를 타일 단위로 선택합니다.

타일 중심 혼합 정밀도

TileMix는 어텐션 행렬을 가속기 커널의 실행 방식에 맞는 스코어 타일로 분할합니다. 이어 인접한 키 타일을 그룹으로 묶고, 각 그룹을 FP16으로 계산할지 INT8로 계산할지를 compact bitmask로 표현합니다. 선택된 그룹은 FlashAttention과 유사한 하나의 융합 커널 안에서 처리됩니다.

이러한 그룹화는 실용적인 절충안입니다. 개별 토큰 상호작용마다 정밀도를 지정하면 메타데이터와 제어 분기의 부담이 커질 수 있습니다. 반대로 너무 큰 단위로 선택하면 중요한 영역에만 정밀도를 남기기 어렵습니다. TileMix는 하나의 라우팅 비트가 여러 인접 키 타일을 제어하도록 해, 하드웨어에 맞는 타일 구조를 유지하면서 긴 컨텍스트에서 라우팅 정보를 줄입니다.

FP16과 INT8 경로는 서로 독립적인 어텐션 결과를 만든 뒤 합치는 방식이 아닙니다. 두 경로 모두 같은 온라인 softmax 상태를 갱신합니다. 따라서 서로 다른 정밀도로 계산된 타일이 하나의 밀집 어텐션 정규화 과정에 함께 참여할 수 있습니다. 정밀도는 단순한 사전 양자화 설정이 아니라, 어텐션 커널 내부에서 실행되는 공간적 결정이 됩니다.

토큰 가지치기와의 차이

TileMix가 선택하는 것은 토큰 상호작용의 존재 여부가 아니라 계산 정밀도입니다. 유효한 타일 그룹은 모두 라우팅되므로 밀집 어텐션의 토큰 연결성을 유지합니다. 또한 정밀도 라우팅을 위해 별도 학습을 수행하거나 모델 구조를 바꿀 필요가 없다고 설명합니다. Grouped-query attention, 가변 길이 배치, INT8 키·값 캐시도 지원 대상에 포함됩니다.

결과가 의미하는 것

논문은 LongEval, LV-Eval, A100 프리필 벤치마크에서 LLaMA, Qwen, Vicuna 계열을 평가했습니다. 제공된 요약에 따르면 TileMix는 균일 INT8에서 손실된 긴 컨텍스트 품질의 일부를 회복하고, FP16보다 프리필 처리량을 높였습니다. 이는 완전한 FP16과 일괄 저정밀 연산 중 하나를 택하는 대신, 설정에 따라 정확도와 효율 사이의 경계를 조절할 수 있음을 뜻합니다.

이 연구의 핵심 의미는 혼합 정밀도를 오프라인 양자화 옵션에서 융합 어텐션 커널의 실행 정책으로 확장했다는 점입니다. 다만 밀집 어텐션의 이차 계산량을 없애거나 모든 상황에서 동일한 이득을 보장하는 방법은 아닙니다. 실제 배포 효과는 라우팅 정책, 컨텍스트 길이, 하드웨어, 모델 계열에 따라 달라질 수 있으므로 공개 구현과 세부 벤치마크 조건을 함께 확인해야 합니다.

출처: Hugging Face Daily Papers