[vllm] NVIDIA RTX PRO 6000 GPU에서 VLLM의 행렬 곱셈 성능 최적화: sm120 아키텍처 지원 추가
PR 링크: vllm-project/vllm#57456 상태: Merged | 변경: +229 / -0
들어가며
최근 대규모 언어 모델(LLM)의 발전과 함께, 모델의 추론 속도 향상은 매우 중요한 과제가 되었습니다. VLLM은 LLM 추론을 위한 고성능 라이브러리로, 특히 행렬 곱셈(Matrix Multiplication, Matmul) 연산의 최적화에 많은 노력을 기울이고 있습니다. 이번 PR은 NVIDIA의 최신 GPU 아키텍처인 Blackwell 워크스테이션/컨슈머 파트(Compute Capability 12.x, 예: RTX PRO 6000, RTX 50 시리즈)를 위한 행렬 곱셈 커널 최적화를 다룹니다. 기존에는 이 아키텍처에서 VLLM_BATCH_INVARIANT=1 설정 시 기본 128x128x64 / 8 warps / 3 stages 설정으로 모든 선형 레이어가 실행되어 최적의 성능을 내지 못했습니다. 이 PR은 sm120 아키텍처 패밀리를 추가하고, Qwen3 모델의 주요 레이어들에 대한 최적화된 설정을 제공함으로써 이러한 문제를 해결합니다.
코드 분석
이번 PR의 핵심 변경 사항은 vllm/model_executor/determinism/batch_invariant_configs.py 파일에 집중되어 있습니다. 이 파일은 다양한 GPU 아키텍처 및 행렬 크기에 따른 최적의 행렬 곱셈 커널 설정을 관리합니다.
batch_invariant_configs.py 파일 변경 분석
가장 중요한 변경은 _MATMUL_SHAPE_CONFIGS 딕셔너리에 sm120 아키텍처에 대한 새로운 설정을 추가한 것입니다. 이 설정은 Qwen3 모델의 15가지 주요 GEMM(General Matrix Multiply) 형태에 대해 최적화된 BLOCK_K, BLOCK_M, BLOCK_N, num_warps, num_stages 값을 정의합니다.
Before:
기존 코드에서는 _get_tuned_matmul_arch_family 함수가 Compute Capability 12.x를 직접적으로 지원하는 sm120 패밀리를 가지고 있지 않았습니다. 따라서 해당 GPU에서는 기본 설정이 사용되었습니다.
def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
if capability is None:
return None
# sm120 family is missing here
if capability.major == 10:
return "blackwell"
if capability.major == 9:
return "ada"
# ... other architectures
After:
_get_tuned_matmul_arch_family 함수에 capability.major == 12일 경우 sm120을 반환하는 로직이 추가되었습니다. 또한, _MATMUL_SHAPE_CONFIGS 딕셔너리에 sm120 키 아래에 Qwen3 모델의 각 레이어(예: gate_up_proj, down_proj, qkv_proj, lm_head, o_proj)에 대한 구체적인 _MatmulShapeConfig가 정의되었습니다. 각 _MatmulShapeConfig는 block_k 값과 m_buckets를 포함합니다. m_buckets는 행렬의 첫 번째 차원(M)의 크기에 따라 다른 BLOCK_M, BLOCK_N, num_warps, num_stages 설정을 적용하도록 합니다. 이는 배치 크기(M)에 따라 동적으로 최적의 커널 설정을 선택할 수 있게 하여 성능을 극대화합니다.
_MATMUL_SHAPE_CONFIGS: dict[str, dict[tuple[int, int], _MatmulShapeConfig]] = {
# ... other architectures
"sm120": {
(12288, 2048): _MatmulShapeConfig( # Qwen3-1.7B gate_up_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 128, 8, 5)),
(4, _MatmulMConfig(16, 128, 8, 4)),
# ... more buckets
),
),
(2048, 6144): _MatmulShapeConfig( # Qwen3-1.7B down_proj
block_k=64,
m_buckets=(
(1, _MatmulMConfig(16, 32, 8, 5)),
# ... more buckets
),
),
# ... other shapes for sm120
},
}
def _get_tuned_matmul_arch_family(capability: DeviceCapability | None) -> str | None:
if capability is None:
return None
if capability.major == 12:
return "sm120"
if capability.major == 10:
return "blackwell"
# ... other architectures
이 변경은 RTX PRO 6000과 같은 sm_120 아키텍처를 사용하는 GPU에서 VLLM이 더 효율적인 행렬 곱셈 커널을 선택하도록 보장합니다. 특히, VLLM_BATCH_INVARIANT=1 모드에서 각 행렬 크기(N, K)에 대해 최적의 BLOCK_K를 선택하고, M 버킷별로 최적의 BLOCK_M, BLOCK_N, num_warps, num_stages를 선택함으로써 성능을 향상시킵니다.
왜 이게 좋은가?
성능 향상
이 PR은 Qwen3 모델을 RTX PRO 6000 (sm_120) GPU에서 실행할 때 상당한 성능 향상을 가져옵니다. PR 설명에 따르면, 다양한 워크로드에서 다음과 같은 성능 개선이 관찰되었습니다:
-
Decode-heavy 워크로드 (기본값: input 32 / output 128 / batch 8):
- Qwen3-1.7B: 10.1% 향상
- Qwen3-4B: 21.1% 향상
- Qwen3-8B: 16.3% 향상
-
Prefill-heavy 워크로드 (input 1024 / output 32 / batch 8):
- Qwen3-1.7B: 7.0% 향상
- Qwen3-4B: 14.9% 향상
- Qwen3-8B: 11.1% 향상
이러한 성능 향상은 VLLM_BATCH_INVARIANT=1 모드에서 기존의 기본 설정 대신 sm_120 아키텍처에 맞게 튜닝된 커널 설정을 사용함으로써 달성되었습니다. 특히, VLLM_USE_BREAKABLE_CUDAGRAPH=1 옵션을 사용하면 decode-heavy 워크로드에서 성능 향상이 최대 64%까지 증가하며, 배치 불변성(batch invariance)으로 인한 오버헤드가 1.00x ~ 1.02x 수준으로 거의 사라집니다. 이는 동적 배치 크기에 따른 최적의 커널 설정을 선택하는 것이 얼마나 중요한지를 보여줍니다.
일반적인 교훈
- 아키텍처별 최적화의 중요성: GPU 아키텍처는 매우 다양하며, 각 아키텍처의 특성에 맞는 커널 튜닝은 성능 향상의 핵심입니다. 새로운 아키텍처(여기서는 sm_120)가 등장하면 해당 아키텍처에 최적화된 설정을 추가하는 것이 필수적입니다.
- 동적 커널 설정의 이점: 행렬 곱셈과 같은 기본 연산에서 배치 크기(M), 행렬 차원(N, K) 등 입력에 따라 최적의 커널 파라미터(
BLOCK_M,BLOCK_N,num_warps,num_stages,BLOCK_K)가 달라집니다.m_buckets와 같이 입력 크기에 따라 다른 설정을 적용하는 메커니즘은 다양한 시나리오에서 일관된 고성능을 보장하는 데 매우 효과적입니다. - 정확성 검증의 중요성: 성능 최적화 과정에서 정확성(correctness)을 보장하는 것이 무엇보다 중요합니다. 이 PR에서는
fp32참조와의allclose및 비트 단위 동일성(bitwise identical) 검증을 통해 최적화된 커널의 정확성을 철저히 검증했습니다. - 점진적 개선: 기존의 최적화된 설정(Ada, Hopper, SM100)과 유사한 구조와 방법을 사용하여 새로운 아키텍처(sm120)에 대한 최적화를 점진적으로 확장하는 것은 효율적인 개발 전략입니다.
리뷰 피드백 반영
리뷰어 LioEinaudi는 특히 엔드-투-엔드(end-to-end) 성능 측정 결과에 대한 중요한 통찰을 제공했습니다. 초기 측정 결과에서 VLLM_COMPILE 사용 시 성능 향상이 제한적이었던 이유는, 컴파일 시점에 고정된 M 값에 대한 커널 설정만 그래프에 포함되었기 때문입니다. 하지만 VLLM_USE_BREAKABLE_CUDAGRAPH=1 옵션을 사용하여 런타임에 동적으로 M 값을 기반으로 커널 설정을 선택하도록 했을 때, PR에서 제시된 성능 향상 수치가 실제로 달성됨을 확인했습니다. 이는 커널 최적화가 실제 추론 환경에서 동적으로 적용될 때 그 효과를 제대로 발휘할 수 있음을 보여줍니다. 또한, TP(Tensor Parallelism) 설정 시 일부 per-rank 연산이 이 테이블에 포함되지 않아 성능 향상이 제한될 수 있다는 점을 지적했으며, 이는 향후 추가적인 최적화 작업의 방향을 제시합니다.
References
- NVIDIA Blackwell Architecture
- VLLM Batch Invariant Matmul
- Triton Kernel Tuning (VLLM이 사용하는 Triton 커널의 예시)
참고 자료
- https://www.nvidia.com/en-us/geforce/technologies/blackwell/
- https://github.com/vllm-project/vllm/wiki/Batch-Invariant-Matmul
- https://github.com/openai/triton/blob/main/python/triton/ops/matmul.py
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
PR Analysis 의 다른글
- 이전글 [sglang] DeepSeek-V4.1 성능 최적화: mHC 및 메타데이터 오버헤드 개선 분석
- 현재글 : [vllm] NVIDIA RTX PRO 6000 GPU에서 VLLM의 행렬 곱셈 성능 최적화: sm120 아키텍처 지원 추가
- 다음글 [sglang] ROCm 환경에서 SGLang HiCache JIT 전송 커널 최적화 및 유연성 개선
댓글