AMD MI300X 커스텀 커널 생성

조금 더 많은 정보: ChatGPT 는 매일 수 억 건의 요청을 처리하며, 이는 곧 줄어들 가능성이 낮은 수치입니다. 각 요청 및 생성된 토큰마다 수십억 개의 파라미터를 가진 모델의 추론 (inference) 을 실행합니다. 이것이 바로 모델 최적화가 모든 수준에서 필수적인 이유입니다: 이러한 규모를 다룰 때, 1% 의 지연 시간 또는 전력 효율 향상도 막대한 절감 효과를 가져옵니다.

그런데 그 효율 향상의 출처는 어디일까요? 모델 아키텍처는 이미 확립되어 있고, 인기 있는 모델들은 오랫동안 양자화된 가중치 (quantized weight) 를 사용하고 있습니다. 그러나 여전히 최적화가 가능한 중요한 수준은 커널 (kernel) 수준입니다. 커널은 네트워크에서 어떤 연산을 수행할 때 실행되는 알고리즘입니다: 행렬 곱셈 커널, 컨볼루션 커널, 배치 정규화 커널 등이 있습니다. 커널은 저수준 (low-level) 으로 고도로 최적화된 알고리즘이며, 주로 해당 장치에 맞게 조정됩니다. 이들을 작성하는 것은 notoriously 길고 어렵습니다. GPU 의 내부를 잘 이해해야 합니다.

커널은 신경망 연산을 실행하는 데 필수적입니다—커널이 없으면 연산은 실제로 사용할 수 없습니다. 따라서 새로운 혁신들은 종종 "Day 0" 커널로 시작하며, 이는 일반적으로 최신 Nvidia 하드웨어에만 최적화됩니다. 이 접근법은 많은 다른 장치를 제외합니다. 특히 AMD GPU 는 비교 가능하거나 우수한 사양을 제공하지만, 커널 개발자들은 종종 이를 간과합니다. Hugging Face 는 AMD 와 협력하여 AMD 플랫폼에서 최첨단 성능을 제공하고 오픈 소스 커뮤니티에 혜택을 주기로 결정했습니다. 이 파트너십의 일환으로 우리는 VLLM 을 사용하여 8 개 MI300X 노드에서 FP8 에서 Llama 3.1 405B 를 서비스하는 데 커널 최적화를 제공하는 데 집중하기로 AMD 와 함께 결정했습니다.

이 블로그 포스트에서는 우리가 MI300X 를 위한 성능을 어떻게 최적화했는지, 각 커널이 개별적으로 미세 조정 (fine-tuned) 되었는지를 탐구합니다. 하지만 먼저, 우리 커널을 사용하여 달성한 성능 향상을 살펴보겠습니다. 다음 세 가지 최적화된 커널을 결합하여:

• Fused residual connection, RMS norm and FP8 conversion kernel
• Fused SwiGLU activation and FP8 conversion kernel
• Skinny GEMM kernel

VLLM 을 MI300X GPU 로 구동할 때 상당한 속도 향상 (speedups) 을 달성했습니다.

입력 크기는 1, 출력 크기는 128 으로 설정하여 디코딩 모드 (decoding regime) 를 모방했습니다. 우리는 30 번의 반복에 대한 중앙값 (median) 을 사용하여 디코딩 지연 시간 (latency) 을 측정합니다.

이러한 성능 향상은 VLLM 에서 측정되었지만, 다음 "How to" 섹션에서 설명된 대로 커널을 별도로 사용하실 수도 있습니다.

前述 모든 커널은 hf-rocm-kernels 저장소에 여기 제공됩니다. 이 저장소에서 설치 방법을 배우고 각 커널의 소스 코드, 각각의 Python 바인딩 (bindings), 다양한 벤치마킹 스크립트 및 테스트 스위트 (test suite) 를 찾을 수 있습니다. 벤치마킹 스크립트와 MI300X 를 사용하여 이 블로그 포스트에서 결과를 재현할 수도 있습니다. Torch 또는 VLLM 에서 동일한 결과를 보장하려면 우리가 사용한 것과 동일한 컨테이너를 사용할 수 있습니다.

저장소를 커널을 구축하는 기초로 사용하실 수도 있습니다: CUDA 스타일의 커널을 Python 에 바인딩하고 간단한 샘플 커널에 대한 지침이 포함되어 있습니다.
새로운 커널 개발 중인 브랜치도 확인해 볼 수 있습니다. 예를 들어, 여기에서 설명된 compute-and-communicate 커널과 같은 것이 있습니다.

이러한 커널은 곧 AMD 포크의 VLLM 프로젝트에 통합될 예정이지만, 만약 여러분이 이를 직접 구현하는 방법을 알아보고 싶다면 이 브랜치와 이 문서를 확인해 보세요.

우선 우리가 작업 중인 장치인 MI300X 의 구조를 빠르게 복습하겠습니다. 그 다음에는 모델 추론의 현재 상태를 검토한 후 최적화 작업을 진행할 것입니다. 이를 통해 병목 현상을 식별하고 작성해야 할 커널을 파악할 수 있습니다. 이후 우리가 작성한 각 커널을 살펴보며, 다양한 관점에서 커널 최적화가 어떻게 수행되는지 탐구할 기회를 가질 것입니다.

GPU 코드를 최적화하기 전에 GPU 의 작동 원리를 알아야 합니다. 이미 GPU 의 내부 구조를 잘 설명하고 있는 리소스가 많이 있는데, 여기 바로 링크를 드리겠습니다. 우리는 여전히 GPU 의 다른 레벨을 빠르게 복습하겠습니다. 만약 여러분이 복습을 건너뛰고 커널의 세부 사항으로 바로 넘어가기를 원한다면, 여기를 클릭하세요!

GPU 에서 가장 작은 작업 단위는 스레드 (thread) 입니다. GPU 에서는 스레드가 명령어를 실행했을 때 작업이 수행됩니다. 명령어는 덧셈, 곱셈, 데이터 타입 변환, 또는 로드와 스토어 같은 기본 연산입니다. 각 스레드는 자신의 전용 메모리인 레지스터 (또는 VGPR) 를 가지며, 이는 오직 해당 스레드만 접근할 수 있습니다. 스레드는 최대 256 개의 32 비트宽的 레지스터를 가질 수 있습니다. 아래는 256 개의 VGPR 에 접근하는 스레드를 나타냅니다.

스레드는 로드 또는 스토어 명령어를 사용하는 경우를 제외하고는 오직 자신의 레지스터에 있는 명령어를 실행할 수 있습니다. 예를 들어, 두 벡터 A 와 B 를 더하기 위해서는 각 스레드가 1) A 의 요소 하나를 레지스터에 로드하고, 2) B 의 다른 요소 하나를 로드한 후, 3) 덧셈을 수행하여 결과를 또 다른 레지스터에 저장하고, 마지막으로 4) 해당 레지스터의 값을 메모리에 스토어해야 합니다. 총 4 개의 명령어입니다.

다음 작업 단위는 워프 (warp) 입니다. 각 워프는 64 개의 스레드로 구성됩니다. 워프는 자체 메모리를 가지지 않지만, 모든 스레드가 동시에 동일한 명령어를 실행해야 한다는 점은 우리에게 중요합니다. 이는 보장이자 제약입니다.

워프는 또한 같은 워프 내의 다른 스레드들이 자신의 레지스터에서 온 정보를 서로 교환할 수 있게 합니다. 워프 내의 다른 스레드는 서로 다른 데이터에 접근할 수 있지만, 모든 명령어를 동시에 실행해야 한다는 사실은 커널을 작성할 때 워프 레벨의 동작을 고려해야 함을 의미합니다.

워프는 스레드 블록 (thread blocks) 으로 묶여 있습니다. 스레드 블록은 소프트웨어 추상화이지만, 하드웨어 구성 요소인 컴퓨트 유닛 (CU) 에서 실행됩니다. 하나의 컴퓨트 유닛은 여러 스레드 블록을 동시에 실행할 수 있지만, 16 개의 워프만 수용할 수 있습니다.

각 컴퓨트 유닛에는 전용 L1 캐시와 공유 메모리가 있습니다. L1 캐시는 제어 또는 할당할 수 없으며, CU 에 위치한 모든 워프의 데이터 재사용에 도움을 줍니다. 반면, 공유 메모리는 할당할 수 있으며 모든 워프로 공유되는 저장소로 사용될 수 있습니다. 예를 들어, 컴퓨트 유닛 내의 모든 워프 (그리고 따라서 스레드) 가 동일한 버퍼에 접근하고 싶을 때, 우리는 이를 공유 메모리에 할당합니다. 공유 메모리와 L1 캐시는 스레드에 '근접'하여 접근하기 때문에 빠릅니다.

Thread blocks 또한 실행 중인 모든 스레드를 동기화할 수 있는 기능을 제공합니다: 이는 공유 메모리에 영향을 미치는 작업 (예: 공유 메모리에서 배열을 0 으로 초기화하거나 reduction operations 수행) 을 처리할 때 매우 유용합니다. 일반적으로 커널을 작성할 때는 thread blocks 을 고려해야 할 가장 높은 수준입니다: 다른 thread blocks 을 동기화하거나 어떤 방식으로든 상호작용을 만드는 것은 매우 어렵습니다. Kernel throughput 는 GPU 에 존재하는 compute unit 의 수와 밀접하게 연결되어 있습니다: CUs 가 많을수록 동시에 실행될 수 있는 thread block 이 더 많아지므로, 모든 CU 를 활용한다면 throughput 가 증가합니다.

Compute units 는 accelerator complex dies (XCDs) 로 그룹화됩니다. 각 XCD 는 38 개의 compute unit 을 포함합니다. CU 들은 서로 상호작용할 수는 없으나, 모두 제어할 수 없는 L2 cache 를 공유하며, 데이터를 재사용할 때 유용할 수 있습니다. 예를 들어, 메모리를 접근할 때 같은 XCD 에 위치해 있는 두 개의 compute unit 이 동일한 데이터를 접근하면 데이터 로딩 latency 가 크게 감소합니다. L2 cache 는 매우 큽니다: 크기는 4MB 인 반면 shared memory 는 64kB, L1 cache 는 32kB 입니다.

8 개 XCD 를 조립하여 (이는 8 * 38 = 304 개의 CUs 를 의미) 마지막 레벨의 cache (infinity cache 라고 하며 256MB) 와 방대한 양의 video ram (192GB) 를 추가하면 MI300X 가 됩니다.

모든 XCD, 즉 모든 스레드는 VRAM 에 접근할 수 있으나, 이를 도달하는 것은 매우 느립니다. 스레드 레벨에서 멀어질수록 메모리 접근 속도는 더 느려지지만 크기와 범위는 더 커지며, 이는 더 많은 스레드를 서비스합니다. Kernel 을 최적화할 때는 많은 작업을 수행하거나 많은 데이터를 로드하는 것 사이에는 항상 균형을 맞추어야 합니다. 그러나 일반적으로 VRAM (일반적으로 global memory 라고 함) 을 최대한 적게 접근하고 싶습니다.

이 그림을 볼 때, GPUs 가 "massively parallel"이라고 불리는 이유를 알 수 있습니다: 여기서는 304 개의 compute unit 이 있으며, 각각 16 개의 warps 를 실행할 수 있고 각 warp 는 64 개의 스레드를 포함합니다. 이는 동시에 최대 311296 개의 스레드를 실행할 수 있음을 의미하며, 각 스레드는 고유한 명령어를 실행합니다. 명령어는 addition 과 같은 기본적인 것을 의미하므로, Newton's method 와 같은 간단한 루틴은 단일 스레드에서 실행하는 데 매우 길게 걸릴 수 있습니다. GPUs 는 명령어가 빠르게 실행되도록 최적화되어 있지 않습니다: 즉, 각 명령어의 latency 를 낮추도록 최적화되어 있지 않습니다: 이는 latency-oriented device 입니다. 많은 스레드를 함께 실행하고, 대량의 데이터를 소비하고 출력하도록 최적화되어 있습니다: 이는 throughput-oriented device 입니다. GPU 를 위한 Kernel 을 최적화할 때 우리는 다음과 같이 적응합니다: 많은 스레드에서 동시에 몇 개의 명령어를 실행하는 것보다, 몇 개의 스레드에서 많은 명령어를 실행하는 것이 더 좋습니다. 따라서 GPU 에서 실행되는 알고리즘을 "parallel"이라고 부릅니다.

이러한 알고리즘이 최적화된 방식으로 실행되지 않는 것을 방해할 수 있는 것은 세 가지입니다: 데이터를 로드해야 할 때 (memory bound), 많은 작업을 수행해야 할 때 (compute bound) 또는 스레드들이 함께 일해야 할 때 (synchronization overhead).

Workload 를 최적화할 때, 코드를 작성하기 전에 먼저 현재 workload 의 상태를 profile 하는 것이 첫 번째 단계입니다. 우리 경우 VLLM 의 model inference 를 profile 하여 각 작업이 얼마나 많은 시간을 차지하는지 파악합니다. 이는 주요 병목 현상을 식별하고 최대 속도 향상을 위해 우선적으로 해결해야 할 커널을 찾을 수 있습니다. 예를 들어, 여기서는 batch size 32 의 분해가 있습니다:

각 슬라이스를 통해 네트워크의 다른 부분을 확인할 수 있습니다:

• "Attention*" 슬라이스: RoPE, attention, KV cache 커널을 그룹화한 부분;
• "Attention GEMMs": QKV 와 Output 두 개의 프로젝션을 포함하는 부분;
• "Communications": Attention 블록과 MLP 블록 후에 각각 수행되는 모든 reduce 연산으로 구성되며, 이는 텐서 병렬 (TP8) 작업 중임을 의미합니다;
• "MLP GEMMs": MLP 에서 수행된 Gate / Up 과 Down 두 개의 프로젝션을 포함하는 부분;
• "RMS norm" 와 "SwiGLU" 슬라이스: 각 커널 하나에 해당하며, RMS norm 커널은 블록당 두 번 호출됩니다 (Attention 전과 MLP 전 각각);
• "Other" 슬라이스: 더 큰 범주에 속하지 않았으며 영향력이 미미한 커널들을 그룹화한 부분.

이미 대부분의 지연 시간이 GEMMs 와 communications 에서 비롯된다는 것을 확인할 수 있으며, attention 과 그 주변 연산이 지연 시간의 주요 기여도가 아님을 알 수 있습니다. 이는 많은 논문들이 attention 과 그 비용을 줄이는 데 집중하고 있다는 사실에 놀라울 수 있습니다. 그러나 VLLM 에 이미 최적화된 KV caching 과 FlashAttention 의 조합을 통해 이 부분이 더 이상 최우선 순위가 아닐 수 있음을 알 수 있습니다. 흥미롭게도 "RMS norm" 커널 호출 두 번은 상당한 비용이 발생하므로, 이 커널을 최적화하는 데 큰 이점이 있을 수 있습니다. SwiGLU 커널과 함께 총 지연 시간의 15% 를 차지하며 이는 무시할 수 없습니다. 결국 이 두 커널에 작업하고 GEMMs 에 작은 속도 향상을 시도하는 것이 우리의 최선의 방안일 것입니다. 이 성능 분석이 우연적이지 않음을 확인하기 위해 다른 배치 크기를 살펴볼 수 있습니다:

배치 크기 32 에서 나타나는 패턴은 다른 배치 크기에 대해서도 유지되지만, 배치 크기가 증가함에 따라 GEMMs 와 communications 의 지연 시간 기여도가 더 커집니다. 또한 배치 크기 32 는 GEMMs 지연 시간에 있어 이상치 (outlier) 로 보입니다. 아마도 배치 크기가 32 일 때 선택된 GEMMs 가 수동으로 튜닝되었거나, 배치 크기 32 가 좋은 메모리 정렬 패턴을 제시하기 때문에 배치 크기 32 의 GEMMs 가 배치 크기 24 나 28 보다 더 빠를 것입니다.

이제 최적화할 핫 스팟을 식별했으니, 우리가 작성한 첫 번째 커널인 RMS norm 커널을 살펴보겠습니다.

각 디코더 블록에는 두 가지 주요 부분이 있습니다: attention 블록과 MLP 블록. 둘 다 현재 숨겨진 상태와 잔차 (residual) 사이의 잔차 연결로 시작합니다. 둘 다 행 수 (토큰 수만큼) 와 열 수를 가진 동일한 모양을 가집니다. 합쳐진 후, 우리는 행별 Root Mean Square (RMS) 노름을 적용하며, 모델이 FP8 이므로 scale 을 사용하여 FP8 으로 양자화합니다. 이 세 가지 연산을 단일 커널로 융합하면 좋은 성능 향상을 얻을 수 있습니다. 수학적으로 수행해야 할 연산은 다음과 같습니다: