[vllm] ViT Full CUDA Graph - 비전 인코더 CUDA Graph 완전 지원

PR 링크: vllm-project/vllm#35963 상태: Merged | 변경: +1,584/-31

들어가며

멀티모달 LLM에서 비전 인코더(ViT)는 이미지 처리의 핵심이지만, 기존에는 CUDA Graph 최적화가 적용되지 않았다. 이 PR은 EncoderCudaGraphManager를 도입하여 ViT 인코더의 forward pass를 CUDA Graph로 캡처하고 리플레이함으로써 커널 런치 오버헤드를 제거한다.

핵심 코드 분석

Budget 기반 CUDA Graph 관리

class EncoderCudaGraphManager:
    @staticmethod
    def _generate_budgets(min_budget: int, max_budget: int) -> list[int]:
        # 2의 거듭제곱 단위로 budget 생성
        # [64, 128, 256, 512, 1024] 등
        ...

    def _find_smallest_fitting_budget_given_tokens(
        self, total_tokens: int
    ) -> int | None:
        # total_tokens 이상인 가장 작은 budget 선택
        # 없으면 None 반환 (eager fallback)
        ...

ViT 인코더는 이미지 크기에 따라 입력 토큰 수가 가변적이다. Budget 시스템은 2의 거듭제곱 크기로 미리 CUDA Graph를 캡처해두고, 실행 시 가장 가까운 budget을 선택하여 리플레이한다.

Hit/Miss 통계 추적

def get_cumulative_stats(self):
    total = self.graph_hits + self.graph_misses
    return {
        "graph_hits": self.graph_hits,
        "graph_misses": self.graph_misses,
        "hit_rate": self.graph_hits / total if total > 0 else 0.0,
        "num_budgets": len(self.budget_graphs),
        "token_budgets": self.token_budgets,
    }

실제 운영 중 CUDA Graph 캐시 적중률을 모니터링할 수 있어, budget 설정을 최적화하는 데 활용할 수 있다.

SupportsEncoderCudaGraph 프로토콜

모델이 SupportsEncoderCudaGraph 프로토콜을 구현하면 자동으로 CUDA Graph 캡처가 활성화된다. 이를 통해 다양한 비전 인코더 아키텍처에서 플러그인 방식으로 사용 가능하다.

왜 이게 좋은가

1. 커널 런치 오버헤드 제거: CUDA Graph 리플레이로 수백 개의 커널을 단일 그래프로 실행
2. 가변 입력 대응: Budget 시스템으로 다양한 이미지 크기를 효율적으로 처리
3. 모니터링: hit/miss 통계로 운영 중 성능 최적화 가능
4. 확장 가능한 설계: 프로토콜 기반으로 새로운 인코더 모델 쉽게 지원

정리

멀티모달 LLM 추론에서 비전 인코더가 병목이 되는 경우가 많다. 이 PR은 CUDA Graph를 비전 인코더에 적용하여 이 병목을 해소하며, 특히 반복적으로 유사한 크기의 이미지를 처리하는 서비스에서 큰 효과를 발휘한다.

참고 자료

• PR #35963
• GPU Model Runner V2 #25266

---

이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.