TorchInductor를 위한 새로운 CuteDSL 백엔드: Blackwell GPU에서 GEMM 성능 극대화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

PyTorch의 컴파일러인 TorchInductor가 NVIDIA의 CuteDSL을 네 번째 행렬 곱셈(GEMM) 백엔드로 통합했다. 기존 CUTLASS C++ 백엔드는 nvcc 컴파일 속도가 느려 오토튜닝 효율이 낮았으나, CuteDSL은 파이썬 기반 인터페이스를 통해 컴파일 시간을 50배 이상 단축하면서도 하드웨어 제어력을 유지한다. NVIDIA B200 GPU 벤치마크 결과, BF16 및 MXFP8 데이터 형식에서 최대 1.78배의 커널 성능 향상을 기록했으며 vLLM 추론 지연 시간도 최대 6.5% 감소했다. 이 백엔드는 Blackwell 아키텍처의 최신 기능인 TMA와 Thread Block Clusters를 효과적으로 활용하여 차세대 AI 모델 추론 최적화의 핵심 도구로 자리 잡을 전망이다.

배경

PyTorch 2.11 이상 (Nightly 권장), NVIDIA Blackwell(B200) 하드웨어 및 CUDA 13.1 이상, GPU 커널 최적화 및 컴파일러 백엔드에 대한 기본 이해

대상 독자

NVIDIA Blackwell GPU 기반으로 고성능 LLM 추론 시스템을 구축하고 최적화하려는 ML 엔지니어 및 컴파일러 개발자

의미 / 영향

이 기술은 파이썬 기반 DSL로도 C++ 수준의 하드웨어 최적화가 가능함을 보여주며, 복잡한 GPU 커널 개발 및 유지보수 장벽을 낮춥니다. 특히 Blackwell과 같은 최신 하드웨어의 기능을 즉각적으로 프레임워크에 통합할 수 있어, 차세대 AI 모델의 배포 효율성을 획기적으로 높일 것으로 기대됩니다.

섹션별 상세

기존 CUTLASS C++ 백엔드는 커널 변형마다 nvcc 컴파일이 필요하여 오토튜닝 과정에서 병목 현상이 발생했다. CuteDSL은 파이썬-MLIR 컴파일러를 사용하여 다른 백엔드와 대등한 수준으로 컴파일 속도를 개선하고 하드웨어 계층 구조를 직접 제어할 수 있게 한다. 이를 통해 런타임에 수백 개의 커널 후보를 신속하게 평가하고 최적의 설정을 선택할 수 있는 환경이 구축됐다.

다양한 컴파일러 백엔드 간의 단일 커널 컴파일 시간 비교 표 — ChartCuteDSL의 컴파일 시간은 0.29초로, nvcc를 사용하는 CUTLASS 4.x(14.41초) 대비 약 50배 빠름을 보여준다. 이는 런타임 오토튜닝 시 CuteDSL이 훨씬 효율적임을 입증하는 핵심 근거이다.

근거

CuteDSL은 nvcc를 사용하는 CUTLASS 4.x GEMM보다 컴파일 속도가 50배 빠르다. — Compile time comparison 차트 (이미지 1)

Triton은 메모리 대역폭 제한 작업에 강점이 있지만, 연산 집약적인 GEMM에서는 하드웨어 특화 기능 활용이 제한적이다. CuteDSL 백엔드는 Blackwell(B200)의 텐서 코어 파이프라인, 공유 메모리 스테이징, 분산 공유 메모리 등을 정교하게 제어하는 데 집중한다. 특히 LLM 추론의 디코드 단계와 같이 특정 형상의 행렬 연산에서 하드웨어 활용도를 극대화하는 전략을 취한다.

TorchInductor는 cutlass_api를 통해 호환 가능한 모든 커널 구성을 조회하고 nvMatmulHeuristics 분석 모델로 후보군을 압축한다. 수백 개의 후보 중 하드웨어 처리량 예측 점수가 높은 상위 5개 내외의 커널만 실제 벤치마킹 대상으로 선정하여 오토튜닝 시간을 최적화한다. 최종 선택된 커널은 메모리에 캐싱되어 이후 동일한 연산 요청 시 컴파일 오버헤드 없이 즉시 실행된다.

NVIDIA B200 GPU 기반 테스트에서 BF16 디코드 형상 연산 속도가 기존 대비 최대 1.73배 향상되는 결과가 확인됐다. MXFP8 및 NVFP4와 같은 최신 데이터 형식에서도 중간 크기 행렬 연산에서 최대 1.78배의 성능 이득을 얻었다. 이는 Blackwell 하드웨어의 잠재력을 소프트웨어 수준에서 효과적으로 끌어내고 있음을 증명하는 수치이다.

B200 GPU에서 BF16 행렬 곱셈 커널의 처리량(TFLOPS) 비교 차트 — Chart다양한 행렬 크기(M, N, K)에 대해 Inductor ATen, Triton, NVGEMM(CuteDSL)의 성능을 비교한다. 특히 작은 M 값(디코드 단계)에서 NVGEMM이 기존 백엔드 대비 최대 1.73배 높은 성능을 보임을 시각화한다.

근거

BF16 데이터 형식의 디코드 형상(M=8~64)에서 최대 1.73배의 성능 향상을 기록했다. — Kernel-Level Speedups 섹션 및 BF16 차트 (이미지 5)

vLLM V1 엔진을 활용한 엔드투엔드 추론 테스트에서 Llama 3.3 70B 모델의 지연 시간이 배치 사이즈 16 기준 6.5% 감소했다. Llama 3.1 8B 모델 또한 배치 사이즈 전반에서 2~4%의 일관된 성능 향상을 보였다. 동적 형상을 사용하는 추론 환경에서도 autotune_batch_hint를 통해 런타임 형상에 최적화된 커널을 선택함으로써 실질적인 서비스 성능 개선이 가능하다.

python

import torch
import torch._inductor.config as config

# NVGEMM 백엔드 활성화
config.max_autotune_gemm_backends = "ATEN,TRITON,NVGEMM"

A = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16)
B = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)

@torch.compile(mode="max-autotune-no-cudagraphs")
def f(a, b):
    return a @ b

out = f(A, B)  # 첫 호출 시 오토튜닝 트리거

TorchInductor에서 CuteDSL(NVGEMM) 백엔드를 활성화하고 행렬 곱셈을 실행하는 예시 코드

vLLM 엔진을 사용한 BF16 데이터 형식의 모델별 추론 속도 향상률 차트 — ChartLlama 3.1, Qwen3, Llama 3.3 모델에 대해 배치 사이즈별 성능 향상을 보여준다. Llama 3.3 70B 모델이 배치 사이즈 16에서 6.5%로 가장 높은 성능 향상을 기록했음을 확인할 수 있다.

근거

vLLM 추론 테스트에서 Llama 3.3 70B 모델은 배치 사이즈 16일 때 6.5%의 속도 향상을 보였다. — End-to-End vLLM Inference 섹션 및 BF16 Inference 차트 (이미지 8)

용어 해설

GEMM: — General Matrix Multiply의 약자로, 딥러닝 모델 연산의 대부분을 차지하는 핵심 연산이다. 하드웨어 가속기의 성능을 최대한 끌어내기 위해 타일링, 워프 스케줄링 등 정교한 최적화가 필수적이다.
Autotuning: — 특정 하드웨어와 데이터 형상에 최적화된 커널 설정을 런타임에 자동으로 탐색하고 선택하는 기법이다. 타일 크기나 워프 구성 등 다양한 후보군을 벤치마킹하여 가장 빠른 결과물을 도출한다.
Epilogue Fusion: — 행렬 곱셈 연산 직후에 이어지는 활성화 함수나 덧셈 연산을 별도의 커널 호출 없이 하나의 커널 안에서 처리하는 최적화 기법이다. 메모리 대역폭 낭비를 줄여 전체 추론 속도를 향상시킨다.
MXFP8: — NVIDIA Blackwell 아키텍처에서 도입된 8비트 부동소수점 데이터 형식이다. 기존 FP8보다 정밀도와 효율성을 개선하여 대규모 언어 모델의 추론 및 학습 속도를 획기적으로 높인다.
JIT Compilation: — 프로그램 실행 시점에 코드를 기계어로 컴파일하는 방식이다. TorchInductor는 JIT 방식을 통해 실제 입력 데이터의 크기와 타입을 확인하고 그에 최적화된 GPU 커널을 생성한다.

코드 예제

bash

pip install nvidia-cutlass-dsl==4.3.5
pip install nvidia-matmul-heuristics
# Clone and install cutlass_api from the cutlass_api branch
git clone --branch cutlass_api https://github.com/NVIDIA/cutlass.git
cd cutlass/python/cutlass_api
pip install -e ".[torch]"

CuteDSL 백엔드 사용을 위한 필수 라이브러리 및 API 설치 명령어

언급된 리소스

GitHubNVIDIA CUTLASS GitHub (cutlass_api branch)

문서nvidia-cutlass-dsl PyPI

import torch import torch._inductor.config as config # NVGEMM 백엔드 활성화 config.max_autotune_gemm_backends = "ATEN,TRITON,NVGEMM" A = torch.randn(128, 4096, device="cuda", dtype=torch.bfloat16) B = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16) @torch.compile(mode="max-autotune-no-cudagraphs") def f(a, b): return a @ b out = f(A, B) # 첫 호출 시 오토튜닝 트리거

pip install nvidia-cutlass-dsl==4.3.5 pip install nvidia-matmul-heuristics # Clone and install cutlass_api from the cutlass_api branch git clone --branch cutlass_api https://github.com/NVIDIA/cutlass.git cd cutlass/python/cutlass_api pip install -e ".[torch]"

TorchInductor를 위한 새로운 CuteDSL 백엔드: Blackwell GPU에서 GEMM 성능 극대화

TL;DR

배경

대상 독자

의미 / 영향

섹션별 상세

용어 해설

코드 예제

언급된 리소스

TorchInductor를 위한 새로운 CuteDSL 백엔드: Blackwell GPU에서 GEMM 성능 극대화

TL;DR

배경

대상 독자

의미 / 영향

섹션별 상세

용어 해설

코드 예제

언급된 리소스

관련 토론

댓글

관련 기사

PyTorch 2.13 출시: MPS용 FlexAttention, CuTeDSL 백엔드, 메모리 최적화 기능 추가

Triton 플러그인 확장으로 TLX와 커스텀 컴파일러 패스 기본 제공

관련 토론

댓글

관련 기사

PyTorch 2.13 출시: MPS용 FlexAttention, CuTeDSL 백엔드, 메모리 최적화 기능 추가

Triton 플러그인 확장으로 TLX와 커스텀 컴파일러 패스 기본 제공