15 분 소요

0. Introduction

Paper link

Official kernel repository

MiniMax-M3 repository

한 줄 요약: MiniMax Sparse Attention, 이하 MSA는 GQA group마다 필요한 KV block을 고르는 가벼운 Index Branch와, 선택된 block에만 exact softmax attention을 수행하는 Main Branch를 결합한다. 핵심은 sparse pattern만 제안한 것이 아니라, indexer 학습과 GPU kernel까지 함께 설계해 1M context에서 실제 wall-clock speedup으로 연결했다는 점이다.

MSA를 단순히 또 하나의 sparse attention으로 읽으면 핵심을 놓치기 쉽다. 이 논문이 겨냥하는 문제는 attention FLOPs만 줄이는 것이 아니다. 어떤 token을 남길지 학습하고, 그 선택을 안정적으로 pretrain하며, irregular sparsity를 GPU가 효율적으로 실행하게 만드는 전체 경로를 함께 다룬다.

이 논문을 지금 볼 가치가 있는 이유는 다음과 같다.

  • Long-context 경쟁의 병목이 context length 자체보다 serving cost와 hardware utilization으로 이동하고 있다.
  • GQA checkpoint를 sparse model로 바꾸는 continued pretraining route와, 처음부터 sparse하게 학습하는 route를 모두 비교한다.
  • Text뿐 아니라 image와 video를 섞은 native multimodal pretraining에서 sparse attention을 검증한다.
  • 28.4x attention compute reduction, 14.2x prefill speedup, 7.6x decode speedup이라는 숫자를 algorithm과 kernel의 co-design으로 설명한다.

1. Problem Setting

1-1. Full attention의 병목

Causal softmax attention에서 query position $t$와 head $h$의 출력은 다음처럼 쓸 수 있다.

\[O_{t,h} = \operatorname{Softmax}\left(\frac{q_{t,h}K_g^T}{\sqrt{d_h}} + M_t\right)V_g\]

여기서 $g$는 해당 query head가 속한 GQA group이다. GQA는 KV head 수를 줄여 KV cache를 절약하지만, 각 query가 여전히 전체 causal context를 훑는다는 점은 바뀌지 않는다. Sequence length가 커지면 prefill의 score matrix와 value aggregation 비용이 빠르게 증가하고, decode에서도 매 token마다 긴 KV cache를 읽어야 한다.

1M context에서는 이 문제가 단순한 benchmark 숫자가 아니다. Repository-scale coding, persistent agent memory, long video understanding처럼 실제 workload가 긴 문맥을 요구하면 attention이 latency와 serving cost를 직접 결정한다.

1-2. Sparse attention이 어려운 세 가지 이유

Sparse attention의 아이디어 자체는 단순하다. 모든 token을 보지 않고 중요한 subset만 선택하면 된다. 하지만 실제 model을 만들 때는 세 문제가 동시에 생긴다.

1) Selection quality

중요한 context를 놓치면 계산량은 줄어도 retrieval과 reasoning 성능이 무너진다. Fixed sliding window는 가까운 token에는 강하지만, 멀리 있는 evidence를 content에 따라 가져오기 어렵다.

2) Trainability

Top-k routing은 discrete operation이라 language modeling loss가 selector를 직접 학습시키기 어렵다. Selector가 초기부터 잘못된 block을 고르면 Main Branch가 나쁜 context를 보고, 그 결과 다시 selector에 약한 supervision만 돌아가는 악순환이 생긴다.

3) Hardware efficiency

Sparse FLOPs가 적다고 GPU가 자동으로 빨라지는 것은 아니다. Query마다 선택한 block이 다르면 memory access가 불규칙해지고, 작은 matrix multiplication이 많아져 tensor core utilization이 떨어질 수 있다. 이 경우 theoretical FLOP reduction과 wall-clock speedup 사이에 큰 차이가 생긴다.

1-3. 기존 접근과 MSA의 위치

기존 방법은 대략 다음 세 계열로 볼 수 있다.

Approach Selection Strength Main caution
Fixed sparse pattern Position-based window or global token Simple and regular Content-dependent remote retrieval가 약함
Inference-time sparsification Pretrained dense model에서 runtime selection 추가 pretraining 부담이 작음 Training cost는 dense이고 일부 phase가 dense에 가까울 수 있음
Native learned sparsity Training 중 selector 학습 Model representation이 sparsity에 적응 가능 Training stability와 kernel complexity가 큼

MSA는 세 번째 계열에 속하지만, 두 가지 operating point를 모두 제공한다. MSA-PT는 처음부터 sparse하게 pretrain하고, MSA-CPT는 기존 GQA checkpoint를 MSA로 교체한 뒤 continued pretraining한다.

2. Core Idea

2-1. Two-branch sparse attention

MSA의 구조는 Index Branch와 Main Branch로 나뉜다.

  1. Index Branch가 현재 query와 causal context의 key를 가볍게 scoring한다.
  2. Token score를 block 단위로 max pooling한다.
  3. 각 GQA group이 독립적으로 Top-k KV block을 선택한다.
  4. Main Branch가 선택된 block의 token에만 exact scaled dot-product softmax attention을 수행한다.

이 구분이 중요하다. Index Branch는 무엇을 볼지 결정하고, Main Branch는 선택된 support 안에서 기존 softmax attention의 계산을 그대로 유지한다. 즉 linear attention처럼 attention operator 자체를 근사하기보다, exact attention을 수행할 support를 줄이는 방식이다.

2-2. Per-GQA-group routing

Per-head token routing은 표현력은 높지만 GPU에서 다루기 어렵다. 반대로 모든 head가 하나의 selection을 공유하면 regularity는 좋아도 서로 다른 semantic pattern을 분리하기 어렵다.

MSA는 중간 지점을 택한다. 같은 KV head를 공유하는 query head 묶음, 즉 GQA group 단위로 selection을 공유한다. 각 group은 독립적인 Top-k block set을 가지지만, group 내부 query head는 같은 KV block을 읽는다.

이 선택은 다음 trade-off를 만든다.

  • Per-head routing보다 index 수와 memory access pattern이 단순하다.
  • Global shared routing보다 group별 retrieval diversity를 유지한다.
  • Block-level access라 KV read가 연속적이고 kernel tile을 구성하기 쉽다.

2-3. Fixed attention budget

공식 구현의 대표 설정은 block size 128, Top-k 16이다. Local block 한 칸을 포함해 query와 GQA group마다 최대 2048 KV token에 Main Branch attention을 수행한다.

선택 block 집합을 다음처럼 쓸 수 있다.

\[s_{t,g,j} = \frac{q^I_{t,g} \cdot k^I_j}{\sqrt{d_I}}\] \[S_{t,g,b} = \max_{j \in b,\ j \le t} s_{t,g,j}\] \[\mathcal{B}_{t,g} = \operatorname{TopK}_k(S_{t,g,:}) \cup \{\operatorname{local}(t)\}\]

Main Branch는 선택된 support에 대해 다음 계산을 수행한다.

\[O_{t,h} = \operatorname{Attn}\left(q_{t,h}, K_{g,\mathcal{B}_{t,g}}, V_{g,\mathcal{B}_{t,g}}\right)\]

Main Branch의 token budget은 context length와 무관하게 고정된다. 다만 Index Branch는 candidate context를 scoring해야 하므로 전체 시스템이 strict linear attention처럼 full-context scan을 완전히 없애는 것은 아니다. MSA가 줄이는 핵심은 query head 수와 value aggregation이 결합된 비싼 Main Branch path다.

2-4. Simplicity as a design constraint

논문은 여러 auxiliary branch를 쌓기보다 standard GQA에 두 projection matrix만 추가하는 방향을 택한다.

  • GQA group마다 하나의 index query projection을 둔다.
  • Index key projection은 group 사이에 공유한다.
  • 별도 index value output은 최종 recipe에서 제거한다.
  • First block sink나 큰 local window를 hard-code하지 않는다.
  • 현재 query가 속한 incomplete local block만 항상 포함한다.

이 simplicity는 단순한 미학이 아니다. Training signal을 indexer에 한정하고, inference kernel의 data path를 단순하게 유지하기 위한 system-level constraint로 읽는 편이 좋다.

3. Architecture / Method

3-1. Overview

Item Description
Base attention Grouped Query Attention
Selector Lightweight Index Branch
Selection unit Contiguous KV block
Selection scope Independent Top-k for each query and GQA group
Main computation Exact softmax attention over selected blocks
Stability devices KL alignment, gradient detach, indexer warmup, forced local block
Kernel ideas Exp-free Top-k, KV-outer execution, pre-scheduled tile chunking, two-phase combine

3-2. Index Branch

1) Lightweight projection

Index Branch는 standard GQA input에서 index query와 index key를 만든다. 논문 설명상 standard GQA에 추가되는 learnable component는 두 projection matrix다.

각 query와 GQA group은 index query를 만들고, 각 context token은 shared index key를 만든다. 이 둘의 dot product가 token-level index score가 된다.

2) Token score to block score

GPU에서는 token마다 제각각 gather하는 것보다 contiguous block을 읽는 편이 훨씬 효율적이다. MSA는 한 block 안에서 causal하게 visible한 token score의 maximum을 block score로 사용한다.

Max pooling은 block 안에 강한 evidence가 하나만 있어도 해당 block을 후보로 올릴 수 있다는 장점이 있다. 반면 block 평균을 쓰면 작은 but critical evidence가 주변 token에 묻힐 수 있다.

3) Group-specific Top-k

각 GQA group은 서로 다른 block set을 선택한다. 예를 들어 한 group은 최근 local syntax에 집중하고, 다른 group은 sequence 초반의 global instruction이나 멀리 있는 visual evidence를 고를 수 있다.

이때 group 내부 query head는 같은 selection을 공유하므로 Main Branch의 KV access는 group 단위로 regular하게 유지된다.

3-3. Main Branch

Main Branch는 선택된 block에서 standard attention score와 softmax를 계산한다. 중요한 점은 attention weight 계산을 근사하지 않는다는 것이다.

  • Q, K, V projection은 기존 GQA의 것을 사용한다.
  • 선택 support 내부에서는 exact softmax normalization을 수행한다.
  • Query head마다 attention weight는 따로 계산한다.
  • Selection만 GQA group 안에서 공유한다.

따라서 품질 손실의 주된 원인은 softmax approximation이 아니라 selector가 필요한 block을 놓치는 selection error다. 논문의 block recall과 score recall 분석이 중요한 이유도 여기에 있다.

3-4. Forced local block

각 query는 자신이 속한 local block을 항상 선택한다. 이 block은 causal boundary 때문에 완전히 채워지지 않은 self block일 수 있다.

이 장치는 초기 indexer가 불안정할 때 immediate neighborhood를 놓치는 퇴행을 막는다. 반면 첫 block을 attention sink로 강제하거나 큰 sliding window를 별도로 유지하지는 않는다. Appendix ablation에서는 이런 pattern이 필요할 때 model이 스스로 학습할 수 있다고 보고한다.

3-5. GPU kernel co-design

1) Exp-free Top-k

Softmax는 score ordering을 바꾸지 않는다. 따라서 Top-k index만 필요하다면 max, exp, sum normalization을 수행할 이유가 없다. MSA kernel은 raw index score에서 바로 Top-k를 뽑는다.

이 최적화는 selector의 semantic behavior는 유지하면서 index path의 unnecessary math를 제거한다.

2) KV-outer sparse attention

Query를 순회하면서 각 query의 selected KV를 gather하는 Q-outer 방식은 query마다 KV set이 달라 memory reuse가 약하다.

MSA는 반대로 KV block을 중심에 놓는다. Reverse sparse index를 통해 해당 block을 선택한 query를 모으고, 같은 KV operand를 공유하는 query position을 concatenation해 tensor core용 matrix multiplication을 크게 만든다.

Sequence prefix나 특정 global evidence block은 매우 많은 query가 선택할 수 있다. One-CTA-per-tile로 처리하면 hot block이 긴 tail latency를 만든다.

MSA는 scheduler가 인기 KV tile의 query dimension을 미리 chunk로 나누고 여러 CTA에 분배한다. 그 뒤 partial result를 two-phase combine으로 합친다. 이 설계는 sparse routing의 load imbalance를 kernel scheduling 문제로 직접 다룬다.

4) Training kernel fusion

KL loss에 필요한 log-sum-exp 값을 Main Branch와 Index Branch forward에서 함께 저장해 별도 KL forward kernel을 줄인다. Backward에서도 variable sparse workload를 persistent scheduling으로 분산한다.

결국 논문의 kernel section은 sparse attention에서 자주 빠지는 질문에 답한다. FLOPs가 줄었다면 GPU가 왜 실제로 빨라지는가라는 질문이다.

4. Training / Data / Recipe

4-1. Experimental backbone and data

Main experiment는 native multimodal 109B MoE model에서 진행된다.

  • 41 layers
  • 약 109B total parameters
  • Token당 약 6B activated parameters
  • First 3 layers는 dense, remaining 38 layers는 MoE
  • 64 query heads, 4 KV heads
  • Head dimension 128, RoPE dimension 64
  • Text와 image/video가 섞인 native multimodal pretraining
  • Total training budget 3T tokens

이 model은 논문 검증용 backbone이다. 공개된 production model MiniMax-M3는 약 428B total parameters와 약 23B activated parameters를 가진 별도 model이므로, 논문 Table 2의 109B 결과와 M3 repository의 수치를 직접 섞어 비교하면 안 된다.

4-2. Why LM loss alone is insufficient

Top-k block selection은 non-differentiable하다. Main Branch의 language modeling loss만으로는 선택되지 않은 block과 index projection에 직접적인 learning signal을 주기 어렵다.

논문은 Index Branch distribution을 Main Branch attention distribution에 맞추는 auxiliary KL loss를 사용한다.

\[\mathcal{L} = \mathcal{L}_{LM} + \lambda \sum_l \mathcal{L}_{KL}^{(l)}\]

Teacher는 같은 GQA group의 Main Branch query head probability를 평균한 distribution이다. Index Branch는 이 teacher를 따라가며 어떤 block이 실제 attention mass를 많이 받는지 학습한다.

4-3. Gradient detach

KL loss를 그대로 backbone까지 흘리면 문제가 생긴다. Backbone이 Main Branch attention을 indexer가 따라가기 쉬운 방향으로 바꿔 auxiliary objective를 우회할 수 있고, gradient spike나 short-context regression이 생길 수 있다.

MSA는 두 곳을 분리한다.

  • Main Branch teacher distribution은 detach한다.
  • Index Branch input에도 stop-gradient를 적용한다.

이렇게 하면 KL loss는 index query와 index key projection만 업데이트한다. Backbone은 language modeling objective로 학습하고, indexer는 Main Branch가 이미 만든 attention pattern을 따라간다.

이 separation은 sparse routing을 추가하면서 base model optimization을 오염시키지 않기 위한 핵심 recipe다.

4-4. Indexer warmup

Training 초반에는 attention entropy가 빠르게 변한다. 이 시점부터 sparse selection을 켜면 아직 random에 가까운 indexer가 Main Branch에 나쁜 context를 전달하고, teacher target도 계속 움직인다.

그래서 MSA는 먼저 full attention으로 indexer를 warmup한다.

Route Starting point Indexer warmup Sparse training
MSA-PT Random initialization First 40B tokens Remaining budget within 3T tokens
MSA-CPT GQA checkpoint at 2.6T tokens First 40B of continued training Remaining 360B within 400B continued training

Warmup 동안 Main Branch는 full attention을 사용하고, Index Branch는 full-context teacher에 맞춰 KL alignment를 학습한다. Warmup 뒤 Top-k routing을 활성화한다.

여기서 40B tokens는 작지 않다. 논문이 near-lossless conversion을 보여주지만, 이는 몇 step의 cheap patch가 아니라 substantial continued pretraining budget을 전제로 한다.

4-5. Two conversion strategies

1) MSA-PT

처음부터 sparse pattern을 학습한다. Model representation이 sparse support에 맞게 전체 pretraining 동안 적응할 수 있다. 실험에서는 math, image, video, retrieval의 여러 benchmark에서 가장 강한 결과를 보인다.

2) MSA-CPT

이미 학습된 Full-Attention GQA checkpoint의 attention을 MSA로 교체하고 continued pretraining한다. Dense model의 behavior를 최대한 유지하면서 sparse serving으로 이동하는 route다.

실무적으로는 기존 checkpoint 자산을 활용할 수 있다는 장점이 있지만, 400B continued tokens와 40B warmup이 필요하다. Conversion cost와 serving saving을 함께 계산해야 한다.

4-6. Engineering notes

1) Block size and Top-k are kernel hyperparameters

Block size는 retrieval granularity와 arithmetic intensity를 동시에 결정한다. 작은 block은 precise selection에 유리하고, 큰 block은 contiguous memory access와 matrix multiplication 효율에 유리하다.

Appendix에서는 total selected token 수를 고정한 채 block size 32, 64, 128을 비교한다. Agent PPL은 거의 비슷하지만 RULER-32K는 66.1, 65.3, 64.6으로 완만하게 낮아진다. 즉 larger block이 완전히 free인 것은 아니며, workload에 따라 quality-efficiency curve를 다시 측정해야 한다.

2) Fixed sliding window보다 dynamic selection이 중요하다

같은 FLOP budget의 sliding-window baseline보다 MSA가 agent-oriented PPL에서 일관되게 낮다. Sparse budget만 맞춘다고 같은 품질이 나오는 것이 아니라, 먼 context를 content에 따라 가져오는 learned routing이 필요하다는 evidence다.

3) Index value head는 최종 recipe에서 제거된다

초기 실험에서는 Index Branch output을 Main Branch에 더하는 value path가 training을 도왔지만, warmup을 도입하면 이 component 없이도 충분했다. 최종 design은 selector 역할만 남겨 inference overhead와 architecture complexity를 줄인다.

5. Evaluation

5-1. Main pretraining results

논문은 같은 3T-token budget에서 Full-Attention, MSA-PT, MSA-CPT를 비교한다. 아래는 전체 Table 2 중 해석 가치가 큰 일부다.

Category Benchmark Full MSA-PT MSA-CPT
General MMLU 67.0 67.2 66.8
Math GSM8K 76.2 77.7 73.7
Code HumanEval 61.0 64.0 57.9
Retrieval RULER-8K 79.8 84.2 77.2
Retrieval RULER-32K 75.0 77.5 75.7
Image VisualWebBench 55.6 68.4 59.4
Video EgoSchema 29.6 37.6 25.8
Video VideoMME 41.11 45.48 39.65
Agent PPL TAU2 1.155 1.148 1.150
Agent PPL SWE 1.216 1.218 1.216

세 가지를 읽어야 한다.

1) MSA-PT는 단순한 quality preservation 이상을 보인다

MSA-PT는 여러 math, retrieval, image, video benchmark에서 Full보다 높다. Sparse training 자체가 regularization처럼 작용했다고 단정할 수는 없지만, model이 전체 pretraining 동안 learned sparse support에 적응하면 dense baseline과 동등한 수준을 넘어서는 경우가 있다는 뜻이다.

2) MSA-CPT는 conservative conversion route다

MSA-CPT는 text와 PPL 영역에서 Full에 가까운 경우가 많지만, GSM8K, HumanEval, EgoSchema처럼 gap이 큰 항목도 있다. Near-lossless라는 표현은 benchmark 전체 평균의 경향으로 읽어야 하고, 모든 capability가 그대로 유지된다는 뜻으로 해석하면 안 된다.

3) Native sparse training과 conversion은 같은 product choice가 아니다

새 foundation model을 처음부터 학습할 수 있다면 MSA-PT가 더 매력적이다. 이미 dense model이 있고 serving cost를 줄여야 한다면 MSA-CPT가 현실적이다. 실험은 두 route의 목적이 다르다는 점을 잘 보여준다.

5-2. Long-context extension

MSA-CPT checkpoint에 약 140B long-context tokens를 추가로 학습한 뒤 HELMET-128K와 RULER-128K를 평가한다.

Benchmark Subset Full MSA-CPT Delta
HELMET-128K Overall 46.53 45.93 -0.60
HELMET-128K ICL 70.40 72.80 +2.40
HELMET-128K Rerank/RAG 34.60 32.50 -2.10
RULER-128K Overall 72.00 72.12 +0.12
RULER-128K QA1/QA2 47.80 46.80 -1.00

Overall score만 보면 dense와 매우 가깝다. 그러나 subtask를 보면 retrieval과 QA에서 작은 gap이 남고, ICL과 multi-key retrieval에서는 오히려 높다.

이 결과의 핵심은 128K에서 무조건 dense와 동일하다는 것이 아니다. Query와 GQA group마다 2048 token만 Main Branch가 보는데도 aggregate long-context capability를 상당 부분 유지한다는 점이다. 동시에 sparse selector의 miss가 민감한 rerank, RAG, QA subtask는 별도 검증이 필요하다.

5-3. Efficiency

논문이 1M context와 H800에서 보고한 대표 결과는 다음과 같다.

Metric Reported result Scope
Per-token attention compute 28.4x reduction 109B experimental model, 1M context
Prefill wall-clock 14.2x speedup Co-designed MSA kernel on H800
Decode wall-clock 7.6x speedup Co-designed MSA kernel on H800

여기서 가장 중요한 것은 28.4x compute reduction과 14.2x 또는 7.6x wall-clock speedup이 같지 않다는 점이다. Sparse execution에는 Top-k, reverse index, query gather, partial output combine 같은 overhead가 남는다. 논문의 가치 중 하나는 이 차이를 숨기지 않고 kernel design으로 상당 부분 줄였다는 데 있다.

또한 이 숫자를 application end-to-end throughput으로 그대로 옮기면 안 된다. MoE routing, MLP, communication, multimodal encoder, scheduler overhead가 전체 serving time에서 차지하는 비중은 배포 환경마다 다르다.

5-4. What really matters in the experiments

1) LM loss and gradient norm

3T-token training 동안 MSA-PT와 Full-Attention의 LM loss가 거의 겹치고 gradient norm도 비슷한 범위를 유지한다. Sparse attention이 large-scale pretraining optimization을 불안정하게 만들지 않았다는 근거다.

2) Selection recall

MSA-CPT에서 warmup 뒤 KL loss가 낮게 유지되고, selected block이 Main Branch attention mass의 대부분을 회수한다. 이 분석은 downstream score보다 mechanism에 직접 연결된다. Sparse model의 성패는 selected token 수보다 중요한 attention mass를 얼마나 회수하는지에 달려 있기 때문이다.

3) Ablation quality

논문은 stop-gradient, warmup, local block, dynamic selection, block size, value head를 분리해 본다. 특히 warmup과 gradient detach가 단순한 training trick이 아니라 divergence와 short-context regression을 막는 필수 component라는 점을 보여준다.

4) Multimodal validation

Image와 video benchmark를 함께 평가한 점도 중요하다. Sparse selector가 text retrieval에만 맞고 visual token sequence에서는 실패할 수 있는데, MSA-PT가 여러 multimodal task에서 강한 결과를 보인다는 것은 group-specific block routing이 modality가 섞인 sequence에서도 작동할 가능성을 보여준다.

6. Limitations

  1. Index Branch는 full context를 scoring한다. Main Branch budget은 고정되지만 selector 자체는 긴 context를 훑는다. 따라서 MSA는 strict linear-time attention이라기보다 expensive attention path를 fixed budget으로 줄인 design에 가깝다.

  2. Conversion cost가 크다. MSA-CPT는 400B continued tokens를 사용하고 그중 40B를 indexer warmup에 쓴다. 기존 checkpoint를 즉시 sparse serving으로 바꾸는 plug-and-play 방법은 아니다.

  3. Long-context gap이 subtask별로 남는다. HELMET Rerank/RAG와 일부 QA에서 Full보다 낮다. Needle retrieval aggregate가 좋더라도 production RAG의 evidence recall을 따로 재야 한다.

  4. Attention speedup과 end-to-end speedup은 다르다. 논문 수치는 MSA attention path와 co-designed kernel의 효율을 보여준다. 실제 model serving에서는 MoE communication, MLP, batching, memory allocator가 병목이 될 수 있다.

  5. Hardware portability를 추가 확인해야 한다. 논문은 H800 결과를 보고하지만, 현재 공개 kernel repository의 requirement는 NVIDIA SM100이다. H800에서 사용한 exact kernel release와 현재 public package의 지원 범위가 동일한지는 배포 전 확인이 필요하다.

  6. 109B experiment와 public M3를 구분해야 한다. 논문의 controlled comparison은 109B model이고, 공개 M3 repository는 약 428B total parameters의 production model이다. M3 repository의 M2 대비 수치와 논문의 GQA 대비 수치를 같은 baseline으로 합치면 안 된다.

  7. Post-training과 agent rollout 검증은 제한적이다. 논문은 pretraining evaluation과 agent-oriented PPL을 폭넓게 제공하지만, RL post-training 이후 tool-use success rate나 multi-turn memory reliability는 후속 검증이 필요하다.

7. My Take

7-1. Why this matters for my work

MSA의 진짜 기여는 새로운 sparse score function보다 algorithm, training recipe, kernel을 같은 design space에서 최적화한 데 있다.

Efficient attention 논문은 종종 두 극단 중 하나에 머문다. Model quality는 자세히 보지만 실제 kernel이 없거나, kernel benchmark는 빠르지만 selector가 foundation model scale에서 어떻게 학습되는지 설명하지 않는다. MSA는 selector supervision, warmup, detach, block granularity, GPU iteration order를 하나의 chain으로 연결한다.

이 관점은 long-context model을 설계할 때 유용하다. Attention operator만 바꾸는 것이 아니라 다음 질문을 함께 봐야 한다.

  • Selection error는 어떤 objective로 줄이는가.
  • Sparse pattern은 hardware tile과 잘 맞는가.
  • Dense checkpoint를 재사용할 수 있는가.
  • Context가 길어질수록 selector cost가 새로운 병목이 되지 않는가.
  • Reported kernel speedup이 end-to-end serving에 얼마나 남는가.

7-2. Reuse potential

1) Existing GQA checkpoint conversion

MSA-CPT route는 이미 보유한 GQA model을 sparse serving으로 옮기는 reference recipe가 될 수 있다. 다만 400B continued training budget을 그대로 재현하기 어려운 조직이라면, smaller warmup budget에서 selection recall과 downstream regression curve를 먼저 측정해야 한다.

2) Repository-scale code and agent memory

Per-GQA-group routing은 syntax, global instruction, file-level evidence, recent execution trace처럼 서로 다른 memory role을 group별로 나눌 가능성이 있다. 실제 agent system에서는 selected block 로그를 남겨 retrieval failure를 debugging하는 것도 유용하다.

3) Multimodal long sequence

Video frame token과 text instruction이 섞인 sequence에서도 block routing을 적용할 수 있다. 다만 fixed block size가 frame boundary나 image patch layout과 맞지 않으면 semantic unit이 잘릴 수 있으므로 modality-aware block packing을 별도로 검토할 가치가 있다.

4) Serving policy

2048-token fixed budget을 모든 request에 동일하게 적용할 필요는 없다. Easy request에는 작은 Top-k를 쓰고 retrieval confidence가 낮을 때 budget을 늘리는 adaptive serving policy가 후속 방향이 될 수 있다. 논문 conclusion도 더 긴 sparse training, larger inference budget, richer indexer를 future direction으로 언급한다.

7-3. Practical evaluation checklist

MSA 계열을 실제 서비스에 도입한다면 benchmark score 하나보다 아래 네 축을 같이 봐야 한다.

Axis What to measure
Quality Task score, evidence recall, selection score recall
Cost Index path FLOPs, Main Branch FLOPs, KV traffic
Latency Prefill, decode, p50, p95, batch-size sensitivity
Portability GPU generation, dtype, compiler, serving framework integration

특히 p95 latency가 중요하다. 특정 popular block에 query가 몰리는 sparse pattern은 average latency는 좋아도 tail을 악화시킬 수 있다. MSA의 pre-scheduled tile chunking은 바로 이 문제를 겨냥한다.

7-4. Follow-up papers

8. Summary

  • MSA는 GQA group마다 KV block을 고르는 Index Branch와 selected support에서 exact softmax를 수행하는 Main Branch를 결합한다.
  • KL alignment, gradient detach, 40B-token warmup, forced local block이 discrete selector의 training stability를 만든다.
  • Exp-free Top-k와 KV-outer kernel이 theoretical sparsity를 H800 wall-clock speedup으로 연결한다.
  • 109B native multimodal model에서 MSA-PT는 여러 retrieval, image, video task에서 Full-Attention과 경쟁하거나 더 높은 결과를 보인다.
  • 가장 중요한 교훈은 sparse attention이 model architecture 하나가 아니라 training and serving system의 co-design 문제라는 점이다.

댓글남기기

참고