TileMix: LLM 추론 가속화를 위한 타일 중심 혼합 정밀도 어텐션

초록

대규모 언어 모델(LLM)의 장문 컨텍스트 프리필은 밀집 자기 주의(dense self-attention)가 이차적인 쿼리-키 점수를 계산하므로 상당한 연산 및 메모리 트래픽을 유발한다. 기존 방법은 균일한 저정밀도 경로를 사용하거나 토큰 상호작용을 선택하는 방식으로, 융합된 밀집 주의 내부에서 하드웨어 정렬 점수 타일에 대한 공간적 정밀도 라우팅을 수행하지 못한다. 우리는 TileMix를 제안한다. TileMix는 타일 중심의 정밀도 라우팅 커널로, 융합된 밀집 주의 내에서 점수 타일 그룹에 대해 수치 정밀도를 실행 가능한 공간적 결정으로 만든다. TileMix는 주의 행렬을 하드웨어 정렬 점수 타일로 분할하고, 라우팅 결정을 컴팩트한 비트마스크로 압축하며, 각 타일 그룹을 FP16 또는 INT8 점수 계산으로 분배한다. 두 경로 모두 공유 온라인 소프트맥스 상태를 업데이트한다. 확장 가능한 정밀도 그룹화를 통해 각 라우팅 비트가 여러 인접 키 타일을 제어할 수 있으므로, 긴 컨텍스트에서도 하드웨어 정렬 연산 타일과 컴팩트 메타데이터를 유지한다. TileMix는 모든 합법적 타일 그룹을 라우팅하므로 밀집 토큰 연결성을 보존하고, 학습이 필요 없으며, 그룹화된 쿼리 주의, 가변 길이 배치, INT8 키/값 캐시를 지원한다. LLaMA, Qwen, Vicuna에 대한 LongEval, LV-Eval, A100 프리필 벤치마크에서 TileMix는 균일 INT8 하에서 손실되는 장문 컨텍스트 품질을 회복하고 FP16 대비 프리필 처리량을 향상시켜, 모델 계열 전반에 걸쳐 제어 가능한 정확도-효율성 프론티어를 제공한다. 구현은 https://github.com/HanzhiZhang-Ulrica/TileMix에서 확인할 수 있다.

English

Long-context prefill in large language models (LLMs) incurs substantial computation and memory traffic because dense self-attention computes quadratic query-key scores. Existing methods either use a uniform low-precision path or select token interactions, leaving spatial precision routing over hardware-aligned score tiles outside fused dense attention. We introduce TileMix, a tile-centric precision-routing kernel that makes numerical precision an executable spatial decision over score-tile groups within fused dense attention. TileMix partitions the attention matrix into hardware-aligned score tiles, packs routing decisions into compact bitmasks, and dispatches each tile group through FP16 or INT8 score computation while both paths update a shared online-softmax state. Scalable precision grouping lets each routing bit govern multiple adjacent key tiles, preserving hardware-aligned compute tiles and compact metadata at long contexts. By routing all legal tile groups, TileMix preserves dense token connectivity, requires no training, and supports grouped-query attention, variable-length batches, and INT8 key/value caches. Across LongEval, LV-Eval, and A100 prefill benchmarks on LLaMA, Qwen, and Vicuna, TileMix recovers long-context quality lost under uniform INT8 and improves prefill throughput over FP16, yielding a controllable accuracy-efficiency frontier across model families. The implementation is available at https://github.com/HanzhiZhang-Ulrica/TileMix.