vLLM Triton Attention Backend 심층 분석

vLLM은 다양한 GPU 하드웨어 전반에 걸쳐 고성능 및 휴대 가능한 어텐션 커널을 제공하기 위해 Triton 기반 어텐션 백엔드를 도입했습니다. 이 백엔드를 통해 vLLM은 NVIDIA, AMD, Intel GPU에서 효율적으로 실행되는 어텐션 연산을 위한 단일 코드베이스를 유지할 수 있어, 하드웨어별 커널 유지 관리 부담을 줄이면서도 특화된 구현과 동등한 성능을 제공합니다.

Triton을 통한 성능 이식성

NVIDIA Hopper, Blackwell, AMD MI300 등 모든 GPU 아키텍처마다 수백 개의 특화된 커널을 유지하는 것은 비현실적입니다. vLLM은 GPU 커널을 Python으로 작성하고 여러 플랫폼에 대해 효율적인 코드로 컴파일할 수 있는 도메인 특화 언어인 Triton을 활용합니다.

Triton은 개발자가 논리적인 연산 타일을 정의하는 타일링 프로그래밍 모델을 사용합니다. Triton 컴파일러와 자동 튜너가 이러한 타일을 하드웨어별 실행 레이아웃에 매핑합니다. 이 추상화 덕분에 백엔드는 하드웨어에 구애받지 않으면서도 저수준 최적화를 수행할 수 있습니다.

Triton 어텐션 백엔드 아키텍처

어텐션 구현을 선형 레이어와 같은 다른 구성 요소와 분리하기 위해 vLLM은 어텐션 백엔드 추상화 레이어를 사용합니다. Triton 어텐션 백엔드는 vLLM에 네이티브이며, PyTorch와 Triton만을 의존하고 전적으로 Triton으로 구현되었습니다.

사용 시나리오

• Default for AMD GPUs: AMD GPU에서 ROCm을 사용할 때 기본 백엔드입니다.
• Intel XPU: float32 연산에 사용되며, 이 플랫폼에서는 FlashAttention이 fp32를 지원하지 않습니다.
• Specific Model Features: ALiBi sqrt( StepFun 오디오 모델에서 사용)와 sink 토큰, GPT-OSS 동작을 지원하며, 특히 Hopper 이전 NVIDIA GPU(예: A100)에서 유용합니다.
• Specialized Requirements: 작은 헤드 크기를 가진 모델, 인코더/디코더 어텐션, 멀티모달 프리픽스 어텐션을 처리합니다.
• Fallback Mechanism: FlashAttention, FlashInfer 또는 기타 종속성이 없을 경우 일반적인 대체 백엔드 역할을 합니다.

페이지드 어텐션의 기술 구현

페이지드 어텐션은 KV 캐시를 페이지 단위로 관리하여 메모리를 최적화합니다. 커널은 쿼리 토큰을 처리하고, 쿼리와 KV 헤드를 순회하며, 페이지드 KV 캐시를 탐색해 어텐션 점수를 계산합니다.

Q 블록 최적화

tl.dot(Triton의 행렬 곱셈) 활용도를 극대화하기 위해 백엔드는 "Q 블록"을 사용합니다. KV 캐시 페이지 크기가 KV 측 타일 크기를 제한하기 때문에, 커널은 여러 쿼리 토큰과 헤드를 하나의 작업 항목(Q 블록)으로 묶어 병렬성을 높이고 캐시 재사용을 증대합니다. 이는 특히 Group Query Attention(GQA)에서 유용합니다.

병렬 타일 Softmax (3D 커널)

Q 블록은 프리필 작업에 유리하지만, 디코드 작업은 단일 쿼리 토큰만 처리합니다. 이를 최적화하기 위해 vLLM은 병렬 타일 Softmax를 활용한 "3D 커널"을 구현했습니다. 이 방식은 KV 캐시 탐색을 여러 커널 인스턴스로 분할하고, 각 인스턴스가 부분 결과를 계산한 뒤 두 번째 커널 실행을 통해 최종 출력을 합산합니다.

지속 커널 및 CUDA 그래프

CUDA 그래프는 고정된 실행 그래프를 기록함으로써 런치 오버헤드를 감소시킵니다. 그러나 표준 어텐션 커널은 배치 크기와 시퀀스 길이에 따라 가변적인 런치 그리드를 사용하므로, CUDA 그래프를 재생할 때 비효율(예: 작업 낭비 또는 스트리밍 멀티프로세서 활용 저조)이 발생합니다.

이를 해결하기 위해 vLLM은 persistent kernels를 개발했습니다. 가변 그리드 대신, 사용 가능한 연산 자원 수와 동일한 고정된 수의 커널 인스턴스를 실행합니다. 이러한 인스턴스는 GPU 메모리에서 메타데이터를 동적으로 읽어 작업량을 결정함으로써 일정한 런치 그리드를 유지하고 CUDA 그래프 재사용을 효율적으로 합니다.

벤치마킹 및 결과

2025년 말의 벤치마크 결과, Triton 백엔드는 특화된 대안에 비해 코드 복잡도가 크게 낮으면서도 높은 효율성을 달성함을 보여줍니다. Triton 페이지드 어텐션 구현은 약 800줄의 코드로 구성되는 반면, FlashAttention-3은 약 70,000줄에 달합니다.

Performance Metrics (Llama 3.1 8B, 입력 토큰 500개, 배치 크기 1):

• NVIDIA H100: 긴 디코드 요청에서 FlashAttention-3 대비 100.7%의 성능을 달성했습니다.
• AMD MI300: 이전 구현 대비 약 5.8배의 속도 향상을 보였습니다.

향후 방향: Helium

실험적인 작업이 PyTorch 팀에서 만든 새로운 도메인 특화 언어인 Helium과 함께 시작되었습니다. Helium은 고수준 Triton 또는 타일링된 PyTorch로 설명됩니다. Helium으로 구현된 간소화된 페이지드 어텐션 커널은 초기 결과가 유망하며 현재 vLLM 저장소의 초안 풀 리퀘스트로 제공되고 있습니다.