LLM WikiAccess-protected knowledge portal

WIKI

FlashInfer 0.6: JIT 컴파일·MLA·MoE EP·Blackwell 지원으로 LLM 서빙의 커널 층을 다시 그은 어텐션 엔진

왜 지금 봐야 하나 LLM 추론의 병목은 GPU의 FLOPS 연산 능력 가 아니라 메모리 대역폭 이다. KV 캐시를 GPU 메모리에서 읽고, 어텐션 연산을 하고, 다시 쓰는 과정이 토큰 생성 속도를 결정한다. FlashInfer는 이 병목을 직접 공략하는 CUDA 커널 라이브러리다. v0.6.x 시리즈 2025년 12월~현재 는 LLM 서빙 스택에서 사실상의 표준 어텐션 엔진이 됐다. vLLM은 Blackwell B200/B3

경로human/study/content/ai-frontier/65-flashinfer-0-6-attention-kernel-mla-moe-blackwell.md
카테고리Study
태그#ai-review #attention #blackwell #kernel #mla #moe #study

왜 지금 봐야 하나

LLM 추론의 병목은 GPU의 FLOPS(연산 능력)가 아니라 메모리 대역폭이다. KV 캐시를 GPU 메모리에서 읽고, 어텐션 연산을 하고, 다시 쓰는 과정이 토큰 생성 속도를 결정한다. FlashInfer는 이 병목을 직접 공략하는 CUDA 커널 라이브러리다.

v0.6.x 시리즈(2025년 12월~현재)는 LLM 서빙 스택에서 사실상의 표준 어텐션 엔진이 됐다. vLLM은 Blackwell(B200/B300) 기본 백엔드로 FlashInfer를 채택했고, SGLang과 TensorRT-LLM도 기본 경로를 FlashInfer로 이전했다. DeepSeek V4-Flash의 MLA 커널, Gemma 4·MiniMax-M3의 MoE Expert Parallelism, Blackwell SM100/SM103 전용 CuTe-DSL 커널이 모두 이 시리즈에 포함됐다.

FlashInfer를 직접 다루지 않더라도, vLLM이나 SGLang을 운영하는 엔지니어라면 어텐션 레이어에서 이 라이브러리가 무슨 일을 하는지 알아야 한다.

배경: LLM 디코드의 병목 구조

Transformer 기반 LLM의 디코드(token generation) 단계를 GPU 관점에서 보면 매 스텝마다 두 단계가 반복된다.

  1. KV 캐시 읽기: 과거 모든 토큰의 Key·Value 행렬을 HBM(GPU 고대역폭 메모리)에서 불러온다.
  2. 어텐션 연산: 현재 토큰의 Query와 캐시된 Key·Value의 스케일드 닷 프로덕트를 계산한다.

단일 토큰을 생성할 때 이미 캐시된 N개 토큰의 KV 전체를 읽는다. 시퀀스가 길수록 읽어야 할 KV 데이터 크기가 선형으로 증가하고, 메모리 이동량이 병목이 된다. 같은 이유로 H100에서 실측 MFU(모델 연산 활용률)가 1% 미만으로 나오는 일도 흔하다.

PagedAttention(vLLM 초기 구현)은 KV 캐시를 고정 크기 페이지로 나눠 메모리 단편화를 줄였다. 하지만 페이지 관리 추상화 레이어와 실제 GPU 커널 사이의 간격이 성능 손실을 만들었다. FlashInfer는 이 레이어를 제거하고 블록-스파스 표현(block-sparse representation)으로 KV 캐시를 직접 모델링한다.

FlashInfer 아키텍처: 두 개의 층

1) 블록-스파스 KV 포맷

FlashInfer는 KV 캐시를 블록-스파스 행렬로 표현한다. Paged KV 캐시, Radix-tree KV 캐시, 연속(contiguous) KV 배열 등 다양한 메모리 레이아웃을 동일한 추상화로 처리한다. 페이지 인덱스를 GPU 공유 메모리에 사전 적재(pre-fetch)해 실제 어텐션 커널이 인덱싱 오버헤드 없이 KV 블록에 접근한다.

2) JIT 커널 컴파일

FlashInfer는 어텐션 커널을 사전 컴파일하지 않는다. 각 고유한 (num_heads, head_dim, page_size, dtype, attention_variant) 조합에 대해 첫 호출 시 CUDA 커널을 JIT 컴파일한다. 이 컴파일이 실제 배포 환경에 정확히 맞는 커널을 생성한다.

트레이드오프: 처음 만나는 형태마다 컴파일 지연이 발생한다(보통 1–3초). 워밍업(warm-up) 요청을 서비스 시작 시 미리 보내거나, JIT 캐시 패키지를 설치해 컴파일 결과물을 재사용하는 방식으로 완화한다.

서빙 프레임워크
vLLM
Blackwell 기본 백엔드
SGLang
기본 어텐션 백엔드
TensorRT-LLM
FP8/FP4 커널
TGI · MLC-LLM 외
→
FlashInfer 0.6 코어
블록-스파스 KV API
JIT 커널 컴파일러
JIT 캐시 (디스크)
MLA 커널 (DeepSeek)
MoE EP GEMM
FP8/FP4 All-reduce
→
GPU 지원 범위
Turing SM75
(T4, RTX 2080)
Ampere SM80
(A100, A6000)
Hopper SM90
(H100, H200)
Blackwell SM100+
(B200, B300)
FlashInfer 0.6 아키텍처 개요

커널 유형별 역할

Decode 어텐션

단일 토큰 생성 단계에서 사용한다. Query 길이가 1이고 KV 캐시 길이가 N이다. 메모리 대역폭이 지배적이라 HBM 읽기 효율이 핵심이다. FlashInfer는 GQA(Grouped Query Attention)와 MQA(Multi-Query Attention) 변형 모두를 지원하며 Blackwell에서는 CuTe-DSL 커널이 자동 선택된다.

Prefill 어텐션

입력 프롬프트 전체를 처리하는 단계다. 컴퓨트 집약적이다. FlashInfer는 청크드 프리필(chunked prefill)을 지원해 긴 프롬프트를 청크로 나눠 처리하고 중간 KV를 캐시한다. 이를 통해 긴 프롬프트가 디코드 중인 배치의 지연을 일으키는 헤드-오브-라인 블로킹을 완화한다.

MLA(Multi-head Latent Attention) 커널

DeepSeek V3/V4가 채택한 어텐션 변형이다. KV를 저차원 잠재 벡터로 압축해 저장한다. 디코드 시 잠재 벡터를 Key·Value로 복원하는 행렬 연산이 필요하다.

FlashInfer의 MLA 커널은 매트릭스 흡수 트릭(Matrix Absorption Trick)을 사용한다. 복원 행렬을 어텐션 연산에 미리 융합(fuse)해 별도의 메모리 쓰기 없이 잠재 벡터에서 직접 어텐션을 계산한다. 결과적으로 KV 캐시 저장 공간이 대폭 줄어들고, latent cache만 저장하는 KV-only 서빙이 가능해진다.

속성표준 MHAMLA (DeepSeek V3 기준)
KV 캐시 크기H × d_head × 2 (대용량)d_c (잠재 차원, 소용량)
캐시 절감 효과기준최대 87% 감소 보고
FlashInfer 지원기본v0.5+ (Blackwell B300: v0.6+)

MoE Expert Parallelism(EP)

Mixture-of-Experts 모델에서 전문가(expert)를 여러 GPU에 분산 배치하는 방식이다. v0.6.x에서 moe_ep가 기본 설치 패키지에 포함됐다. TRTLLM-GEN MoE 스택과 통합되며 Gemma 4, MiniMax-M3, DeepSeek-class 모델의 대규모 서빙에 적용된다.

Blackwell(SM100/SM103) 전용 경로

Hopper(SM90)까지는 CUTLASS와 Triton 기반 커널이 주를 이뤘다. Blackwell에서는 NVIDIA CuTe DSL(CUDA Template Dialect)과 새로운 WGMMA(Warpgroup Matrix Multiply-Accumulate) 명령을 활용해야 피크 성능에 근접할 수 있다.

CuTe-DSL GQA Decode 커널: CuTe DSL로 작성한 GQA 디코드 커널이다. 슬라이딩 윈도(sliding-window)와 어텐션 싱크(attention-sink) 마스킹을 지원해 Llama3, Gemma 4 등 긴 컨텍스트 모델에 적합하다.

SM103(B300) MLA 디코드: B300 아키텍처에 맞춘 MLA 디코드 경로다. CUTLASS split_kv를 클러스터 인식(cluster-aware) 방식으로 구현해 B300의 고대역폭 메모리에서 MLA 잠재 벡터 접근을 최적화한다.

CUDA 그래프 캡처 가능 FP8 All-reduce: 분산 추론(DP/TP)에서 CUDA 그래프로 캡처 가능한 FP8 all-reduce가 추가됐다. CUDA 그래프는 CPU 오버헤드를 제거하는 핵심 기법인데, 이전에는 FP8 all-reduce가 그래프 캡처 경계 밖에 있어 혼합 정밀도 분산 서빙에서 그래프를 쓸 수 없었다.

서빙 프레임워크별 통합 현황

프레임워크FlashInfer 역할기본 활성 여부
vLLMBlackwell SM100+ 기본 어텐션 백엔드Blackwell에서 자동 선택
SGLang기본 어텐션 백엔드 (MLA 포함)모든 GPU에서 기본
TensorRT-LLMFP8/FP4 어텐션·MoE GEMM 커널 공급TRTLLM-GEN 스택 통합
TGI선택 가능한 어텐션 백엔드설정에 따라 선택
MLC-LLM어텐션 커널 옵션 중 하나TVM 경로와 병행

운영 시 확인할 사항

JIT 컴파일 워밍업

서비스 시작 후 처음 만나는 배치 형태마다 컴파일 지연이 발생한다. 두 가지 완화 방법이 있다.

# 방법 1: JIT 캐시 패키지 설치 (사전 컴파일된 커널 번들)
pip install flashinfer-python[cache]

# 방법 2: 서버 시작 후 더미 워밍업 요청
# 일반적인 batch_size, seq_len 조합으로 몇 가지 요청을 미리 전송한다
curl -X POST http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"...", "prompt":"warmup", "max_tokens":1}'

버전 고정

FlashInfer는 vLLM·SGLang과 버전을 맞춰야 한다. 패치 버전 불일치가 CUDA 커널 ABI 충돌을 유발할 수 있다. 컨테이너 이미지 빌드 시 특정 버전을 명시적으로 고정한다.

# vLLM이 테스트한 버전을 확인하고 동일 버전으로 고정
# vLLM requirements.txt 또는 pyproject.toml에서 flashinfer 버전 확인 후 일치
pip install flashinfer-python==0.6.14

Blackwell 배포

B200/B300에서는 cute-dsl 백엔드가 자동 선택된다. SM90 이하 환경과 같은 FlashInfer 버전을 쓰더라도 커널 경로가 다르다. 스테이징에서 Blackwell 환경을 따로 검증한다.

지원 GPU 범위 확인

FlashInfer 0.6.x는 SM75(Turing) 이상만 지원한다. V100(SM70) 이하 GPU는 지원하지 않는다. 이전 세대 GPU를 운영 중이라면 대안 백엔드(CUTLASS, Triton)를 확인해야 한다.

v0.6.x 운영 체크리스트

한 문장으로

FlashInfer 0.6.x는 JIT 컴파일 기반 블록-스파스 KV 어텐션 엔진으로, MLA 매트릭스 흡수 트릭·MoE Expert Parallelism·Blackwell CuTe-DSL 커널을 통합해 vLLM·SGLang·TensorRT-LLM의 기본 어텐션 백엔드 자리를 잡은 LLM 서빙 커널 라이브러리다.

References