Command Palette
Search for a command to run...
텐센트 팔콘GEMM은 스트라센(Strassen)과 LCMA의 저복잡도 행렬 곱셈 원리를 활용하여 하드웨어 한계를 뛰어넘는 행렬 곱셈 최적화를 연구합니다.
8월 1일, HyperAI가 주최하는 제9회 AI 컴파일러 기술 살롱이 베이징에서 개최되었습니다. 이번 행사는 AI 컴파일 기술의 최신 발전 동향을 집중 조명했으며, 업계 및 연구 기관의 전문가들이 프로그래밍 언어, 연산자 개발, 컴파일 최적화, 추론 실행 등에 대한 통찰력을 공유하고, 고수준 언어 표현에서 하드웨어 실행에 이르기까지 AI 컴파일러의 협력적 진화를 보여주었습니다.
안에,텐센트의 고성능 컴퓨팅 엔지니어인 주훙린은 "FalconGEMM: 저복잡도 행렬 곱셈으로 하드웨어 한계를 뛰어넘다"라는 제목의 발표에서 저복잡도 행렬 곱셈을 위한 팀의 알고리즘 및 연산자 최적화 방식을 공유했습니다.
cuBLAS와 같은 성숙한 연산자 라이브러리가 행렬 곱셈 성능을 하드웨어 최고 수준에 가깝게 끌어올렸고, 기존 커널 수준 최적화 공간이 점점 좁아지고 있다는 문제에 직면하여, 연구팀은 알고리즘 복잡도부터 다시 접근했습니다. Strassen 및 AlphaTensor와 같은 저복잡도 행렬 곱셈 알고리즘을 기반으로 통합 LCMA(저복잡도 행렬 알고리즘) 프레임워크를 구축했습니다. QDSL, 연산자 융합, 지속 커널, 세밀한 스케줄링 및 비용 모델을 결합하여 "곱셈 횟수 감소"라는 이론적 이점을 GPU에서 실제 성능 향상으로 전환했습니다.
NVIDIA H20의 FP16 및 BF16 테스트에서 FalconGEMM은 다양한 행렬 형태에서 cuBLAS보다 우수한 성능을 보였으며, 최대 약 10%~16%의 성능 향상을 달성했습니다. 또한, 언어 모델 벤치마크에서 표준 행렬 곱셈과 기본적으로 동일한 수치 정확도를 유지했습니다.
위챗 공식 계정 "HyperAI"를 팔로우하고 백그라운드에서 키워드 ""를 입력하여 답글을 남겨주세요.0801 AI 컴파일러"..."을 클릭하시면 발표자의 공식 프레젠테이션 PPT를 다운로드하실 수 있습니다.
스트라센의 연구를 시작으로, 행렬 곱셈에 대한 최적화 공간을 재검토해 보겠습니다.
행렬 곱셈은 딥러닝에서 가장 중요한 기본 연산자 중 하나이며, 일반적으로 모델의 계산 시간 대부분을 차지합니다. CUDA, MKL, cuBLAS와 같은 소프트웨어 스택은 수년간 최적화되어 왔으며, 많은 시나리오에서 단일 GEMM 커널의 성능은 이미 하드웨어 최고 성능에 매우 근접해 있습니다. 이는 명령어, 파이프라인, 메모리 접근 수준에서만 부분적인 최적화를 계속 진행한다면 개선의 여지가 점점 줄어든다는 것을 의미합니다. 따라서 연구팀은 고전적인 스트라센 알고리즘으로 다시 주목했습니다.
스트라센 알고리즘은 1969년 볼커 스트라센(Volker Strassen)에 의해 제안되었습니다. 가장 기본적인 2×2 행렬 곱셈의 경우, 기존 방식은 8번의 곱셈 연산을 필요로 하는 반면, 스트라센 알고리즘은 입력 행렬을 재조합함으로써 7번의 곱셈 연산만으로 결과를 얻고, 추가적인 덧셈과 뺄셈 연산을 통해 최종 결과를 복원합니다. 이는 곱셈 연산량을 1/8로 줄이는 것과 같습니다. 피연산자가 스칼라일 경우에는 이러한 효율성 향상이 큰 이점이 되지 않지만, 피연산자가 부분 행렬이 될 경우 행렬 덧셈(O(N²))과 행렬 곱셈(O(N³)) 사이의 시간 복잡도 차이가 상당해집니다.이로써 "행렬 곱셈을 한 번 덜 하고 행렬 덧셈을 여러 번 더 하는 것"이 실질적인 가치를 지니게 됩니다.
스트라센 알고리즘을 재귀적으로 사용하면 곱셈 횟수를 더욱 줄일 수 있습니다. 예를 들어, 4×4 블록 행렬 곱셈은 기존 방식에서는 64번의 블록 곱셈이 필요하지만, 스트라센 알고리즘을 두 단계로 적용하면 49번의 블록 곱셈만으로 충분합니다. 그러나 재귀 단계가 많아질수록 덧셈, 데이터 구성, 메모리 접근 오버헤드도 증가합니다.따라서 실제 시스템에서는 계산량 감소와 추가 오버헤드 사이의 균형을 맞추기 위해 제한된 수의 레이어만 사용하는 경우가 많습니다.
2022년, DeepMind의 AlphaTensor는 이러한 알고리즘 영역을 더욱 확장했습니다. AlphaTensor는 행렬 곱셈을 텐서 분해 문제로 변환하고 강화 학습을 사용하여 더 낮은 랭크의 분해 방법을 탐색함으로써, 기존의 Strassen 방법 외에도 다양한 M, N, K 형태에 대해 복잡도가 낮은 다양한 행렬 곱셈 알고리즘이 존재할 수 있음을 보여주었습니다.
하지만 알고리즘 발견에서 엔지니어링 응용으로 나아가는 데에는 여전히 실질적인 문제가 있습니다.만약 복잡도가 낮은 알고리즘 하나하나에 대해 별도의 GPU 커널을 직접 작성해야 한다면, 개발 및 유지 관리 비용이 지나치게 높아질 것입니다.이 문제를 해결하기 위해 연구팀은 이러한 알고리즘들을 LCMA(저복잡도 행렬 알고리즘)로 추상화했습니다. LCMA는 입력 행렬에서 어떤 하위 블록을 미리 결합해야 하는지, 실제로 몇 번의 행렬 곱셈이 수행되는지, 그리고 중간 결과들이 최종적으로 어떻게 출력 행렬로 결합되는지를 일관되게 설명합니다. 그런 다음 Codegen을 사용하여 해당 구현을 자동으로 생성합니다.
따라서 "스트라센 커널을 어떻게 구현할 것인가"라는 질문이 제기됩니다.초점은 "높은 성능을 유지하면서 다양한 저복잡도 행렬 알고리즘을 지원할 수 있는 통합 프레임워크를 구축하는 방법"으로 옮겨갔습니다.
동시에, 저복잡도 알고리즘은 수치 정밀도 문제도 해결해야 합니다. 스트라센 곱셈은 대수적으로 표준 행렬 곱셈과 동일하지만, 부동 소수점 연산은 결합 법칙을 엄격하게 준수하지 않으며 계산 순서 변경으로 인해 추가적인 반올림 오차가 발생할 수 있습니다. 따라서 LCMA는 성능을 추구하는 동시에 저정밀도 계산에서 오차 전파를 제어해야 합니다.
통일된 알고리즘 설명이 마련되자 다음 단계는 적합한 GPU 구현을 찾는 것이었습니다. 팀은 CUDA, Triton, TiLang, QDSL을 시도했습니다. CUDA는 가장 강력한 하드웨어 제어 기능을 제공했지만, 다양한 LCMA 알고리즘, 레지스터, 공유 메모리, 중간 합산 구조 등을 처리해야 하는 경우 각각 특정 조정이 필요하여 확장 및 유지 관리 비용이 많이 드는 것으로 나타났습니다.
Triton은 기본적인 Strassen 시나리오에서 CUDA에 근접한 성능을 보이지만, 알고리즘이 더 큰 블록 구조로 확장될 경우 여러 중간 계산에 걸쳐 레지스터 버퍼를 정확하게 재사용해야 하므로 Triton은 추가적인 메모리 누수(spill) 발생 가능성이 높습니다. TiLang은 레지스터 및 공유 메모리 제어에서 더 많은 유연성을 제공하지만, 팀 테스트에서 성능은 여전히 Triton보다 약 51~101 TP3T 낮았습니다. 이론상 12.51 TP3T의 성능 향상만 있는 기본적인 Strassen의 경우, 이러한 성능 손실은 알고리즘의 이점을 상쇄할 만큼 충분히 큽니다.
결정적인,해당 팀은 FalconGEMM의 주요 구현 백엔드로 QDSL을 선택했습니다. QDSL은 CUDA에 가까운 개발 세분성을 제공하는 동시에 코드 생성 기능과 내장 PTX를 지원합니다. 이를 통해 기존의 고성능 구현을 쉽게 마이그레이션할 수 있으며, 다양한 LCMA 설명에 기반한 일괄 코드 생성이 가능하여 후속 통합 및 맞춤형 최적화에 더 큰 유연성을 제공합니다.
LCMA에서 FalconGEMM까지, 알고리즘 성능 향상을 GPU로 이전하기
가장 간단한 Strassen GPU 구현은 다음과 같은 몇 단계로 나눌 수 있습니다. 부분 행렬 A와 B를 결합하여 7개의 새로운 입력 쌍을 생성하고, 7개의 GEMM을 일괄 처리 방식으로 실행한 다음, 마지막으로 7개의 중간 결과 세트를 결합하여 최종 행렬 C를 생성합니다. 일반적인 GEMM과 비교했을 때, 실제로 계산 집약적인 행렬 곱셈 부분은 원래의 7/8에 불과합니다.따라서 사전 및 사후 처리에 소요되는 추가 시간이 절감된 계산 부하의 1/8 미만인 한, 전반적으로 이점을 얻을 수 있는 가능성이 있습니다.
연구팀은 먼저 NVIDIA H20에서 테스트를 진행했습니다. H20은 높은 메모리 대역폭과 상대적으로 낮은 최대 연산 부하를 가지고 있어 "계산 부하를 줄이는 대신 데이터 처리량을 늘리는" 이러한 접근 방식에 적합합니다. 약 2048³ 이상의 행렬 크기에서는 기본 구현만으로도 안정적인 성능 향상을 보였습니다. 그러나 행렬 크기가 작아질수록 입력 조합, 중간 결과 쓰기, 출력 조합의 비율이 급격히 증가하여 절감된 계산 부하가 쉽게 소모됩니다.
따라서 후속 최적화의 초점은 GEMM 자체에서 중간 메모리 접근으로 옮겨갔습니다. 가장 직접적인 접근 방식은 연산자 융합입니다.중간 결과를 글로벌 메모리에 반복적으로 기록하는 대신, 가능한 한 칩에 저장해 두도록 하십시오.그러나 입력 A/B를 GEMM에 직접 결합하는 것은 동일한 서브블록이 여러 SM에서 사용될 수 있어 중복 로딩 및 합산이 발생하기 쉽기 때문에 적합하지 않습니다. 이와 대조적으로 배치 GEMM과 Combine H를 사후 처리하는 융합 방식이 더 실현 가능합니다.
진정한 어려움은 스트라센의 7개 중간 결과가 최종 4개 출력 부분 행렬에 각기 다른 방식으로 기여한다는 사실에 있습니다. H를 병렬 단위로 사용하는 경우, 여러 SM이 동시에 동일한 C에 값을 기록하여 심각한 원자적 충돌이 발생할 수 있습니다. 반대로 C를 병렬 단위로 사용하는 경우, 일부 H가 서로 다른 SM에 의해 반복적으로 계산될 수 있습니다. 두 접근 방식 모두 곱셈 연산을 줄이는 이점을 무효화합니다.
팀은 결국 스트라센의 중간 결과를 기반으로 작업을 구성하는 방식을 포기했습니다. 대신, 행렬의 공간 좌표에 따라 작업을 그룹화했습니다. 7개의 배치된 GEMM에서 동일한 위치에 있는 7개의 곱셈 타일을 하나의 그룹으로 묶어 동일한 SM에서 실행했습니다. 이러한 방식으로 한 그룹이 계산을 완료하면 결과를 온칩 최종 C 코드에 직접 누적할 수 있어 중간 결과를 글로벌 메모리에 다시 기록할 필요가 없어지고 SM 간의 쓰기 충돌을 방지할 수 있었습니다.
이 융합 접근 방식은 Strassen 알고리즘에서 발생하는 추가 메모리 접근 횟수를 크게 줄여주지만, 그룹 세분화 정도가 커지면 부하 불균형이 발생합니다. 예를 들어, 4096³ 행렬 곱셈의 경우, 세분화되지 않은 스케줄링 방식은 약 211 TP3T의 추가 파동 낭비를 초래할 수 있으며, 이는 Strassen 알고리즘 자체의 12.51 TP3T 계산량 감소 효과를 초과할 수도 있습니다.
이를 위해,팀은 필요에 따라 실행을 위해 그룹을 두 개의 SM으로 나누는 Stream-K의 아이디어에서 영감을 얻었습니다.스케줄링 계층은 여전히 그룹을 기본 단위로 사용하지만, 실제 실행은 타일 단위로 더욱 세분화할 수 있으므로, 테일 SM의 유휴 시간을 줄이고 하드웨어 활용률을 향상시키면서 그룹 수준의 데이터 재사용이라는 장점을 유지할 수 있습니다.
하지만 로드 밸런싱 문제를 해결한 후 새로운 문제가 발생했습니다. 바로 L2 캐시 스래싱이었습니다. 그룹 분할 후, 서로 다른 유형의 중간 곱셈 연산이 동일한 웨이브 내에서 혼합될 수 있었고, 접근하는 데이터가 독립적이어서 L2 적중률이 크게 떨어졌습니다. 동시에 GEMM은 이미 텐서 코어를 과도하게 사용하고 있었고, 메모리 접근 부하가 거의 최대치에 도달하면 H20은 전력 한계에 도달했습니다. 실제 테스트에서 코어 주파수는 약 1.8GHz에서 1.6GHz로 떨어져 연산 성능이 저하되었고, 융합의 이점 중 일부가 다시 한번 상쇄되었습니다.
L2 캐시 스래싱 문제를 해결하기 위해 팀은 분할 그룹의 순서를 추가로 조정하여 동일한 유형의 중간 결과가 동일한 웨이브 내에서 처리되도록 하고, 혼합 처리는 몇몇 후반 웨이브에서만 발생하도록 했습니다. 이를 통해 세밀한 스케줄링이 제공하는 로드 밸런싱을 유지하면서 L2 데이터 지역성을 개선하여 궁극적으로 빈도 감소 문제를 해결했습니다.
언급할 가치가 있는 것은 다음과 같습니다.이러한 스케줄링 최적화의 핵심 기반은 영구 커널입니다.일반적인 커널에서는 태스크 어드바이저(CTA)가 블록 처리를 완료한 후 종료되는 것과 달리, 퍼시스턴트 커널에서는 CTA가 스트리밍 서비스(SM)에 장기간 상주하며 후속 태스크를 지속적으로 수락할 수 있습니다. 이를 통해 개발자는 그룹 및 타일의 실행 순서를 더욱 유연하게 제어하고 온칩 리소스를 재사용할 수 있습니다. 따라서 태스크 분할, 스케줄링 재배치 및 캐싱 최적화 또한 동일한 커널 내에서 수행할 수 있습니다.
캐시 재정렬부터 비용 모델까지, 최대 성능 향상은 10%~16%입니다.
병합, 로드 밸런싱 및 캐시 재배열을 거치면 FalconGEMM은 더 많은 형태에서 LCMA의 계산 이점을 활용할 수 있습니다. 그러나 LCMA가 모든 경우에 일반 GEMM보다 우수한 것은 아닙니다.본질은 동일합니다. 추가적인 데이터 처리량을 감수하는 대신 곱셈 계산량을 줄이는 것입니다.기존 GEMM이 이미 메모리 접근에 제약을 받는 상황에서 계산량을 더 줄인다고 해서 충분한 이점을 얻을 수는 없습니다. 저복잡도 알고리즘은 계산 밀도가 높을 때만 더 유리합니다.
그러므로,팀은 LCMA를 언제 사용해야 하는지, 그리고 다양한 LCMA 중에서 어떤 것을 선택해야 하는지를 결정하기 위해 지붕선과 유사한 비용 모델을 추가로 설계했습니다.목표는 실행 시간을 정확하게 예측하는 것이 아니라 "적절한 알고리즘을 선택하는 것"이므로, 이 모델은 주로 다양한 방식의 계산 및 메모리 접근 비용을 분석하고, 대상 GPU의 컴퓨팅 성능과 대역폭을 결합하여 해당 GPU의 계산/메모리 접근 병목 현상 범위를 추정합니다.
이 모델에서 저복잡도 알고리즘의 곱셈 횟수 감소는 계산 효율성 향상으로 이어지지만, 추가적인 데이터 조합과 반복적인 메모리 접근은 새로운 메모리 오버헤드를 발생시킵니다. 앞서 언급한 융합 최적화는 이러한 메모리 접근 비용을 더욱 줄여줍니다. 따라서 FalconGEMM은 다양한 M, N, K 형태에 따라 기존 GEMM과 여러 LCMA 방식의 이점 범위를 판단하고, 보다 적합한 구현 방식을 자동으로 선택할 수 있습니다.
QDSL의 코드 생성 기능을 활용함으로써 전체 프레임워크는 궁극적으로 비교적 완전한 실행 흐름을 형성합니다. 먼저, LCMA 설명에 기반하여 해당 융합 영구 커널이 생성되어 융합을 통해 중간 메모리 접근 횟수를 줄입니다. 그런 다음, 비용 모델이 특정 형상에 적합한 행렬 곱셈 알고리즘을 선택합니다. 마지막으로, QDSL이 대상 코드를 자동으로 생성하고 컴파일합니다. 이러한 방식으로 LCMA는 더 이상 고정된 Strassen 구현에 그치지 않고, 작업 부하에 따라 동적으로 선택 가능한 알고리즘 공간을 형성합니다.
성능 테스트는 주로 NVIDIA H20에서 수행되었습니다. 결과는 다양한 저정밀 행렬 곱셈 시나리오에서 다음과 같은 결과를 보여줍니다.FalconGEMM은 다양한 형상에서 cuBLAS보다 우수한 성능을 보이며, 최대 성능 향상은 약 10%~16%에 달합니다.그룹 분할 및 캐시 재배열이 완료되면, 큰 형상은 저복잡도 알고리즘으로 인한 연산 이점을 안정적으로 얻을 수 있고, 작은 형상의 성능 또한 향상되며, L2 캐시 지터로 인해 발생하는 전력 병목 현상으로 인한 주파수 감소를 방지할 수 있습니다.
비용 모델 선택 결과에 따르면, 해당 모델은 LCMA 반환 임계값을 넘는 대부분의 형태에 대해 더 나은 성능을 보이는 구현을 선택할 수 있었으며, 이는 "알고리즘 선택 + 융합 커널" 접근 방식이 다양한 계산 밀도를 가진 행렬 곱셈 시나리오를 효과적으로 다룰 수 있음을 나타냅니다.
성능 외에도 수치 정확도는 FalconGEMM이 검증해야 할 중요한 문제입니다. 연구팀은 초기 저정밀도 실험에서 주로 부동 소수점 덧셈 순서 변경으로 인한 상당한 오류를 발견했습니다. 예를 들어, A+C+B−C는 대수적으로 A+B와 같지만, 유한 정밀도 부동 소수점 계산에서는 반드시 엄밀히 같지는 않습니다.
추가 분석 결과, 실제로 중요한 오류는 FP32 누적 자체보다는 저정밀도 변환에서 주로 발생하는 것으로 나타났습니다. 일반적인 FP16/BF16 입력 시나리오에서는 행렬 곱셈을 먼저 FP32 누적을 사용하여 수행한 다음 다시 저정밀도로 변환합니다. 이때 중간 결과가 자주 변환되면 FP32 가수 정보가 지속적으로 버려지게 됩니다.
융합 방식은 실제로 이러한 문제를 완화합니다. WGMMA의 출력은 FP32 정밀도를 유지하는 반면, FalconGEMM은 최종 C 값을 FP32를 사용하여 온칩에서 직접 결합하고 누적한 다음 계산이 완료된 후에야 목표 정밀도로 다시 변환합니다. 여러 개의 독립적인 커널 간에 저정밀도 중간 결과를 반복적으로 기록하는 방식과 비교했을 때, 이 방법은 하나 이상의 정밀도 변환을 줄여 계산 순서 변경으로 인한 오류가 FP32의 하위 비트에 더 많이 남도록 합니다.
언어 모델 벤치마크에서 FalconGEMM과 표준 행렬 곱셈을 사용하여 얻은 최종 점수는 매우 미미한 차이만 보일 뿐 거의 동일하며, 이는 현재 구현 방식이 모델 정확도를 크게 저하시키지 않음을 나타냅니다.
다음 단계에서 팀은 두 가지 방향으로 연구를 진행할 계획입니다. 첫째, 배치 그룹과 K 루프의 순서를 조정하여 중간 입력 결과에 대한 메모리 접근을 더욱 줄임으로써 Combine A/B를 더욱 통합하는 것입니다. 둘째, LCMA를 어텐션 알고리즘으로 확장하는 것입니다. 플래시 어텐션 역시 높은 연산 효율과 메모리 사용량을 가지고 있으며, 저복잡도 행렬 알고리즘을 블록 분해 및 파이프라인과 결합할 경우 새로운 성능 향상 가능성을 가져올 수 있습니다.
Strassen과 AlphaTensor부터 LCMA와 FalconGEMM에 이르기까지, 이 연구의 중요성은 이미 고도로 최적화된 GEMM을 단순히 몇 퍼센트 더 빠르게 만드는 것 이상입니다. 이는 새로운 접근 방식을 제시합니다. 커널 자체가 이미 하드웨어 한계에 근접했을 때, 성능 최적화는 명령어와 파이프라인을 더욱 심층적으로 분석할 뿐만 아니라 알고리즘 복잡성에서도 개선의 여지를 찾아낼 수 있습니다. 그런 다음 컴파일, 융합 및 스케줄링을 통해 이론적인 계산량 감소를 실제 실행 시간 단축으로 전환할 수 있습니다.
