MiniMax Sparse Attention

초장기 컨텍스트 (Ultra-long-context) 능력은 에이전트 워크플로우 (agentic workflows), 저장소 규모의 코드 추론 (repository-scale code reasoning), 그리고 지속적 메모리 (persistent memory)를 위해 프런티어 LLM (frontier LLMs)에 필수적인 요소가 되고 있습니다. 이 모든 기능은 모델이 수십만에서 수백만 개의 토큰을 공동으로 어텐드 (attend)할 것을 요구하지만, 소프트맥스 어텐션 (softmax attention)의 이차 비용 (quadratic cost)은 배포 규모에서 이를 실행 불가능하게 만듭니다. 우리는 Grouped Query Attention (GQA)를 기반으로 구축된 블록 단위 희소 어텐션 (blockwise sparse attention)인 MiniMax Sparse Attention (MSA)을 소개합니다. 경량화된 인덱스 브랜치 (Index Branch)가 키-값 (key-value) 블록의 점수를 매기고 각 GQA 그룹에 대해 독립적으로 Top-k 서브셋을 선택하여, 효율적인 블록 수준 실행을 유지하면서도 그룹별 희소 검색 (sparse retrieval)을 가능하게 합니다. 이후 메인 브랜치 (Main Branch)가 선택된 블록들에 대해서만 정확한 블록 희소 어텐션 (block-sparse attention)을 수행합니다. 단순성과 확장성 원칙을 중심으로 설계된 MSA는 의도적으로 간소화되어, 광범위한 GPU에 걸쳐 효율적으로 배포하기 용이합니다. 희소성 (sparsity)을 실질적인 속도 향상으로 전환하기 위해, 우리는 블록 단위 액세스 하에서 텐서 코어 (tensor-core) 활용도를 높일 수 있도록 exp-free Top-k 선택 및 KV-outer 희소 어텐션을 사용하는 GPU 실행 경로와 MSA를 공동 설계했습니다. 네이티브 멀티모달 (natively multimodal) 학습이 적용된 109B 파라미터 모델에서, MSA는 1M 컨텍스트에서 토큰당 어텐션 연산량을 28.4배 줄이면서도 GQA와 대등한 성능을 보여줍니다. 공동 설계된 커널 (kernel)과 결차된 MSA는 H800에서 프리필 (prefill) 14.2배, 디코딩 (decoding) 7.6배의 실제 시간 (wall-clock) 속도 향상을 달성합니다. 우리의 추론 커널은 다음에서 확인할 수 있습니다: https://github.com/MiniMax-AI/MSA. MSA를 기반으로 하는 프로덕션 등급의 네이티브 멀티모달 모델은 다음에서 공개되었습니다: https://huggingface.co/MiniMaxAI/MiniMax-M3.