FLASHINFER: EFFICIENT AND CUSTOMIZABLE ATTENTION ENGINE FOR LLM INFERENCE SERVING

배경 및 개요

LLM(거대언어모델)의 핵심에는 트랜스포머 아키텍처와 그 심장부인 어텐션 메커니즘이 자리 잡고 있습니다. 특히 LLM을 실제 서비스로 제공하는 ‘추론 서빙(Inference Serving)’ 환경에서, 어텐션 연산의 효율성은 전체 시스템의 처리량(Throughput)과 지연 시간(Latency)을 결정하는 가장 중요한 요소입니다.

그러나 LLM 추론 환경은 다음과 같은 두 가지 큰 도전 과제 때문에 기존의 어텐션 최적화 방식으로는 한계가 명확했습니다.

다양하고 동적인 워크로드: LLM 서빙은 사용자의 초기 프롬프트를 처리하는 ‘프리필(Prefill)’ 단계와, 다음 토큰을 순차적으로 생성하는 ‘디코딩(Decoding)’ 단계를 포함하며, 접두사 공유(Prefix-sharing), 추측적 디코딩(Speculative Decoding) 등 다양한 최적화 기법으로 인해 어텐션 연산 패턴이 매우 복잡하고 동적으로 변합니다.

하드웨어 및 스토리지 맞춤화 요구: 현대 GPU 아키텍처의 성능을 최대한 활용하고, 페이지드 어텐션(Paged Attention)과 같은 효율적인 KV 캐시 스토리지 포맷을 지원하기 위해서는 어텐션 연산자(Kernel)의 세밀한 맞춤화가 필수적입니다.

기존의 LLM 서빙 프레임워크들은 이러한 복잡성에 대응하기 위해 각자 파편화된 전용 어텐션 솔루션을 구현해야 했고, 이는 높은 유지보수 비용과 비효율성을 야기했습니다.

이러한 문제를 해결하기 위해, 이 논문은 FlashInfer를 제안합니다. FlashInfer는 LLM 추론 서빙을 위해 특별히 설계된 고도로 효율적이고, 유연하며, 사용자 정의 가능한 차세대 어텐션 엔진입니다. 이는 단순히 또 하나의 최적화된 어텐션 커널이 아니라, 다양한 KV 캐시 스토리지 형식, 변화무쌍한 어텐션 연산 패턴, 그리고 최신 GPU 하드웨어의 특성을 모두 아우르는 통합적인 코드 생성 기반(code-generation based) 프레임워크입니다. FlashInfer는 SGLang, vLLM, MLC-Engine과 같은 주요 LLM 서빙 프레임워크에 이미 통합되었으며, 기존 SOTA 솔루션 대비 획기적인 성능 향상을 입증했습니다.

주요 성능 테스트 결과

FlashInfer의 성능은 커널 수준과 엔드투엔드(End-to-end) 서빙 시스템 수준 모두에서 기존 SOTA(State-of-the-Art) 솔루션과 비교하여 평가되었습니다.

1. 엔드투엔드 LLM 서빙 성능 (SGLang 통합 기준)

일반 서빙 벤치마크: Triton 백엔드를 사용하는 SGLang 대비, FlashInfer 백엔드는 토큰 간 지연 시간(Inter-token Latency, ITL)을 29–69% 감소시키는 압도적인 성능 향상을 보였습니다.

긴 컨텍스트 추론 (Long-Context Inference): Streaming-LLM과 같은 긴 컨텍스트 추론 시나리오에서, FlashInfer의 사용자 정의 융합 커널(Fused Kernel)은 기존 방식 대비 28–30%의 지연 시간 감소 효과를 가져왔습니다.

병렬 생성 (Parallel Generation): 여러 개의 응답을 동시에 생성하는 병렬 디코딩 작업에서, FlashInfer의 ‘조합 가능한 포맷(Composable Formats)’은 기존 방식 대비 13–17%의 속도 향상을 달성했습니다.

2. 커널 수준 성능

동적 워크로드 처리: 균일하지 않은 다양한 시퀀스 길이의 요청이 섞여 들어오는 현실적인 워크로드에서, FlashInfer의 동적 스케줄링 커널은 기존 FlashAttention 커널보다 월등한 하드웨어 활용률(Bandwidth/FLOPS Utilization)을 보였습니다.

사용자 정의 커널의 효과: RoPE와 어텐션을 융합한 커스텀 커널은, 이를 분리해서 실행하는 기존 FlashAttention 방식보다 1.6배에서 3.7배 더 높은 대역폭 활용률을 달성하여, JIT 컴파일의 강력한 효과를 입증했습니다.

연구 방법론: FlashInfer의 아키텍처

FlashInfer의 혁신적인 성능은 세 가지 핵심적인 설계 원칙에 기반합니다.

1. 통합된 KV 캐시 추상화: 블록-희소(Block-Sparse) 포맷

문제점: 페이지드 어텐션, 레딕스 트리 등 다양한 KV 캐시 스토리지 구조는 각각 다른 메모리 접근 패턴을 요구하여 통합된 최적화가 어려웠습니다.

FlashInfer의 해결책: 이 모든 다양한 구조를 ‘블록-희소 행렬(Block-Sparse Matrix)’이라는 단일한 수학적 개념으로 추상화합니다. 이를 통해 어떤 복잡한 스토리지 구조라도 통일된 방식으로 처리할 수 있는 범용 어텐션 커널을 설계할 수 있게 됩니다. 또한, ‘조합 가능한 포맷(Composable Formats)’을 도입하여, 공유 접두사(shared prefix)와 같이 밀집된(dense) 부분과 개별적인 희소(sparse) 부분을 분리하여 각각 최적의 블록 크기로 처리함으로써 메모리 접근 효율을 극대화합니다.

2. 유연한 어텐션 변형 지원: 사용자 정의 가능한 템플릿과 JIT 컴파일

문제점: 그룹드-쿼리 어텐션(GQA), 슬라이딩 윈도우 어텐션 등 새로운 어텐션 변형이 계속 등장하여 기존 라이브러리로는 신속한 대응이 어려웠습니다.

FlashInfer의 해결책: 사용자가 어텐션 연산의 각 단계를(예: Logits 변환, 마스킹) CUDA 코드로 직접 정의할 수 있는 ‘어텐션 템플릿’ 프로그래밍 인터페이스를 제공합니다. 개발자가 이 템플릿을 제출하면, FlashInfer의 JIT(Just-In-Time) 컴파일러가 이를 즉석에서 고도로 최적화된 하드웨어 특화 커널 코드로 변환하고 컴파일합니다. 이는 라이브러리를 재컴파일할 필요 없이, 거의 무한한 종류의 커스텀 어텐션을 최고의 성능으로 실행할 수 있게 합니다.

3. 동적 워크로드 대응: 동적 부하 분산 스케줄링

문제점: 실시간으로 변하는 요청들의 시퀀스 길이에 맞춰 작업을 GPU 코어에 효율적으로 분배하는 것은 매우 어려운 문제입니다. 특히 정적인 실행 계획을 요구하는 CUDAGraph와 같은 최적화 기법과의 호환성 문제가 있었습니다.

FlashInfer의 해결책: 컴파일 시점의 타일 크기 결정과 런타임 시점의 스케줄링을 분리하는 동적 부하 분산 스케줄러를 도입했습니다. 이 스케줄러는 런타임에 각 요청의 시퀀스 길이를 입력받아, 모든 GPU 코어에 작업량이 균등하게 분배되도록 실시간으로 작업 계획을 수립합니다. 이 ‘계획’ 자체는 CPU에서 생성되지만, 생성된 계획은 GPU에 캐싱되어 CUDAGraph와 호환되는 고정된 커널 호출을 통해 실행되므로, 유연성과 고성능을 동시에 달성합니다.

결론 및 시사점

결론:

FlashInfer는 현대 LLM 추론 서빙 환경의 복잡성과 동적인 요구사항을 해결하기 위한 근본적인 아키텍처 혁신을 제시합니다. 블록-희소 포맷을 통한 통일된 스토리지 추상화, JIT 컴파일을 통한 무한한 커스터마이징, 그리고 동적 스케줄링을 통한 런타임 최적화를 결합함으로써, FlashInfer는 기존의 어떠한 어텐션 라이브러리보다 더 빠르고, 더 유연하며, 더 효율적인 솔루션을 제공합니다.

주요 서빙 엔진들과의 성공적인 통합과 압도적인 성능 향상은 FlashInfer가 LLM 추론, 특히 대규모, 롱-컨텍스트, 에이전트 기반 애플리케이션의 성능을 한 단계 끌어올릴 핵심적인 기반 기술임을 명백히 증명합니다.

시사점:

어텐션 엔진의 새로운 표준: FlashInfer의 코드 생성 및 JIT 컴파일 접근 방식은, 미리 컴파일된 고정된 커널 라이브러리를 제공하는 기존 방식(예: FlashAttention)을 넘어, 차세대 어텐션 엔진의 새로운 표준이 될 가능성이 높습니다.

소프트웨어와 하드웨어의 공동 설계 가속화: 개발자가 JIT 컴파일을 통해 최신 GPU 하드웨어의 기능(예: TMA, 워프 특수화)을 직접 활용할 수 있게 함으로써, 소프트웨어가 하드웨어의 발전을 더 빠르게 따라잡고 그 성능을 최대한 활용할 수 있는 길을 열었습니다.

LLM 서빙 기술의 민주화 및 혁신 촉진: 연구자와 개발자들이 자신만의 혁신적인 어텐션 메커니즘을 아이디어 단계에서 즉시 고성능 커널로 구현하고 테스트할 수 있게 되었습니다. 이는 커뮤니티 주도의 LLM 최적화 연구를 크게 활성화시킬 것입니다.

시스템의 ‘정적’ 부분과 ‘동적’ 부분의 분리: CUDAGraph와 같은 정적 최적화 기법을 활용하면서도 동적인 워크로드를 처리하기 위해, ‘계획(CPU, 동적)’과 ‘실행(GPU, 정적)’을 분리한 FlashInfer의 런타임 아키텍처는 다른 AI 시스템 설계에도 중요한 영감을 줄 수 있습니다.

LLM 추론의 병목 지점 이동: FlashInfer가 어텐션 연산의 병목을 상당 부분 해결함에 따라, 이제 LLM 추론 시스템의 다음 최적화 대상은 다른 부분(예: 모델 가중치 로딩, 네트워크 오버헤드 등)으로 이동하게 될 것입니다.

출처: https://arxiv.org/pdf/2501.01005

FlashInfer #LLMInference #AttentionMechanism #GPU #HighPerformanceComputing #SystemsForML

플래시인퍼 #LLM추론 #어텐션메커니즘 #GPU #고성능컴퓨팅 #머신러닝시스템