MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

The computational challenges of Large Language Model (LLM) inference remain a significant barrier to their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on a single A100 GPU. Existing methods for speeding up prefilling often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs. To address this gap, we introduce MInference (Milliontokens Inference), a sparse calculation method designed to accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices-the A-shape, Vertical-Slash, and Block-Sparsethat can be leveraged for efficient sparse computation on GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse indices based on the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels to significantly reduce the latency in the pre-filling stage of long-context LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications to the pre-training setup or additional fine-tuning. By evaluating on a wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, and Qwen2-128K, we demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy. Our code is available at https://aka.ms/MInference.

한 줄 요약

MInference는 10배 가속화된 A100 GPU 기반 1M 토큰 프리필링을 위한 동적 희소 어텐션 기법이다.

핵심 기여도

A-shape, Vertical-Slash, Block-Sparse의 세 가지 희소 어텐션 패턴을 식별하여 GPU 최적화 희소 계산을 구현.
각 어텐션 헤드에 대해 오프라인 최적 패턴 결정 및 인퍼런스 시 동적 인덱스 생성.
LLaMA-3-8B-1M 모델에서 A100 GPU 기준 30분 → 3분으로 10× 가속화.
기존 모델에 추가 학습 없이 바로 적용 가능.

핵심 아이디어

기존의 고정 희소 어텐션 방법은 입력에 따라 유연성이 부족하여 정확도나 효율성이 떨어진다. MInference는 이 문제를 해결하기 위해, 어텐션 행렬의 패턴이 입력에 따라 동적으로 변한다는 점을 활용한다. 본 연구는 Long-context LLM에서 일반적으로 나타나는 A-shape, Vertical-Slash, Block-Sparse 세 가지 희소 패턴을 발견하고, 이를 기반으로 각 어텐션 헤드에 최적의 패턴을 오프라인 결정한 후, 인퍼런스 시 동적으로 희소 인덱스를 생성한다. 이는 기존의 low-rank 기반 희소 추정 방법과 달리, 추정 오버헤드를 최소화하면서도 정확도를 유지한다. 특히, Block-Sparse 헤드는 64개 블록 단위로 mean pooling을 수행하여 블록 수준 어텐션 가중치를 계산하고, Vertical-Slash 헤드는 마지막 `last_q` 개의 쿼리와 키 벡터를 사용해 전역적으로 중요한 인덱스를 추정한다.

기술적 접근법

**패턴 탐지**: A-shape, Vertical-Slash, Block-Sparse 패턴을 Long-context 어텐션 행렬에서 식별.
**오프라인 최적화**: 각 어텐션 헤드에 대해 최적의 패턴을 결정.
**동적 마스크 생성**: 입력에 따라 Vertical-Slash, Block-Sparse 패턴 기반으로 희소 인덱스를 생성.
**GPU 최적화**: PIT, Triton, FlashAttention 기반의 세 가지 GPU 커널을 개발하여 희소 어텐션 계산 효율화.
**추정 오버헤드 최소화**: low-rank 기반 희소 추정 대신, 빠른 근사법으로 인덱스 생성.

주요 결과

**LLaMA-3-8B-1M 모델**: A100 GPU 기준 1M 토큰 프리필링 시간 30분 → 3분 (10× 가속).
**InfiniteBench, RULER, Needle In A Haystack, PG-19 등**: 정확도 유지 또는 개선.
**GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, Qwen-2-7B-128K 등 다양한 모델에서 실험 성공.**
**정적 마스크 사용 시 성능 급격히 저하** (예: KV retrieval에서 정확도 0%에 가까움).

의의 및 한계

MInference는 Long-context LLM의 프리필링 단계에서 발생하는 90% 이상의 어텐션 계산 오버헤드를 95%까지 줄이며, 기존 모델에 추가 학습 없이 바로 적용 가능한 점에서 실용적 가치가 크다. 특히, 다양한 모델과 벤치마크에서 정확도를 유지하면서 10배 가속화를 달성한 것은 학술적·산업적 측면에서 모두 의미가 있다. 그러나, 모든 어텐션 헤드가 동일한 패턴을 따르지 않으며, 일부 패턴 제거 시 성능 저하가 발생하므로 패턴 선택의 정확성에 의존한다는 한계가 있다. 또한, 다모달 LLM이나 인코더-디코더 모델에서의 확장성은 추가 연구가 필요하다.

실용적 활용

MInference는 대규모 문서 처리, 코드 이해, 멀티턴 QA, 장기 기억 기반 추론 등 Long-context가 필수적인 산업 및 연구 분야에 적용 가능하다. 특히, 실시간 응답이 필요한 애플리케이션에서 프리필링 지연을 최소화하여 사용자 경험을 개선할 수 있다.