Gaming GPU에서 hipBLASLt를 넘은 fp16 GEMM

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

이 글은 24 GB VRAM을 제공하는 RX 7900 XTX에서 Mojo로 작성한 fp16 WMMA GEMM 커널을 hipBLASLt와 같은 C++ shim으로 비교한 결과를 담고 있습니다. 256³부터 4096³까지의 정사각형 입력 10개에서 커널은 모든 크기에서 hipBLASLt보다 높은 처리량을 기록했으며, 1536³에서는 97,974 GFLOP/s 대 69,332 GFLOP/s로 1.35배 앞섰습니다. 성능을 뒤집은 핵심은 128x128 타일 하나를 고정하지 않고 grid 크기에 따라 128x128, 64x128, 64x64 타일과 warp 배치를 선택한 방식이었습니다. 글은 workspace 캐싱, heuristic 후보 실측, split-K와 wgm 설정, 10초 clock warm-up을 포함한 공정한 benchmark protocol이 과장된 비교를 막는다고 설명하며, 결과의 범위는 단일 RX 7900 XTX와 fp16 정사각형 GEMM으로 제한합니다.

실용적 조언

GPU GEMM kernel을 비교할 때는 입력 크기 하나만 고정하지 말고 grid 규모에 맞춰 tile shape와 warp layout을 선택해야 합니다. 이 글에서는 128x128, 64x128, 64x64 타일을 compile-time parameter로 만들고 M·N에서 계산한 블록 수로 launcher가 선택하게 했습니다.
Vendor library가 쉽게 이기는 것처럼 보이더라도 workspace 캐싱 여부, heuristic 후보의 실제 timing, split-K와 workgroup mapping 설정을 먼저 확인해야 합니다. 이 프로젝트에서는 해당 harness 결함을 고친 뒤 hipBLASLt의 512³ 처리량이 2497 GFLOP/s에서 5201 GFLOP/s로 상승했습니다.
GPU benchmark에는 측정 전 clock warm-up 시간을 고정해 기록해야 합니다. 동일한 fp16 kernel이 1초 warm-up에서 66,000 GFLOP/s, 10초 warm-up에서 90,705 GFLOP/s를 냈으므로 이 프로젝트는 모든 측정 전에 10초를 기다리고 warmup_s를 결과에 남깁니다.

섹션별 상세

RX 7900 XTX는 24 GB VRAM을 데이터센터 가격보다 낮은 비용으로 제공하지만, AMD의 hipBLASLt 튜닝 범위는 주로 Instinct 계열의 CDNA에 집중되어 gfx1100 기반 RDNA3에서는 상대적으로 제한적이라는 문제에서 출발합니다. mojo-baro는 이 차이를 확인하기 위해 Mojo로 작성한 GPU kernel을 같은 C++ shim을 거친 hipBLASLt와 비교했습니다. 비교 대상은 RX 7900 XTX 한 장에서 실행한 정사각형 fp16 GEMM이며, 결과의 적용 범위도 이 카드와 해당 shape family로 한정됩니다.

WMMA 기반 커널은 4x2 warp가 128x128 출력 블록을 계산하고, 두 개의 LDS 버퍼와 K-step마다 하나의 barrier를 사용해 현재 타일을 소비하는 동안 다음 A·B 타일을 적재합니다. A에는 XOR swizzle을 적용해 bank conflict를 피하고 B는 transpose 저장했으며, tile k+2를 레지스터에 미리 가져오는 2단계 global prefetch를 사용했습니다. 커널은 188 VGPR을 register spill 없이 사용하고 32 KB LDS에 들어가 RDNA3의 64 KB LDS-per-WGP 예산에서 WGP마다 두 블록을 상주시키며, 256³부터 4096³까지 모든 측정 크기에서 hipBLASLt보다 높은 처리량을 냈습니다.

첫 번째 커널 버전은 128x128 타일을 고정해 1024³ 이하에서 hipBLASLt의 0.78~0.98배에 그쳤고, 작은 grid가 충분한 작업 블록을 만들지 못해 일부 compute unit이 유휴 상태가 되는 문제가 있었습니다. 수정 버전은 M과 N에서 계산한 128x128 등가 블록 수에 따라 96개 이상이면 128x128과 4x2 warp, 64개 이상이면 64x128과 2x4 warp, 그보다 작으면 64x64와 2x2 warp를 compile-time parameter로 선택했습니다. 이 변경 뒤 1536³에서 97,974 GFLOP/s를 기록해 hipBLASLt의 69,332 GFLOP/s보다 1.35배 높아졌지만, 글은 이를 RDNA3 전체나 다른 GPU에 대한 일반적인 우위로 확대하지 않습니다.

이 프로젝트는 이전에 fp32 GEMM이 hipBLASLt보다 약 2배 빠르다는 수치를 냈지만, vendor 호출마다 workspace를 새로 할당하고 heuristic 후보를 직접 측정하지 않았으며 split-K와 wgm도 설정하지 않았다는 문제가 있었습니다. 세 결함을 고친 뒤 512³에서 hipBLASLt 처리량이 2497 GFLOP/s에서 5201 GFLOP/s로 올라가 기존의 2배 우위가 사라졌고, 최종 결론은 크기별로 두 커널이 엇갈리며 전체적으로 비슷하다는 더 작은 주장으로 수정됐습니다. fp16 benchmark에는 캐시된 workspace, 모든 heuristic 후보의 timing, 명시적인 splitK·wgm 설정을 처음부터 적용해 같은 비교 오류를 피했습니다.

측정값은 실행 전에 질문, 측정 도구, 예측 범위, 예측을 반증할 조건을 등록한 뒤 결과가 예측과 달라도 그대로 기록하는 방식으로 생성됐습니다. bench/run.py는 처리량을 출력하기 전에 kernel 결과를 correctness reference와 비교하므로, 빠르지만 틀린 결과에는 GFLOP/s 수치를 부여하지 않습니다. 또한 동일한 fp16 binary라도 1초 warm-up에서는 4096³에서 66,000 GFLOP/s, 10초 warm-up에서는 90,705 GFLOP/s를 기록했기 때문에 모든 benchmark가 10초 동안 clock을 안정화하고 warmup_s를 함께 기록합니다.

용어 해설

GEMM: — GEMM은 행렬 A와 B를 곱해 출력 행렬을 만드는 연산으로, 대규모 신경망 추론과 학습의 핵심 계산량을 차지합니다. 이 글에서는 fp16 입력을 사용해 GPU 타일 단위로 처리량을 측정합니다.
WMMA: — WMMA는 GPU의 행렬 곱셈·누산 명령을 활용하는 방식입니다. 여러 warp가 출력 타일을 나눠 계산하며, 일반적인 스칼라 연산보다 행렬 연산 하드웨어를 직접 활용해 GEMM 처리량을 높입니다.
LDS: — LDS는 GPU 작업 그룹 내부의 스레드가 공유하는 빠른 로컬 메모리입니다. 이 커널은 A와 B의 타일을 두 개의 LDS 버퍼에 번갈아 적재해 한 타일을 계산하는 동안 다음 데이터를 준비합니다.
VGPR: — VGPR은 GPU 실행 단위가 사용하는 벡터 범용 레지스터입니다. 레지스터 사용량이 지나치게 크면 spill로 인해 메모리 접근이 늘어나지만, 이 커널은 188 VGPR을 사용하면서 register spill을 발생시키지 않았습니다.
split-K: — split-K는 GEMM의 K 차원 계산을 여러 작업 블록으로 나눠 병렬 처리하는 튜닝 매개변수입니다. hipBLASLt의 성능을 공정하게 측정하려면 이 값과 workgroup mapping을 후보별로 설정하고 실제 실행 시간을 비교해야 합니다.

언급된 도구

hipBLASLt중립

GEMM의 tile size, split-K, workgroup mapping 후보를 선택하고 GPU 행렬 연산을 실행하는 AMD math library입니다.

mojo-baro중립링크

Mojo GPU kernel, benchmark harness, protocol 파일을 포함해 RX 7900 XTX의 inference와 GEMM 성능을 측정하는 프로젝트입니다.

llvm-mc중립

gfx1100 아키텍처가 허용하는 fp32 WMMA 명령 mnemonic을 확인하는 검증 도구입니다.

언급된 리소스

GitHubmojo-baro

Gaming GPU에서 hipBLASLt를 넘은 fp16 GEMM

TL;DR

실용적 조언

섹션별 상세

용어 해설

언급된 도구

언급된 리소스

관련 토론

댓글

관련 기사

AMD V620에서 Qwen과 Gemma의 llama cpp 벤치마크

AMD GPU의 FP8 학습 최적화

AMD llama.cpp 브랜치의 Strix Halo 성능

Radeon 680M에서 양자화 포맷별 속도 역전

GPU 텍스처 유닛으로 LLM 추론 가속