PyTorch, Helion에 컴파일 타임 휴리스틱 추가해 커널 성능 향상

PyTorch 팀이 Helion DSL에 컴파일 타임 휴리스틱을 도입해 빠른 기본 구성을 제공함. 기존 오토튜닝이 단일 paged-attention 커널에 10~25시간 소요되는 병목 문제를 해결함.
휴리스틱은 정적 분석과 데이터플로우를 통해 커널을 matmul, reduction, pointwise 카테고리로 분류하고 블록 크기, 워프, 스테이지를 결정함. 컴파일 타임 오버헤드는 무시할 수 있는 수준임.
H100 환경에서 reduction 커널 기준 vLLM의 수동 작성 CUDA 연산자 대비 1.27x, Flash Linear Attention(FLA) 대비 1.18x의 전방 속도 향상을 달성함.
matmul 성능은 torch.compile(max-autotune) 대비 1.88x를 기록했으나, jagged HSTU attention에서는 휴리스틱(3.27x)이 이전 기본값(3.72x)보다 낮은 성능을 보임.

Implementation Details

라이브니스 분석을 통해 레지스터, 공유 메모리, 텐서 메모리 사용량을 추정하여 점유율과 레지스터 압력을 계산함.
matmul 휴리스틱은 CTA 수에 따라 블록 크기를 조정하고 공유 메모리 한계 내에서 num_stages를 최대화함. num_warps는 Triton의 MMA 명령어 선택에 의해 결정됨.
reduction 휴리스틱은 리덕션 폭에 기반해 워프 수를 최적화하고, 여러 읽기가 필요한 RMSNorm 같은 커널에는 persistent reduction을 적용함.

현재 NVIDIA H100과 B200에서 Triton 백엔드를 사용할 때만 활성화됨.
다른 컴파일러의 메모리 레이아웃 결정을 예측하는 데 의존하므로 오류가 발생할 수 있으며, 특히 matmul 결과가 텐서 메모리에 저장되는 B200에서 두드러짐.
휴리스틱 결과를 오토튜너 시드로 사용하는 것은 초기 탐색 속도를 35~41x 높이지만 최종 구성 품질에는 영향을 미치지 않음.