CoWA와 MALA, 장문 컨텍스트 연산 부담 대폭 감소

장문 컨텍스트 모델의 연산 중복을 해결하기 위해 **CoWindow Attention (CoWA)**와 **MassAlloc Attention (MALA)**라는 두 가지 새로운 어텐션 메커니즘이 등장했다. 이들은 어텐션 헤드가 먼 컨텍스트와 기여도가 낮은 타일을 처리하는 방식을 최적화한다.

CoWindow Attention (CoWA)

CoWA는 상보적인 윈도우를 사용해 먼 컨텍스트를 KV 헤드에 분배하되, 로컬 및 프리픽스 싱크 윈도우는 공유한다. 각 헤드는 희소하게 어텐션하지만, 가시 위치의 합집합은 전체 인과적 이력을 커버한다. 이 패턴은 위치 기반으로 정의되며 학습된 라우터나 인덱서가 필요 없다.

MassAlloc Attention (MALA)

MALA는 완전한 인과적 QK 스코어링을 유지하되, 어텐션 자체의 softmax 통계를 사용해 타일의 후속 연산 실행 여부를 결정한다. 학습과 추론 모두에서 공통 허용 오차를 사용하여 기여도가 낮은 스코어링 이후 작업을 줄인다.

성능 벤치마크

두 방법 모두 학습 전후방 및 추론 프리필/디코딩을 지원한다. 8 H100 GPU에서 TP=8로 128K 토큰 처리 시 FullAttn 대비 어텐션 연산자 속도 향상은 다음과 같다:

| 방법 | 전방 | 후방 | 디코딩 | | :--- | :--- | :--- | :--- | | CoWA | 7.4x | 8.6x | 3.0x | | MALA | 2.2x | 3.0x | 1.6x |

14B 파라미터와 32K 컨텍스트 환경에서 CoWA는 총 학습 FLOPs를 28.5%, MALA는 23.1% 줄였으며, 보고된 평가에서 모델 성능은 FullAttn과 유사했다. 0.6B에서 14B까지 확장성을 검증했고, 32B 규모에서 별도 지속 학습 실험을 진행했다.

한계점

두 결과 모두 밀집 어텐션과의 보편적 무손실 동등성을 입증하지는 않았다. CoWA의 집단적 커버리지가 FullAttn과 동일한 헤드 단위 상호작용이나 출력을 보장하지 않으며, MALA는 여전히 완전한 인과적 QK 스코어링 비용을 지불한다.