tile-ai/tilelang-ascend

해결 과제

TileLang-Ascend는 Huawei Ascend NPU를 위한 고성능 AI 컴퓨팅 커널을 생성할 때 겪는 어려움을 해결합니다. 개발자는 Python 스타일의 구문을 사용하여 최적화된 커널(GEMM, 어텐션 메커니즘, 벡터 연산 등)을 작성할 수 있으며, 복잡한 저수준 NPU 최적화를 수동으로 처리할 필요 없이 최상의 성능을 유지할 수 있습니다.

작동 원리

이 시스템은 TVM 컴파일러 인프라 위에 구축된 도메인 특화 언어(DSL)입니다. 메모리 계층(Global, L1, L0 버퍼)을 관리하고 특수 컴파일러 백엔드(Ascend C & PTO 및 AscendNPU IR)를 제공함으로써, 고수준의 Python 코드를 Ascend NPU 아키텍처에 매핑합니다. 이 시스템은 자동 동기화 및 스코프 분리를 위한 '개발자 모드'와 실행 스코프 및 동기화 플래그를 수동으로 제어하는 '전문가 모드'를 모두 지원합니다.

대상 사용자

Huawei Ascend A2 및 A3 프로세서에 맞춰 컴퓨팅 워크로드를 최적화해야 하는 AI 엔지니어 및 커널 개발자.

주요 특징

• 광범위한 연산자 지원: GEMM, Flash Attention, Sparse Flash Attention, 컨볼루션, Softmax 구현을 포함합니다.
• 메모리 최적화: 온칩 메모리 사용량을 줄이기 위한 자동 버퍼 재사용 및 자동 작업 공간 할당 기능을 제공합니다.
• 성능 도구: 소프트웨어 파이프라이닝 (T.Pipelined), 자동 벡터화 (T.Parallel), L2 캐시 스위즐링을 지원합니다.
• PyTorch 통합: PyTorch와 NPU 간의 원활한 워크플로우를 위해 torch_tl_ascend를 제공합니다.