기술 요약: CUDA-L2

문제 정의

행렬 곱셈(matmul)은 대규모 언어 모델(LLM)에서 매우 중요한 연산으로, 훈련과 추론 모두에서 계산 시간의 상당 부분을 차지합니다. 수동으로 최적화된 CUDA 커널이 존재하지만, 다양한 행렬 차원($M, N, K$)과 GPU 아키텍처 전반에 걸쳐 최적의 성능을 달고 구현하는 것은 여전히 큰 과제입니다. 서로 다른 차원 조합은 종-종 별도의 최적화 전략을 필요로 하며, 이러한 전략은 하드웨어 간에 잘 전이되지 않습니다. 또한, 동일한 차원에 대해서도 누산기 정밀도(FP16 vs FP32)와 같은 설정 세부 사항에 따라 최적의 전략이 달라집니다.

LLM을 사용하여 커널을 생성하는 것과 같은 기존의 자동화된 접근 방식은 주로 단일 고정 구성(예: KernelBench)을 대상으로 한 벤치마크에 집중해 왔습니다. 이러한 방법들이 NVIDIA의 cuBLAS 및 cuBLASLt와 같이 고도로 최적화된 폐쇄형 라이브러리와 경쟁해야 하는 프로덕션 환경에서 어떻게 적용될 수 있는지는 여전히 불분명합니다. 인간 전문가가 최적화한 베이스라인을 능가하기 위해 반정밀도 일반 행렬 곱셈(HGEMM)의 방대한 구성 공간을 체계적으로 탐색할 수 있는 자동화된 시스템이 부족한 실정입니다.

방법론: CUDA-L2

저자들은 HGEMM CUDA 커널을 자동으로 최적화하기 위해 대규모 언어 모델(LLM)과 강화 학습(RL)을 결합한 시스템인 CUDA-L2를 제안합니다. 이 시스템은 $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$의 모든 $M, N, K$ 값의 조합을 나타내는 1,000개의 뚜렷한 구성을 처리하도록 설계되었습니다.

시스템 아키텍처 및 훈련

CUDA-L2는 이전의 CUDA-L1 프레임워크를 네 가지 핵심 기술적 향상 사항으로 확장했습니다:

1. 지속적 사전 학습(Continued Pretraining): 모델은 표준 벤치마크를 넘어 PyTorch, ATen, CUTLASS 및 NVIDIA 튜토리얼의 소스를 포함한 다양한 CUDA 코드 코퍼스에서 사전 학습됩니다. 원시 코드 스니펫에 대해 지시문 설명이 생성되며, 새로운 지식과 아키텍처 특성을 통합하기 위해 검색 증강 컨텍스트(retrieval-augmented context)가 추가됩니다.
2. 다단계 RL 훈련: 훈련은 광범위한 일반 커널 유형에서 특정 도메인(구체적으로 matmul)으로 진행됩니다.
   • 일반 커널 RL: 실행 속도를 보상으로 사용하여 광범위한 연산(선형 대수, 컨볼루션, 리덕션 등)에 대해 모델을 훈련합니다.
   • HGEMM RL: 다양한 $M, N, K$ 구성을 가진 HGEMM 커널에 특화하여 훈련합니다.
3. 프로파일링 메트릭을 통한 컨텍스트 강화: RL 컨텍스트에는 단순한 엔드 투 엔드 실행 시간을 넘어 더 나은 최적화 결정을 내릴 수 있도록 NVIDIA Nsight Compute(NCU) 프로파일링 메트릭(예: 메모리 처리량, SM 점유율, 캐시 효율성)이 포함됩니다.
4. 보상 함수 설계: 생성된 커널에 대한 보상은 테스트 반복 간의 평균 속도 향상 점수를 기반으로 합니다. 견고성을 보장하기 위해 다음 항목에 대해 페널티를 부여합니다:
   • FP32 CPU 그라운드 트루스(ground truth)로부터의 수치적 편차.
   • 과도한 코드 길이(간결함을 장려하기 위함).
     시스템은 바이너리 입력에 대한 정확한 일치와 신뢰할 수 있는 베이스라인(cuBLAS, cuBLASLt)에 대한 제한된 편차를 포함한 엄격한 정확성 기준을 강제합니다.

평가 프로토ル

시스템은 서로 다른 배포 현실을 반영하기 위해 두 가지 시나리오에서 평가됩니다:

• 오프라인 모드(Offline Mode): 커널이 시간 간격 없이 연속적으로 실행되어 피크 처리량을 시뮬레이션합니다.
• 서버 모드(Server Mode): 커널이 무작위 간격으로 실행되어 GPU 열 역학(냉각 기간 및 클록 부스팅)을 고려한 실시간 추론을 시뮬레이션합니다.

평가는 CUDA-L2를 다음 네 가지 베이스라인과 비교합니다:

1. torch.matmul (PyTorch의 기본 디스패치).
2. cuBLAS (NN과 TN 사이의 최적 레이아웃 선택 사용).
3. cuBLASLt-heuristic (NVIDIA의 휴리스틱 API에서 상위 랭크된 알고리즘 선택).
4. cuBLASLt-AutoTuning (가장 빠른 것을 찾기 위해 최대 100개의 후보를 전수 조사).

주요 기여 및 발견된 기술

CUDA-L2는 LLM 유도 RL이 수동 튜닝이 불가능한 구성 공간을 체계적으로 탐색함으로써 우수한 구현을 발견할 수 있음을 입증합니다. 시스템은 다음과 같은 여러 고급 최적화 기술을 자율적으로 발견하고 적용합니다:

• 추상화 선택(Abstraction Selection): 모델은 작은 행렬을 위한 원시 WMMA 인트린직(intrinsics)과 더 크고 복잡한 타일링 연산을 위한 CuTe 추상화 사이에서 자동으로 선택합니다.
• 입력 패딩(Input Padding): 시스템은 나누어 떨어지지 않는 차원(예: $M=8192$를 $8320$으로 패딩)으로 입력 행렬을 패딩하여, 그렇지 않으면 유효하지 않을 효율적인 타일 크기($B_M$)를 활성화함으로써 추가 계산 오버헤드와 성능 이득 사이의 균형을 맞추는 전략을 발견합니다.
• 고급 메모리 및 연산 중첩(Overlap):
  • 더블 버퍼링 레지스터 프래그먼트(Double-Buffered Register Fragments): 데이터 이동과 연산을 중첩시키기 위해 두 개의 레지스터 프래그먼트 세트 사이를 교대로 전환(핑퐁 실행)합니다.
  • 공격적인 프리페칭(Aggressive Prefetching): 레지스터 압박이 허용될 때 단일 단계 프리페칭을 넘어 다단계 프리페칭(여러 반복을 앞서 로드)으로 확장합니다.
  • 스테거드 프리페칭 스케줄링(Staggered Prefetch Scheduling): 행렬 곱셈-누산(MMA) 연산 주변으로 행렬 A와 B의 프리페칭을 인터리빙하여 명령어 수준 병렬성을 극대화합니다.
  • 직접 레지스터-투-공유 메모리 복사(Direct Register-to-Shared-Memory Copy): 레지스터 레이아웃이 공유 메모리 레이아웃과 일치할 때 에필로그(epilogue)에서 중간 임시 텐서를 제거하기 위해 넓은 데이터 타입(예: uint128_t)을 사용하여 복사 작업을 줄입니다.
  • 블록 스위즐링(Block Swizzling): 특히 큰 문제 크기에 대해 L2 캐시 히트율을 개선하기 위해 스레드 블록 실행 순서를 재정렬합니다.

결과

CUDA-L2는 A100 GPU에서 1,000개의 구성에 대해 모든 주요 베이스라인을 체계적으로 능가합니다.

오프라인 모드 성능 (베이스라인 대비 평균 속도 향상):

• torch.matmul 대비 +22.0%
• cuBLAS (최적 레이아웃 사용) 대비 +19.2%
• cuBLASLt-heuristic 대비 +16.8%
• cuBLASLt-AutoTuning (가장 강력한 베이스라인, 100개 후보 전수 조사) 대비 +11.4%

서버 모드 성능:
시스템이 열 역학 및 캐시 냉각 효과를 더 잘 처리할 수 있기 때문에 서버 시나리오에서 속도 향상이 증가합니다:

• torch.matmul 대비 +28.7%
• cuBLAS 대비 +26.0%
• cuBLASLt-heuristic 대비 +22.4%
• cuBLASLt-AutoTuning 대비 +15.9%

승률은 모든 베이스라인에 대해 79.3%에서 95.7% 사이로, 이는 개선이 이상치가 아닌 체계적으로 발생했음을 나타냅니다. 또한, GPU 활용도가 낮은 작은 행렬 크기에서 속도 향상이 일반적으로 더 높으며, 큰 행렬은 GPU 처리량을 포화시켜 베이스라인과의 격차를 좁히는 경 양상을 보입니다.

의의

본 논문은 CUDA-L2가 HGEMM과 같이 성능이 매우 중요하고 고도로 최적화된 커널조차 LLM 유도 RL 자동화를 통해 개선될 수 있음을 보여준다고 주장합니다. 인간이 수동으로 튜닝하기 어려운 구성 공간 전체로 최적화를 확장함으로써, 이 시스템은 수동 튜닝이나 휴리스틱 기반 라이브러리가 놓칠 수 있는 최적화를 찾아내어 최고 수준의 폐쇄형 라이브러리를 능가하는 성능 이득을 달 achieves 합니다.

현재 버전은 A100 (Ampere) 아키텍처로 제한되어 있지만, 프레임워크는 광범위한 적용 가능성을 갖도록 설계되었습니다. 저자들은 Ada Lovelace, Hopper, Blackwell을 포함한 다른 GPU 아키텍처로 시스템을 확장하기 위한 지속적인 연구를 언급했습니다. 이 연구는 자동화된 방법이 수동 튜닝이나 휴리스틱 기반 라이브러리가 놓칠 수 있는 최적화를 찾기 위해 방대한 구성 공간을 체계적으로 탐색할 수 있음을 시사합니다.