[vllm] vLLM의 혁신: Breakable CUDA Graph로 LLM 추론 성능 최적화

PR 링크: vllm-project/vllm#42304 상태: Merged | 변경: +0 / -0

들어가며\n대규모 언어 모델(LLM)의 추론 성능을 극대화하는 것은 매우 중요한 과제입니다. vLLM은 높은 처리량과 낮은 지연 시간을 제공하여 이 분야에서 선두를 달리고 있습니다. 이러한 성능 향상을 위해 torch.compile과 CUDA Graph와 같은 최적화 기술이 활발히 연구되고 적용됩니다. 그러나 복잡한 LLM 모델 아키텍처와 동적인 연산들로 인해 torch.compile이 항상 완벽한 CUDA Graph를 생성하기 어렵다는 문제가 있었습니다.\n\n이번에 vLLM 프로젝트에 병합된 PR "[Experimental] Breakable CUDA graph"는 이러한 문제를 해결하기 위한 혁신적인 접근 방식을 제시합니다. 이 PR은 torch.compile의 fullgraph 모드가 달성하기 어려운 시나리오에서, CUDA Graph 내부에 'eager break' 지점을 명시적으로 허용함으로써, 전체 연산의 상당 부분을 CUDA Graph로 실행하면서도 특정 동적 연산을 유연하게 처리할 수 있도록 합니다. 이는 LLM 추론의 안정성과 성능을 동시에 향상시키는 중요한 최적화입니다.\n\n## 문제점: torch.compile과 CUDA Graph의 한계\nPyTorch의 torch.compile은 모델을 컴파일하여 더 효율적인 커널로 변환함으로써 성능을 향상시키는 강력한 도구입니다. 특히 fullgraph=True 옵션은 모델의 전체 순전파를 단일 CUDA Graph로 캡처하여 오버헤드를 최소화하고 최대의 성능을 끌어내려 합니다. CUDA Graph는 GPU 커널 실행의 CPU 오버헤드를 제거하여 반복적인 워크로드에서 상당한 속도 향상을 제공합니다.\n\n하지만 LLM 모델은 종종 동적인 텐서 크기 변경, 조건부 분기, 또는 torch.cat, torch.split과 같이 그래프 캡처에 어려움을 주는 연산들을 포함합니다. 이러한 연산들은 torch.compile의 dynamo 백엔드에서 'graph break'를 유발하거나, fullgraph 모드 달성을 방해하여 결국 모델의 일부 또는 전체가 eager execution으로 폴백되게 만듭니다. 이는 CUDA Graph의 잠재력을 완전히 활용하지 못하게 하며, 결과적으로 추론 성능 저하로 이어집니다.\n\n이번 PR의 설명에 따르면, VLLM_USE_BREAKABLE_CUDAGRAPH=1을 설정하면 torch.compile의 dynamo 및 inductor가 관여하지 않도록 breakable cudagraph와 torch.compile fullgraph가 상호 배타적으로 동작합니다. 이는 breakable cudagraph가 torch.compile의 대안적인 CUDA Graph 캡처 메커니즘으로 작동하며, torch.compile이 fullgraph를 달성하기 어려운 모델에서 더 나은 성능을 제공할 수 있음을 시사합니다.\n\n## 해결책: Breakable CUDA Graph의 등장\n'Breakable CUDA Graph'는 이러한 문제를 해결하기 위해, CUDA Graph 캡처 도중에 특정 연산을 eager mode로 실행할 수 있도록 허용하는 메커니즘입니다. 이는 전체 연산 흐름을 하나의 큰 그래프로 캡처하려 시도하되, 그래프 캡처를 방해하는 연산들은 일시적으로 그래프 캡처를 중단하고 eager mode로 실행한 후, 다시 그래프 캡처를 재개하는 방식입니다. 이를 통해 대부분의 연산은 CUDA Graph의 이점을 누리면서도, 그래프에 포함될 수 없는 연산들은 유연하게 처리할 수 있게 됩니다.\n\n핵심 아이디어는 다음과 같습니다:\n1. Context Manager 기반 캡처: BreakableCUDAGraphCapture라는 컨텍스트 매니저를 사용하여 CUDA Graph 캡처 세션을 시작하고 종료합니다.\n2. Eager Break 지점 명시: add_eager 메서드나 @eager_break_during_capture 데코레이터를 사용하여 그래프 캡처 도중 eager mode로 실행할 함수를 명시합니다.\n3. 세그먼트 관리: 캡처된 그래프와 eager 연산들을 순서대로 세그먼트 리스트에 저장합니다.\n4. Replay: replay() 메서드를 호출하면 저장된 세그먼트 리스트를 순서대로 실행합니다. 이때 그래프 세그먼트는 CUDA Graph로, eager 세그먼트는 일반 PyTorch 연산으로 실행됩니다.\n\n## 코드 분석\n### vllm/compilation/breakable_cudagraph.py: 핵심 로직 구현\n이 파일은 Breakable CUDA Graph의 핵심 로직을 담고 있습니다. BreakableCUDAGraphCapture 클래스는 컨텍스트 매니저로 작동하며, eager_break_during_capture 데코레이터는 특정 함수를 그래프 브레이크 지점으로 표시합니다.\n\n**BreakableCUDAGraphCapture 클래스:**\n```python\nclass BreakableCUDAGraphCapture:\n _tls = threading.local()\n _tls.active = None\n\n def enter(self):\n if BreakableCUDAGraphCapture._tls.active is not None:\n raise RuntimeError(

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs
• https://pytorch.org/docs/stable/generated/torch.cat.html
• https://pytorch.org/docs/stable/generated/torch.split.html

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