동적 블록 주의 희소화와 소프트맥스 임계값 기반 가속

초록

BLASST는 기존 FlashAttention 흐름에 단일 비교 연산만 추가해, 블록별 최대값이 현재 진행 중인 최대값보다 일정 임계값 이하이면 해당 블록의 소프트맥스와 값 로딩, 행렬 곱을 완전히 건너뛰는 동적 희소화 기법이다. 사전 계산이나 프록시 점수가 필요 없으며, 자동 보정 절차를 통해 컨텍스트 길이에 따라 역비례하는 임계값 λ = a/L을 사용해 목표 희소율을 정확히 맞춘다. 프리필과 디코드 모두에서 70% 이상 희소화 시 1.5배 이상 속도 향상을 달성하면서 정확도 손실을 최소화한다.

상세 분석

BLASST의 핵심 아이디어는 FlashAttention이 블록 단위로 수행하는 온라인 소프트맥스 과정에서 이미 계산된 통계값을 활용한다는 점이다. 구체적으로 각 블록 j에 대해 Q·Kᵀ를 계산해 블록 내 최대 점수 ˜m₍j₎를 얻고, 이를 누적 최대값 m₍j₎와 비교한다. 만약 ˜m₍j₎ − m₍j₎ < ln λ이면 exp (˜m₍j₎ − m₍j₎) < λ가 되므로 소프트맥스 정규화 후 해당 블록의 가중치는 거의 0이 된다. 따라서 BLASST는 다음 세 가지 연산을 전부 생략한다. ① exp 연산 및 행합 계산, ② 값(Value) 블록을 HBM에서 SRAM으로 로드, ③ 로드된 값과 가중치의 행렬 곱(MMA). 이 세 단계는 특히 프리필 단계에서 CUDA 코어와 Tensor Core를, 디코드 단계에서는 메모리 대역폭을 크게 차지하므로, 블록을 건너뛰는 것만으로도 전체 연산량을 70% 이상 감소시킨다.

알고리즘 1은 기존 FlashAttention 루프에 ‘if condition skip’ 구문을 삽입한 형태이며, 구현상의 오버헤드는 블록 최대값을 구하는 과정에서 이미 수행되는 rowmax 연산을 재활용함으로써 거의 무시할 수준이다. 또한 BLASST는 블록 단위이므로 토큰 수준의 미세 조정보다 하드웨어 친화적이며, 다양한 어텐션 변형(MHA, GQA, MQA, MLA)과 슬라이딩 윈도우 어텐션에도 그대로 적용 가능하다.

희소율과 정확도 사이의 관계를 정량화하기 위해 저자들은 Llama‑3.1‑8B 모델을 여러 컨텍스트 길이(8K64K)와 RULER 벤치마크 데이터셋에 적용해 실험하였다. 결과는 6070% 희소화까지는 정확도 손실이 거의 없으며, 75% 이상에서는 급격히 떨어진다는 일관된 패턴을 보였다. 중요한 점은 동일 희소율을 유지하려면 임계값 λ이 컨텍스트 길이에 역비례해야 한다는 것이다. 이를 기반으로 자동 보정 절차(알고리즘 2)를 제시했으며, 여러 길이에 대해 λ = a/L 형태의 선형 회귀를 통해 a 값을 추정한다. 이렇게 하면 목표 희소율을 정확히 달성하면서도 예측 가능한 속도 향상을 얻을 수 있다.

추가 실험으로 sparsity‑aware training을 수행했는데, 학습 단계에서 BLASST와 동일한 임계값 기반 마스크를 적용해 모델이 희소 패턴에 내성을 갖도록 했다. 결과적으로 사후 적용 시보다 더 높은 희소율(80% 이상)에서도 정확도 저하가 최소화되었으며, 이는 BLASST가 단순한 추론 최적화 도구를 넘어 훈련 단계에서도 활용 가능함을 시사한다.

전체적으로 BLASST는 (1) 사전 계산이 필요 없는 완전 온라인 방식, (2) 기존 FlashAttention 커널에 최소한의 수정만으로 적용 가능, (3) 프리필·디코드 모두에서 효과적인 두 단계(계산·메모리) 최적화, (4) 자동 보정과 sparsity‑aware training을 통한 실용성 강화라는 네 가지 강점을 가진다. 이러한 설계는 현재 GPU 아키텍처에서 장점이 극대화되며, 장기 컨텍스트를 요구하는 LLM 서비스에 바로 적용할 수 있는 현실적인 솔루션이다.