Command Palette
Search for a command to run...
베이징 인공지능 연구원(BAAI)은 Triton-TLE 계층형 언어 확장 기능을 제안하고 FlagTree 컴파일 최적화를 통해 자동 튜닝 속도를 100배 향상시켰습니다.
8월 1일, 베이징에서 제9회 AI 컴파일러 기술 살롱이 개최되었습니다. 이번 행사는 AI 컴파일 기술의 최신 발전 동향을 집중 조명했으며, 업계 및 연구 기관의 전문가들이 프로그래밍 언어, 연산자 개발, 컴파일 최적화, 추론 실행 등에 대한 통찰력을 공유하고, 고수준 언어 표현에서 하드웨어 실행에 이르기까지 AI 컴파일러의 협력적 진화를 보여주었습니다.
안에,BAAI AI 컴파일러의 연구원인 궈후이와 샤오항은 "FlagTree: Triton-TLE 언어 확장, 타일 IR 백엔드 및 컴파일러 최적화 사례"라는 주제로 발표를 진행하며, FlagTree 팀이 Triton을 중심으로 수행한 일련의 연구 결과를 소개했습니다.
이 작업의 핵심 목표는 Triton의 사용 편의성과 커뮤니티 생태계를 유지하면서 개발자가 필요에 따라 하드웨어 세부 정보를 더 깊이 탐구할 수 있도록 하고, 비교적 통합된 개발 시스템을 통해 신속한 연산자 최적화, 아키텍처 인식 튜닝 및 네이티브 코드 수준 최적화를 지원하는 것입니다.
위챗 공식 계정 "HyperAI"를 팔로우하고 백그라운드에서 키워드 ""를 입력하여 답글을 남겨주세요.0801 AI 컴파일러"..."을 클릭하시면 발표자의 공식 프레젠테이션 PPT를 다운로드하실 수 있습니다.
Triton을 기반으로 추상성과 성능의 균형을 재조정합니다.
최근 몇 년 동안 AI 컴파일러는 점점 더 두드러지는 모순에 직면해 왔습니다.하드웨어 아키텍처와 모델 연산자는 점점 더 복잡해지고 있지만, 개발자들은 여전히 고급 DSL을 사용하여 전문가가 작성한 커널에 가까운 성능을 달성하기를 희망합니다.
추상화 수준이 너무 높으면 컴파일러가 심층적인 최적화를 완료하는 데 필요한 충분한 정보를 얻지 못할 수 있고, 추상화 수준이 너무 낮으면 개발자는 CUDA나 벤더 고유 언어와 같은 복잡한 개발 방식으로 돌아가게 될 것입니다.
Triton의 성공은 GPU 연산자 개발을 스레드 수준에서 타일 수준으로 끌어올린 데 있습니다.사용자는 파이썬 DSL을 사용하여 데이터 블록 간의 계산 관계를 설명하는 반면, 스레드 매핑, 레지스터 할당, 데이터 레이아웃, 파이프라이닝 및 동기화와 같은 작업은 주로 컴파일러가 처리합니다.
이러한 접근 방식은 고성능 운영자를 위한 개발 장벽을 낮추고 대규모 커뮤니티 생태계를 조성하는 데 기여해 왔습니다. 그러나 차세대 GPU, 도메인 특화 아키텍처, 그리고 국내 생산 AI 칩의 지속적인 발전으로 인해 Triton의 기존 추상화 방식은 한계에 부딪히기 시작했습니다.
한편으로는,컴파일러 백엔드가 새로운 하드웨어의 저장 구조, 통신 메커니즘 또는 연산 장치를 아직 지원하지 않는 경우, 프런트엔드 개발자가 이러한 기능을 자체적으로 활용하기 어려울 것입니다.반면에,점점 더 많은 중요 운영자들이 저장 수준, 병렬 처리 세분성, CTA 협업, 통신과 연산의 중첩에 대한 정밀한 제어를 요구하고 있으며, 기존 Triton 코드는 이러한 의도를 완전히 표현하는 데 어려움을 겪는 경우가 있습니다.
Gluon, TLX, TileLang과 같은 새로운 언어 및 DSL의 등장은 동일한 추세를 반영합니다. 즉, AI 연산자 개발은 더 이상 단일 커널을 작성하는 것에 그치지 않고 데이터 레이아웃, 병렬 계층 구조, 파이프라인, 통신 토폴로지 및 하드웨어 특성을 표현하는 것까지 포함하게 되었습니다.
TLE는 Triton을 대체하려는 것이 아니라, Triton의 구문과 생태계를 계층적으로 확장하려는 것입니다.이 시스템은 TLE-Lite, TLE-Struct, TLE-Raw의 세 가지 계층으로 구성되어 있으며, 각각 경량 의미 힌트, 아키텍처 인식 제어, 네이티브 코드 수준 최적화에 해당합니다.
TLE-Lite는 알고리즘 엔지니어와 신속한 최적화 시나리오를 위해 설계되었습니다.개발자는 기본 하드웨어에 직접적으로 신경 쓸 필요가 없습니다. 대신 컴파일러에 보다 명확한 구조적 정보를 추가하여 보완할 수 있습니다.예를 들어, 텐서에 서브타일 단위로 접근해야 하거나, 분산 메쉬에서 연산을 실행해야 하거나, CTA(Consumer-Producer Announcement)가 생산자-소비자 파이프라인을 사용해야 할 수도 있습니다. 서브타일 연산을 예로 들면, 개발자는 더 큰 텐서에서 논리적 하위 블록을 직접 추출하고, 활성화, 정규화 또는 통계 연산을 수행한 다음, 오프셋을 수동으로 계산하거나 마스크를 구성하고 경계를 처리할 필요 없이 다시 저장할 수 있습니다.
컴파일러는 이것이 일반적인 타일 접근임을 인식할 수 있으므로 데이터 레이아웃, 벡터화, 뱅크 충돌 및 레지스터 재사용 측면에서 추가적인 최적화를 수행할 수 있습니다. 이 접근 방식은 특히 희소 어텐션, 로컬 정규화, 블록 통계 및 라우팅 연산자에 적합합니다. 분산 환경에서 TLE는 디바이스 메시를 사용하여 노드, GPU, 블록 클러스터 및 블록과 같은 다양한 레벨을 설명하고 이를 통합된 다차원 토폴로지로 구성합니다.
개발자는 메시 기반 접근 방식을 사용하여 프로그램을 작성하고, 컴파일러와 런타임은 논리적 토폴로지를 실제 하드웨어 및 통신 메커니즘에 매핑합니다. 결과적으로 링 통신, 배리어 동기화 및 샤딩된 액세스는 더 이상 코드에 흩어져 있는 랭크와 통신 그룹이 아니라 분석 가능하고 구조화된 의미론이 됩니다. 통신 관계가 명시적으로 표현되면 컴파일러는 토폴로지를 인식하는 스케줄링, 통신-계산 중첩, 배리어 병합 및 교착 상태 검사를 수행할 수 있게 됩니다.
TLE-Lite는 또한 파이프라인 기본 요소를 통해 CTA의 내부 협업을 생산자-소비자 모델로 추상화합니다. 개발자는 주로 누가 데이터를 생산하고 누가 소비하는지 정의하며, 기본 장벽, 버퍼 재사용 및 동기화 메커니즘은 컴파일러가 처리합니다.이는 기본 기능을 숨기는 것을 의미하는 것이 아니라, 분석, 검증 및 최적화가 가능한 프로그램 구조로 변환하는 것을 의미합니다.
경량화된 의미 체계부터 네이티브 패스스루까지, 다양한 수준의 최적화 깊이를 다룹니다.
TLE-Lite가 주로 플랫폼 간 의미 표현을 다룬다면,반면 TLE-Struct는 아키텍처에 대한 인식과 세밀한 조정에 더 중점을 둡니다.
다양한 GPU, DSA 및 AI 가속기는 스토리지 계층, 실행 장치, 동기화 메커니즘 및 온칩 네트워크에서 상당한 차이를 보입니다. 따라서 TLE-Struct는 개발자에게 계층적 병렬 및 스토리지 구조를 제공하여 개발자가 특정 공급업체의 독점 인터페이스에 직접 종속되지 않고 데이터 레이아웃, 연산 매핑 및 메모리 계층 구조를 명시적으로 정의할 수 있도록 합니다.
예를 들어, 동일한 로컬 버퍼를 GPU 백엔드의 공유 메모리와 DSA 백엔드의 스크래치패드 또는 온칩 SRAM에 매핑할 수 있습니다. 사용자는 구조화된 메모리 의도를 표현하고, 컴파일러는 이를 대상 하드웨어에 적합한 주소 공간 및 메모리 접근 명령어로 변환하는 역할을 합니다.
MoE에서 전문가 카운팅은 일반적인 시나리오입니다. 이 연산자는 기본적으로 서로 다른 전문가에게 라우팅된 토큰 수를 계산하며, 공유 메모리 레이아웃, 동시 업데이트, 뱅크 충돌 및 블록 간 집계와 같은 요소의 영향을 쉽게 받습니다.
TLE-Struct를 사용하면 개발자는 카운터 레이아웃을 명시적으로 구성하여 서로 다른 Expert 또는 Token을 로컬 저장소의 서로 다른 영역에 매핑할 수 있습니다. 컴파일러는 이러한 구조적 정보를 얻은 후 적절한 액세스 및 동기화 메서드를 생성합니다.
TLE-Raw는 성능 최적화 전문가를 위해 설계되었으며, 벤더 네이티브 코드의 인터페이스를 유지합니다.
일부 성능 향상 경로는 CUDA, 어셈블리 또는 전용 내장 함수를 직접 사용해야 합니다. 이러한 함수를 상위 수준 DSL로 재패키징하도록 강제하면 성능이 저하되거나 마이그레이션 비용이 증가할 수 있습니다. TLE-Raw를 사용하면 개발자는 Triton/TLE 아키텍처 내에 네이티브 코드를 인라인하거나 공급업체 컴파일 파이프라인에 직접 액세스할 수 있습니다.
All-Gather GEMM을 예로 들면,이 연산자는 통신, 행렬 계산, 버퍼링 및 동기화 관리를 포함합니다. 개발자는 기본 통신 기능을 재사용하면서 구조화된 텐서와 타일을 사용하여 계산을 표현할 수 있으며, 최종적으로 통합 컴파일 파이프라인이 다양한 부분을 호출 가능한 연산자로 구성합니다.
이러한 계층형 설계 덕분에 알고리즘 엔지니어, 연산자 개발자 및 성능 전문가가 가장 낮은 수준의 프로그래밍부터 시작하지 않고도 동일한 시스템 내에서 서로 다른 최적화 깊이를 선택할 수 있습니다.
성능 테스트 결과 또한 다음과 같습니다.TLE의 전반적인 추상화 오버헤드는 제어 가능합니다.
Radix Select 테스트에서 TLE는 TensorRT-LLM 알고리즘을 재현하여 다양한 형태에서 약 85%~97%의 성능을 달성했습니다. 플랫폼 간 유지 관리 및 빠른 반복 개발이 필요한 팀에게 있어, 더욱 통합된 코드로 전문가 수준에 가까운 성능을 달성하는 것은 상당한 엔지니어링 가치를 지닙니다.
128K 컨텍스트를 사용하는 SparseMLA 시나리오에서,TLE는 파이프라인 기본 요소를 사용하여 서로 다른 실행 역할 간의 협업을 표현하며, FlashMLA의 기준 성능인 약 90% 수준의 성능을 달성합니다.
또한, 연구팀은 8개의 NVIDIA H100 프로세서가 탑재된 단일 노드에서 All-Gather를 테스트하고 GEMM과 All-Gather의 통합 가능성을 추가적으로 탐구했습니다. 목표는 단순히 통신 라이브러리를 교체하는 것이 아니라, 통신 기능을 연산자 수준 표현식과 컴파일 최적화에 통합하여 데이터 전송, 로컬 연산, 동기화 및 후속 데이터 소비를 파이프라인 형태로 구현하는 것입니다.
추론 시스템의 경우, 이러한 종류의 통신-계산 융합은 개별 GEMM의 최고 성능을 개별적으로 개선하는 것보다 훨씬 더 의미가 있는 경우가 많습니다. 왜냐하면 사용자가 궁극적으로 체감하는 것은 종단 간 지연 시간이기 때문입니다.
플래그트리 컴파일 최적화 사례
FlagTree 팀은 TLE 언어 확장 외에도 두 가지 추가 작업을 수행했습니다.첫째, CUDA Tile IR 백엔드와 통합했습니다. 둘째, 실제 모델 워크로드에 맞춰 Triton 컴파일 최적화를 여러 가지 구현했습니다.
**CUDA Tile IR의 핵심 개념은 프로그램이 기본 스레드 매핑을 미리 결정하는 대신 타일을 표현할 수 있도록 하는 것입니다.**Dialect의 입력은 타일 프로그램과 3차원 타일 그리드로 구성됩니다. Dialect 내부에서는 타일 컴퓨트, 뷰 및 토큰 순서 연산(TKO)을 통해 계산, 데이터 보기 및 필요한 종속성이 표현됩니다.
이 중에서 TensorView는 전역 포인터, Shape 및 Stride를 설명하고, PartitionView는 그 위에 블록 매핑 기능을 추가하며, Token은 관련 TKO 간의 종속성을 제한하는 데 사용되고, Memory Model은 메모리 의미론과 범위를 별도로 지정합니다.
FlagTree는 NVIDIA의 네이티브 백엔드인 Triton을 완전히 대체하는 것이 아닙니다. 대신 기존 시스템에 TileIR 경로를 추가하고 TLE 기본 요소를 통해 TileIR 뷰 및 토큰 인터페이스를 재정의합니다. 요구 사항을 충족하는 커널은 TileIR 백엔드에 액세스할 수 있으며, 아직 지원되지 않는 커널은 네이티브 CUDA로 대체됩니다. 두 경로는 컴파일러에서 공존할 수 있습니다.
자동 최적화와 관련하여,해당 팀은 Triton Autotune의 실제 모델에서 적용 범위와 비용 간의 충돌을 해결하기 위해 FlagOSTune을 제안했습니다.
실제 모델 추론에서 동일한 연산자는 매우 다양한 형태에 대응될 수 있습니다. 연구팀의 통계에 따르면 6개의 모델과 4가지 유형의 추론 시나리오에서 1994개의 고유한 MM 형태가 존재하며, 이는 일일 벤치마크에서 다루는 범위를 훨씬 뛰어넘는 수치입니다. 후보 구성의 범위를 직접 확장하면 이론적인 탐색 규모는 수백만 개의 집합으로 급격히 증가할 것입니다.
FlagOSTune은 탐색 공간을 확장하여 성능을 향상시키고, 모델 예측과 제한적인 실제 테스트를 결합하여 탐색 비용을 절감합니다. 이 시스템은 먼저 XGBoost를 사용하여 후보 구성의 순위를 매긴 다음, 가장 유망한 소수의 구성만 실제 컴파일 및 GPU 테스트에 사용하고, 그 후에 유전 알고리즘을 사용하여 탐색을 계속합니다.
NVIDIA, Moore Threads, Muxi 등 다양한 컴퓨팅 시스템에서 여러 연산자의 성능이 향상되어 1.21배에서 7.35배까지 속도가 향상되었습니다.NVIDIA H20 MM 연산자를 사용한 다양한 형상 실험에서 검색 공간 구성이 62만 개 이상에서 4,070개로 압축되어 성능 손실 없이 튜닝 효율성이 120배 향상되었습니다.
데이터 레이아웃 측면에서, 팀은 데이터 레이아웃 변환의 성능 오버헤드를 최적화하는 데 집중했습니다.
Triton에서 `convert_layout` 함수는 일반적으로 스레드 간 데이터 재배열을 수반하며, 공유 메모리에 쓰기, 동기화, 그리고 다시 읽어오는 작업을 필요로 합니다. 이는 비용이 수반되는 작업입니다. 작은 연산자나 메모리 사용량이 많은 연산자의 경우, 레이아웃 변환 횟수가 적더라도 성능 병목 현상이 발생할 수 있습니다.
FlagTree는 비용 모델링, 역전파 및 로컬 조인트 해결과 같은 기술을 사용하여 불필요한 데이터 레이아웃 변환을 줄임으로써 레이아웃 변환 제거 메커니즘을 향상시킵니다.100명 이상의 작업자를 대상으로 한 테스트에서 순 전환 제거율은 약 68%~79%였으며, 성능 향상은 최대 71%에 달했습니다.
또 다른 최적화 방법은 명령어 재정렬입니다. 루프 언롤링은 단순히 루프 본문을 복사할 뿐이며 여러 개의 로드 작업을 미리 실행하지 않습니다. 컴파일러 재정렬을 활성화하면 독립적인 로드 작업을 미리 실행하여 후속 계산을 인터리빙할 수 있으므로 하드웨어 파이프라인을 활용하여 메모리 접근 지연 시간을 줄일 수 있습니다.세 가지 대표적인 연산자에서 이러한 최적화는 평균적으로 약 1.19배에서 1.61배의 속도 향상을 가져오며, 최대 2배의 속도 향상을 보입니다.
또한, 연구팀은 DeepSeek-V4-Flash에서 실제 형상에 대한 Fused Marlin MoE 연산자의 성능을 최적화하여 조각 병합을 통해 데이터 로딩 시간을 단축했습니다. 알고리즘 수준의 최적화와 결합하여,vLLM CUDA와 비교했을 때, FlagGems는 테스트된 53개 모양 모두에서 속도 향상을 달성했으며, 모양 빈도를 가중치로 적용했을 때 평균 약 1.208배의 개선을 보였습니다.
NVIDIA H20의 DeepSeek-V4-Flash 엔드투엔드 테스트에서,TP=4일 때 첫 번째 토큰의 지연 시간은 20.221 TP3T만큼 감소하고 전체 처리량은 12.631 TP3T만큼 증가합니다. TP=8일 때 첫 번째 토큰의 지연 시간은 12.871 TP3T만큼 감소하고 전체 처리량은 7.651 TP3T만큼 증가합니다.
미래의 직업 전망
다음 단계는,FlagTree는 NUMA 프로그래밍 모델과 MegaKernel 컴파일러를 발전시키는 데 집중할 것입니다.
멀티 노드, 멀티 GPU 구성, GPU 내부 클러스터 및 로컬 SRAM의 경우 데이터 접근은 근접성 측면에서 상당한 차이를 보입니다. TLE는 텐서가 메시를 따라 어떻게 분할되고 서로 다른 분산 방식 간에 어떻게 전환되는지를 명시적으로 기술하여 컴파일러가 어떤 데이터를 계산에 근접시켜야 하는지, 어떤 통신을 계산과 병합하거나 중첩할 수 있는지를 판단할 수 있도록 합니다.
메가커널 컴파일러는 최적화 범위를 단일 커널에서 모델 실행 프로세스로 확장하려고 시도합니다. 이 컴파일러는 공동 최적화가 가능한 계산을 식별하고, 이를 구체적인 작업으로 분해한 다음, 이러한 작업에 대한 할당, 병렬화 및 파이프라인 방식을 일관되게 결정하여 최종적으로 하나 또는 몇 개의 메가커널을 생성합니다.
TLE는 프런트엔드 시맨틱스와 하드웨어 아키텍처 사이의 가교 역할을 합니다. 타일, 메시, 파이프라인, 메모리 레이아웃 및 네이티브 기능 인터페이스를 통해 컴파일러는 일련의 독립적인 연산자뿐만 아니라 스케줄링, 배치, 동기화 및 병합이 가능한 태스크 그래프를 인식할 수 있습니다.
FlagTree는 TLE-Lite의 경량 시맨틱 힌트부터 TLE-Struct의 아키텍처 인식 제어, TLE-Raw의 네이티브 기능 통과에 이르기까지 다양한 컴파일 최적화 기법을 결합하여 Triton의 사용 편의성, 크로스 플랫폼 마이그레이션 및 궁극적인 성능 간에 보다 유연한 계층 시스템을 구축하고, 새로운 하드웨어 기능을 모델 및 연산자의 실제 성능으로 더 빠르게 전환할 수 있도록 노력하고 있습니다.
