PyTorch FlyDSL: 연산 지원, 성능 향상

1. 기존 연산 경로에 더한 선택지

지원 대상은 일반적인 밀집 행렬곱, 빈칸 없이 값을 계산하는 방식과 여러 묶음을 함께 처리하는 그룹형 행렬곱, 저정밀 스케일 행렬곱이다. 값을 정규화하는 RMSNorm, 큰 값 K개를 고르는 TopK, 입력 사이의 관계를 계산하는 다중 헤드 어텐션도 포함된다.

컴파일 방식에서는 TorchInductor가 조건에 맞는 행렬곱 구현들을 직접 측정해 가장 빠른 후보를 저장한다. 기존 구현과 새 구현을 작업별로 경주시켜 고르는 방식이다. 즉시 실행에서는 선택적 런타임이 설치되고 활성화되면 지원 조건에 맞는 정규화·선택 연산을 새 커널, 실행할 계산 절차를 담은 코드로 보낸다.

어텐션은 AOTriton 라이브러리를 거쳐 사전 조정 데이터에 따라 구현이 선택된다.

FlyDSL은 데이터 이동과 계산의 배치를 커널 작성자가 직접 제어하게 해 하드웨어별 최적화를 지원한다.

2. 연산별 측정 결과는 범위가 다르다

연산 수준 측정은 Instinct MI355X(gfx950)에서 컴파일과 최초 자동 튜닝을 제외하고 진행했다. 저정밀 측정은 이미 양자화된, 값을 적은 비트로 표현한 입력에서 시작했다. 아래 기하평균은 각 사례에 같은 가중치를 준 집계다.

회사 발표 기준 BF16 밀집 행렬곱은 왼쪽 행렬이 행 우선, 오른쪽이 열 우선인 NT 배치의 15개 형상에서 기존 두 구현 중 빠른 쪽 대비 기하평균 1.10배였다. 두 비교 대상은 ATen과 Triton이다.

회사 발표 기준 MXFP8과 MXFP4 스케일 행렬곱은 NT 배치의 각 17개 형상에서 기존 기본 구현 대비 각각 1.58배와 1.68배였다.

그룹형 행렬곱은 회사 발표 기준 24개 사례에서 Triton 대비 1.20배, ATen 대비 1.95배의 기하평균을 기록했다.

회사 발표 기준 RMSNorm은 22개 사례에서 ATen 대비 1.17~3.66배였다. 정렬된 차원에서는 대체로 1.2~1.5배였고, 최고치는 특정 FP16 입력에서 나왔다.

회사 발표 기준 작은 K의 TopK는 10개 사례에서 ATen 대비 기하평균 4.82배였다. 이는 결정적 알고리즘, 같은 입력에 같은 결과를 내는 절차를 끈 조건이며, 켠 조건에서는 4.02배였다.

3. 어텐션과 모델 추론에서 확인한 효과

모델 전체 측정은 MI355X 한 대에서 vLLM으로 진행했다. 입력 256토큰과 출력 512토큰을 고정하고 동시 요청을 8~256개로 바꿨다. 비교 양쪽의 서빙 설정과 모델별 캐시는 같고 행렬곱 후보 목록만 달랐다.

회사 발표 기준 BF16 추론은 기존 두 백엔드에 새 후보를 더하자 Llama-3.1-8B-Instruct가 최대 1.13배, Qwen3-32B가 최대 1.06배였다. Llama-3.3-70B-Instruct는 측정 잡음 범위에 대체로 머물렀다.

저정밀 8비트 추론은 회사 발표 기준 아래와 같았다. 기준은 hipBLASLt를 쓰는 ATen이며, 모델·동시 요청 수의 18개 조합 모두 개선됐다.

MI355X 한 대(TP=1), MXFP8 A8W8, 입력 256토큰·출력 512토큰, 동시 요청 8~256개, ATen(hipBLASLt) 대비 FlyDSL 추가, 조합별 1회 측정

동시 요청 수에 따라 모델별 이득의 폭이 달랐다. 회사 발표 기준 같은 요청 조건의 MXFP4 추론도 18개 조합에서 1.25~3.01배였다. 이 측정은 최종 출력층을 양자화하지 않고 조합별 한 번씩 실행했다.

어텐션은 AOTriton의 이전 버전 대비 처리량 집계다. 회사 발표 기준 gfx950에서 0.13b 대비 0.14.2b는 순방향 0.92~2.34배, 역방향 0.93~2.47배였다.

4. 사용 조건과 성능의 대가

최근 ROCm PyTorch nightly와 gfx950 기반 MI350 계열에서 시작한다. 컴파일 행렬곱용 선택적 런타임은 0.3.x 계열이며, 아래는 원문 예시의 후보 설정 부분이다.

import torch
import torch._inductor.config as config

config.max_autotune_gemm_backends = "ATEN,TRITON,FLYDSL"
config.flydsl_enable_autotuning = True

설정 뒤 행렬곱을 torch.compile(mode="max-autotune")로 컴파일하면 자동 튜닝 후보가 된다. 어텐션은 별도 패키지나 설정 없이 AOTriton 0.14.2b 이상과 연결된 빌드로 이용한다.

반복 실행 속도를 얻는 대신 첫 호출에는 컴파일과 튜닝 비용이 든다. 저정밀 입력이 요구 배치에 맞지 않으면 변환 복사도 추가된다. 지원 연산이 모델 실행 시간에서 차지하는 비중이 클수록 최적화의 가치가 커진다.

어텐션 버전 비교에는 다른 커널과 조정 데이터 변경도 포함되며, 현재 새 어텐션 커널은 FP32를 지원하지 않는다. 정규화 역방향은 기존 구현을 사용하고, 상위 K개 선택에서 같은 값의 인덱스 순서는 백엔드마다 달라질 수 있다.

5. 지원 확대는 진행 중

현재 통합은 gfx950의 MI350 계열을 대상으로 한다. 다음 단계로 MI450 지원, 저정밀·어텐션 커널 확장, 후처리 융합, 뒤따르는 계산을 하나로 묶는 처리와 역방향 지원 확대 등이 진행 중이라고 원문은 설명한다.

FlyDSL은 AMD GPU용 PyTorch에 최적화 선택지를 보탰으며, 이득은 특히 저정밀 행렬곱과 일부 선택 연산 및 추론 조건에서 두드러진다.