AMD: ROCm 저지연 GEMM 커널로 Instinct MI355X에서 LLM 추론 최대 1.79배 가속

GEMM과 LLM 디코딩 단계란?

GEMM(일반 행렬 곱셈)은 모든 대형 언어 모델을 통과하는 계산을 지배하는 기본 연산입니다. 프리필 단계에서 모델은 입력 프롬프트를 병렬로 처리하지만, 디코딩 단계 — 출력 토큰을 하나씩 생성할 때 — 배치 크기 M은 일반적으로 작습니다: 1, 2, 4, 8 또는 16행. 이 비대칭성(작은 M, 하지만 수천 단위의 큰 K와 N)이 디코딩 단계를 핵심 병목으로 만듭니다. 높은 처리량에 최적화된 표준 GEMM 루틴은 이 경우 낮은 지연 성능을 제공합니다.

AITER 프레임워크 내 FlyDSL의 세 가지 기법

AMD는 ROCm 생태계의 일부로 AITER(AI Tensor Engine for ROCm)와 그 안의 FlyDSL — 특수 GEMM 커널을 자동으로 합성하는 생성기를 개발했습니다. FlyDSL은 세 가지 보완적 기법을 결합합니다:

1. Inter-CTA Split-K 병렬화 — K 차원을 따라 실행 그리드를 확장하여 여러 블록(CTA)에 작업을 분산하고 GPU 자원 유휴 상태를 제거합니다.
2. Intra-CTA K-슬라이스 분할 — 단일 CTA 내에서 K 축을 더 작은 조각으로 나누어 추가 동기화 비용 없이 유효 병렬성을 높입니다.
3. LDS 파이프라인(다단계) — 전역 메모리에서 로컬 공유 메모리 버퍼(LDS)로의 데이터 전송과 활성 계산을 중첩시켜 AMD Instinct MI355X 아키텍처(gfx950, 256 컴퓨트 유닛)의 메모리 지연을 숨깁니다.

결과 및 하드웨어: 평균 1.64배, 핵심 시나리오에서 1.79배

벤치마킹은 실제 프로덕션 모델의 32개 기본 형태와 48개 추가 변형에서 수행되었습니다 — DeepSeek V3, Llama 70B, Llama 450B, Qwen32B — FlyDSL 커널을 세 가지 기준 구현과 비교했습니다: HipblasLT, AITER Triton, AITER ASM. 핵심 형태(K=7168)에서 평균 지연 시간 감소는 1.64배이며, 디코딩에 중요한 M≤8 토큰 시나리오에서 가속은 1.79배에 달합니다. 특정 형태에서 최대 이득은 2.37배로 측정되었습니다. 프로덕션 모델의 더 넓은 BF16 형태 집합에서 평균은 1.49배입니다.

AMD는 프로그래매틱 접근으로 소프트웨어 격차를 좁힐 수 있을까?

FlyDSL과 AITER는 ROCm 생태계의 소프트웨어 격차에 대한 AMD의 체계적인 대응입니다. NVIDIA의 cuBLAS가 수년간의 우위를 갖고 있는 반면, AMD는 이제 고성능 커널을 프로그래매틱하게 생성합니다 — 즉, 어셈블리 코드를 수동으로 작성하지 않고도 최적화를 새로운 GPU 아키텍처에 신속하게 확장할 수 있음을 의미합니다. AMD Instinct 인프라로의 전환을 고려하는 운영자에게 이 디코딩 지연 시간 개선은 생성된 토큰당 비용에 직접적인 영향을 미칩니다.