[vllm] vLLM의 멀티모달 추론 성능 극대화: Triton/FlashInfer 복합 어텐션 도입

PR 링크: vllm-project/vllm#56305 상태: Merged | 변경: +1095 / -3

들어가며

vLLM은 고성능 LLM 서빙을 위한 표준으로 자리 잡았지만, 멀티모달 모델(이미지+텍스트)의 프리픽스 어텐션(Prefix Attention) 처리 시 성능 병목이 존재했습니다. 특히 Blackwell 아키텍처와 같은 최신 GPU 환경에서 기존 Triton 기반 어텐션은 순수 FlashInfer 대비 효율성이 떨어지는 경우가 많았습니다. 이번 PR은 Triton의 유연한 이미지 마스킹 처리와 FlashInfer의 고속 인과적(Causal) 어텐션을 결합한 'Composite Backend'를 도입하여, 멀티모달 추론의 TTFT(Time To First Token)를 획기적으로 개선했습니다.

코드 분석

1. Composite Backend 설계 (vllm/v1/attention/backends/composite.py)

이 PR의 핵심은 create_composite_attention_backend를 통해 두 개의 서로 다른 어텐션 구현체를 하나의 워크플로우로 묶는 것입니다.

# Before: 단일 백엔드 선택
backend = get_attn_backend(head_size, dtype, kv_cache_dtype, use_mm_prefix=True)

# After: Composite 라우팅을 통한 동적 선택
# Triton은 이미지 마스킹이 필요한 쿼리에, FlashInfer는 텍스트 프리필/디코드에 할당
backend = create_composite_attention_backend(Triton, FlashInfer, MMPrefixAttentionRouting)

2. 자동 백엔드 선택 로직 (tests/kernels/attention/test_attention_selector.py)

시스템이 자동으로 최적의 백엔드를 선택하도록 로직이 개선되었습니다. 특히 Blackwell(SM100) 환경에서 TRITON_FLASHINFER가 우선순위를 갖도록 설정되었습니다.

# After: Blackwell 환경에서의 자동 선택 테스트
expected = "TRITON_FLASHINFER" if use_mm_prefix else "FLASHINFER"
assert backend.get_name() == expected

왜 이게 좋은가

이번 최적화의 핵심은 '적재적소의 커널 활용'입니다.

1. 성능 수치: FP8 텍스트 추론 시 TTFT가 기존 대비 34.5% 감소했습니다. 또한, 1024 requests 환경에서 처리량(Total tok/s)이 비약적으로 상승했습니다.
2. 유연성: 기존의 단일 백엔드 방식은 특정 커널이 지원하지 않는 기능(예: 특정 head size에서의 마스킹)이 있을 때 전체 성능이 저하되었습니다. Composite 방식은 이를 분리하여 각 커널의 강점만을 취합니다.
3. 일반적 교훈: 복잡한 추론 워크로드에서는 하나의 최적화된 커널을 고집하기보다, 데이터의 특성(이미지 마스킹 여부 등)에 따라 런타임에 커널을 교체하는 'Composite' 패턴이 성능 최적화의 핵심 전략이 될 수 있음을 보여줍니다.

리뷰어 피드백 반영

리뷰 과정에서 lesj0610님은 향후 도입될 Native FlashInfer 구현체와의 우선순위 충돌 문제를 제기했습니다. 이에 대해 커뮤니티는 현재의 Composite 구현을 먼저 도입하되, 향후 Native 지원이 완성되면 백엔드 우선순위 로직을 유연하게 조정하기로 합의했습니다. 이는 오픈소스 프로젝트에서 점진적 성능 개선을 달성하는 모범적인 협업 사례입니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.vllm.ai/en/latest/models/engine_args.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.