NVIDIA H100에서 기본 CUDA 커널부터 shared memory, 레지스터 타일링, 벡터화, warp 타일링, Tensor Core, TMA, WGMMA까지 단계적으로 적용하며 GEMM 성능을 cuBLAS 수준까지 끌어올리는 과정을 기록한 워크로그입니다.
2026년 1월 12일
🚧 작업 진행 중입니다. 실수가 보이면 LinkedIn으로 연락 주세요.
행렬 곱셈은 현대 딥러닝의 핵심에 있습니다. 트랜스포머, CNN, 단순한 MLP까지 무엇이든 결국 GEMM으로 귀결됩니다. GPU는 이 연산을 대규모로 수행하도록 만들어졌고, cuBLAS 같은 라이브러리는 마지막 명령어 하나까지 조정된 커널로 성능 기준을 세웁니다.
이 블로그에서는 NVIDIA H100 위에서 그 경로를 바닥부터 다시 쌓아 올립니다. 가장 기본적인 커널에서 시작해 최적화를 점진적으로 덧붙입니다. shared memory 타일링, 레지스터 블로킹, 벡터화, warp 타일링, 그리고 Tensor Core, Tensor Memory Accelerator 같은 Hopper 전용 기능까지 다룹니다. 이 프로젝트는 Pranjal Shankhdhar와 Simon Boehm의 훌륭한 작업에서 영감을 받았고, 여기에 제 기여를 더해 전체 최적화 경로를 탐색하면서 결과를 재현할 수 있는 일관되고 재현 가능한 저장소를 제공하려고 합니다. 처음 일곱 개 커널에서는 FP32 정밀도만 사용합니다. 이 단계에서는 GEMM 성능 튜닝의 기반이 되는 핵심 최적화 기법에 집중하고 싶었고, 이 기법들은 대체로 아키텍처에 크게 의존하지 않습니다. FP32를 사용하면 Nsight Compute로 디버깅하기 쉬워지고 PTX와 SASS를 깔끔하게 살펴볼 수 있습니다. 두 번째 단계로 넘어가 Tensor Core와 H100 전용 기능을 활용하게 되면 혼합 정밀도로 전환합니다. 그 시점부터는 모든 벤치마크를 Tensor Core가 활성화된 cuBLAS와 비교합니다. 반면 첫 번째 단계에서는 순수 FP32 모드로 동작하는 cuBLAS와 비교합니다(저장소에는 혼합 정밀도 구현도 포함되어 있습니다).
목표는 단순히 최고 속도를 내는 것이 아닙니다. 각 변경이 실제로 무엇을 가져오는지, 각 단계에서 프로파일러가 무엇을 말해주는지, 그리고 커널이 순진한 형태에서 고도로 튜닝된 형태로 어떻게 진화하는지를 보는 것입니다. 마지막에는 직접 작성한 CUDA가 cuBLAS에 얼마나 가까이 갈 수 있는지, 그리고 고정된 행렬 크기에서는 오히려 앞설 수 있는지도 측정해볼 것입니다.
전체 코드는 GitHub에 있습니다. 모든 코드는 FP32와 BF16+FP32 혼합 정밀도를 지원하는 제 GitHub에서 확인할 수 있습니다.
시작해봅시다.
코드로 들어가기 전에 GPU 내부 하드웨어 구성 요소에 대한 명확한 정신 모델을 먼저 잡아두는 것이 도움이 됩니다. GPU 내부의 메모리 계층, 온칩 메모리와 오프칩 메모리의 크기 및 지연 시간 차이, 그리고 Hopper 아키텍처 계열에서 새로 도입된 구성 요소를 이해하면 이후 내용을 훨씬 쉽게 따라갈 수 있습니다. 이 섹션에서는 아직 CUDA 프로그래밍 모델을 본격적으로 다루지 않겠습니다. 대신 커널을 따라가며 개념을 점진적으로 소개할 것입니다. 저는 이것을 일종의 워크로그로 생각하고 있기 때문입니다. 따라서 이 섹션은 일종의 입문서 역할에 더 가깝습니다. 아래 그림은 Aleksa의 글을 확장해 전체 아키텍처를 자세히 보여줍니다.
최상위 수준에서 H100은 여러 개의 Graphics Processing Cluster(GPC)로 구성되어 있으며, 총 8개의 GPC가 있고 각 GPC는 18개의 Streaming Multiprocessor(SM)를 포함합니다. 네 개의 GPC는 하나의 L2 파티션에 직접 연결되고, 나머지 네 개는 두 번째 파티션에 연결됩니다. 이 SM들에는 칩의 주요 연산 유닛과 일부 "온칩" 메모리 구성 요소가 들어 있습니다. H100의 SXM 버전에는 132개의 SM이 있고(여기서 사용하는 버전), PCIe 버전에는 114개가 있습니다. 사실 8 * 18 = 144이므로 132개보다 많아야 할 것 같지만, 144는 전체 GH100 다이에 해당합니다. 실제 제품에서는 일부 SM이 비활성화되어 SXM 변형에서 132개의 기능하는 SM만 남습니다. H100 같은 현대 GPU는 매우 거대하고 극도로 복잡한 실리콘이기 때문에 결함 없이 제조하는 것이 사실상 불가능합니다. 단 하나의 SM만 불량이어도 칩 전체를 사용할 수 없게 됩니다. 이런 낭비를 피하기 위해 NVIDIA는 결함이 있거나 부분적으로 결함이 있는 SM을 비활성화해 더 적은 SM으로도 칩이 정상 동작하도록 만듭니다. 이 과정은 제조 수율을 높여줍니다. 아래는 SM 내부를 더 자세히 본 모습입니다.
SM 내부에는 위 그림에서 보이는 것처럼 네 개의 파티션이 있습니다. 각 SM은 다음과 같은 핵심 자원을 포함합니다.
CUDA 코어: 표준 부동소수점 연산(FLOPS)과 정수 연산(IOPS)을 처리합니다.
4세대 Tensor Core: 각 SM에는 4개의 특수 유닛이 포함됩니다. 이들은 고처리량 행렬 곱셈-누산 연산을 위해 설계되었으며, 현대 GPU 워크로드의 최고 성능을 달성하는 데 필수적입니다.
Load/Store (LD/ST) 유닛: SM과 메모리 계층 사이에서 데이터를 이동시키는 역할을 담당합니다.
SFU 유닛: sin, cos, sqrt, exp 같은 복잡한 수학 연산을 처리해 CUDA 코어의 부담을 덜어줍니다. 각 SM 파티션에는 자체 SFU가 있어 이런 연산을 일반 산술 연산과 병렬로 수행할 수 있습니다. MUFU로 시작하는 SASS 명령어(MUFU.SQRT, MUFU.EX2 등)를 보게 되면 이것들은 SFU에서 실행됩니다.
Dispatch 유닛: warp 스케줄러와 실행 파이프라인 사이의 다리 역할을 합니다. warp 스케줄러가 warp와 다음 명령을 선택하면 dispatch 유닛이 그 명령을 SM 안의 적절한 기능 유닛으로 보냅니다. 각 SM 파티션은 자체 dispatch 유닛을 가지므로 서로 다른 warp의 여러 명령이 서로 다른 실행 유닛으로 동시에 발행될 수 있습니다.
Warp 스케줄러: 각 SM에는 네 개의 warp 스케줄러가 있습니다(파티션당 하나). 각각은 32개 스레드 묶음인 warp에 명령을 발행하는 역할을 담당합니다(이건 조금 뒤에서 더 설명합니다!). warp 스케줄러는 클럭 사이클당 단 하나의 warp에만 하나의 명령을 발행할 수 있습니다. 따라서 네 파티션을 합치면 한 SM은 사이클당 최대 네 개의 warp 명령을 발행할 수 있고, 이는 어느 순간이든 128개 스레드가 병렬로 실행될 수 있음을 의미합니다. 모든 스케줄러를 완전히 활용하려면 블록당 충분한 활성 warp가 있어 아무 스케줄러도 놀지 않도록 해야 합니다. 그래서 일반적으로 블록당 128개 미만의 스레드를 시작하는 것을 피합니다. 그래야 모든 스케줄러가 작업할 warp를 하나씩 갖게 되기 때문입니다. 실제로 SM은 여러 스레드 블록을 동시에 올릴 수 있고 필요하면 다른 블록의 warp를 가져올 수도 있지만, SM 자원이 우연히 단 하나의 블록만 수용할 수 있는 상황에서는 이 경험칙을 기억해둘 만합니다.
이제 메모리 계층을 살펴보겠습니다. 각 메모리 타입이 GPU 안에서 물리적으로 어디에 존재하는지, 그리고 접근 지연 시간 면에서 어떻게 다른지를 봅시다. 이번에도 Aleksa의 글에 있는 피라미드 그림을 가져와 사용하겠습니다.
계층의 맨 아래부터 시작해 위로 올라가며, 가장 크고 느린 메모리에서 가장 작고 빠른 메모리로 이동해봅시다.
L2 캐시: 글로벌 메모리 위에는 L2 캐시가 있습니다. 이는 모든 SM이 공유하는 큰 온칩 캐시(SRAM으로 구성)입니다. 연산 코어와 느린 오프칩 HBM 사이의 주요 다리 역할을 하며, 최근 접근한 데이터를 캐싱해 지연 시간을 줄입니다. 물리적으로 두 부분으로 나뉘어 있고, 각 SM은 하나의 파티션에는 직접 연결되고 다른 파티션에는 crossbar를 통해 간접 연결됩니다.
Distributed Shared Memory (DSMEM): 메모리 계층에서 새로 등장한 요소입니다. DSMEM은 같은 GPC 안에서 여러 스레드 블록이 SM 간에 직접 데이터를 공유할 수 있게 해줍니다. 전통적인 shared memory를 단일 SM 바깥으로 확장한 것으로, 하나의 thread block cluster 안에서 최대 16개 블록이 협력할 수 있게 합니다. 지연 시간은 L2보다는 낮지만 개별 SM의 shared memory나 당연히 L1보다는 높습니다.
Shared Memory (SMEM) & L1 캐시: 둘은 같은 물리적 온칩 저장소 위에 공존하기 때문에 함께 묶어 설명합니다. 둘 다 SRAM 셀로 이루어져 있어 매우 빠르며, 피라미드 아래쪽의 다른 메모리 타입들보다 훨씬 낮은 지연 시간과 높은 대역폭을 가집니다. 둘을 합친 최대 총 용량은 256 KiB이고 메모리 대역폭은 31 TB/s입니다. L1 데이터 캐시는 SM의 LD/ST 유닛이 접근합니다. 이 256 KiB는 설정 가능하며, 더 큰 shared memory와 더 작은 L1 캐시를 맞바꾸거나 그 반대로 설정할 수 있습니다. 다만 shared memory에 할당할 수 있는 최대치는 228 KiB 정도입니다. L1 캐시를 위해 남겨둬야 하는 공간이 필요하기 때문입니다. 사실 위 H100 아키텍처 그림에서 보듯 이 228 KiB도 정확한 수치는 아닙니다. 블록당 1 KiB의 SMEM이 시스템 용도로 빠지므로, 실질적인 최대 설정 가능 크기는 228 − num_blocks * 1 KiB 입니다.
Register Memory (RMEM): 메모리 계층의 최하위이자 피라미드의 꼭대기에 있는 것이 레지스터입니다. 여기에는 단일 스레드가 조작하는 값이 저장됩니다. 레지스터는 각 스레드에 사적으로 속하지만 예외가 하나 있습니다. 같은 warp 안에 있는 스레드라면 다른 스레드의 레지스터를 읽을 수 있습니다. 이는 warp level shuffle primitives를 통해 가능합니다. 예를 들어 reduction 커널에서 자주 볼 수 있는데, 스레드 간 극도로 빠른 통신을 가능하게 합니다. 레지스터는 매우 빨라서 유효 대역폭이 124 TB/s 수준이고 지연 시간은 대략 한 클럭 사이클 정도입니다. 스레드의 레지스터 사용량이 가용 레지스터 파일을 초과하면 컴파일러는 값을 local memory로 spill하는데, 이것은 글로벌 메모리에 있으므로 훨씬 느립니다. CPU 프로그래밍과 마찬가지로 레지스터는 CUDA C/C++ 수준에서 직접 다루지 않습니다. PTX에서만 보이며 궁극적으로는 컴파일 시 ptxas가 할당합니다(아래 Compilation Story 참고). 컴파일러의 목표 중 하나는 스레드당 레지스터 사용량을 충분히 낮게 유지해 더 많은 스레드 블록이 동시에 하나의 SM에 상주할 수 있도록 하는 것입니다. 레지스터 압박이 높으면 occupancy가 줄어들기 때문입니다.
Tensor Memory Accelerator (TMA): Hopper 아키텍처에서 도입되었습니다. 글로벌 메모리와 shared memory 사이, 그리고 thread block cluster 안의 shared memory들 사이의 비동기 데이터 전송을 가능하게 합니다. 또한 shared memory bank conflict를 방지하기 위한 swizzling을 자동으로 수행해, 이전에는 개발자가 수동으로 관리해야 했던 복잡한 데이터 이동 및 레이아웃 패턴을 추상화합니다.
📖 컴파일 이야기
CUDA 프로그램이 소스 코드에서 최종 실행에 이르기까지의 여정은 NVCC 컴파일러 드라이버가 조율하는 다단계 컴파일 과정을 따릅니다. NVCC는 프로그램을 Host Code(CPU)와 Device Code(GPU)로 분리하며 전체 과정을 관리합니다.
Device Code는 먼저 PTX(Parallel Thread Execution, 저는 "피티엑스"라고 읽습니다 :))로 컴파일됩니다. PTX는 NVIDIA의 가상 ISA(Instruction Set Architecture)로, 아키텍처 독립적인 중간 표현(IR)을 제공합니다. 그다음 ptxas 어셈블러가 PTX 코드를 받아 필요한 최적화를 수행하고 SASS(Streaming ASSembler)라고 불리는 네이티브 ISA로 변환합니다. 이것이 사람이 읽을 수 있는 코드가 작성될 수 있는 가장 저수준 형식입니다. SASS 코드와 기타 메타데이터는 특정 GPU 아키텍처용 실행 컨테이너인 CUBIN(CUDA Binary)으로 묶입니다. 마지막으로 NVCC는 하나 이상의 CUBIN과 원래 PTX를 Fat Binary에 함께 묶고, 이것이 CPU 바이너리 코드와 함께 최종 실행 파일 안에 포함됩니다.
PTX를 포함하는 것은 전방 호환성을 위해 매우 중요합니다. Fat Binary가 실행될 때 미래 GPU에 일치하는 CUBIN이 없다면, 런타임이 포함된 PTX를 사용해 JIT(Just-In-Time) 컴파일을 수행하고 필요한 SASS를 생성해 실행을 보장합니다. 우리는 커널 2와 5에서 PTX와 SASS를 분석하며 왜 이것들이 유용한지 보게 될 것입니다.
이제 탄탄한 정신 모델을 세웠으니, 지금까지 이야기한 모든 것을 한눈에 모아 H100 아키텍처를 전체적으로 시각화한 그림으로 이 섹션을 마무리해봅시다.
CUDA 프로그래밍 모델에서 계산은 2단계 계층으로 조직됩니다. CUDA 커널을 한 번 호출하면 새로운 grid가 생성되고, 이 grid는 여러 block으로 이루어집니다. 각 block 안에는 스레드가 1D, 2D, 3D로 배치되었는지와 관계없이 총 1024개까지 있을 수 있습니다. 즉 blockDim.x * blockDim.y * blockDim.z <= 1024 입니다. grid 안의 모든 스레드는 같은 커널 함수를 실행하고, 스레드 인덱스를 이용해 자신을 구분하고 처리해야 할 데이터 부분을 식별합니다. 일반적으로는 하드웨어 효율성을 위해 thread block 각 차원의 스레드 수를 32의 배수로 두는 것이 권장됩니다. 이것은 곧 소개할 warp라는 개념과 정렬되기 때문입니다. 지금은 warp가 32개 스레드 묶음이라는 점만 기억하시면 됩니다. 따라서 차원을 이에 맞춰 두는 것이 유리합니다. 커널은 SIMT(Single Instruction, Multiple Threads) 실행 모델을 따르므로, 하나의 스레드 관점에서 작성됩니다. 따라서 CUDA 프로그래밍은 SPMD(Single Program, Multiple Data) 패러다임의 한 형태입니다.
커널 내부의 스레드가 __device__ 함수를 호출하면 그 함수는 바로 그 스레드 자신이 실행합니다. 이 함수는 자신을 호출한 스레드만 알고 있습니다. 기본적으로는 C++의 일반 함수와 같지만, 단 하나의 GPU 스레드 내부에서 일어난다는 점만 다르고, 이런 함수 인스턴스가 수천 개 병렬로 돌아간다고 생각하면 됩니다.
__global__ 함수는 커널입니다. GPU 실행용으로 컴파일되지만 CPU(호스트)에서 시작됩니다. 커널 호출은 block의 grid를 만들고, 그 block 안의 각 스레드는 독립적으로 커널 코드를 실행하기 시작합니다.
모든 thread block은 입력의 서로 다른 부분을 처리하므로, 임의의 순서로 실행될 수 있습니다. 따라서 block의 실행 순서나, block 안에서 어떤 스레드가 먼저 실행될지에 대해 절대 가정해서는 안 됩니다.
이제 막 익힌 CUDA 프로그래밍 모델과 그 아래 하드웨어에 대한 정신 모델을 연결해보면, 다음 그림은 단일 스레드의 시점을 시각화합니다. 이 스레드가 커널과 하드웨어 안에서 어디에 위치하는지, 어떤 메모리 공간과 상호작용하는지, 그리고 전체 grid 구조에 어떻게 들어맞는지를 보여줍니다. L1과 L2 캐시는 하드웨어가 관리하고 우리가 직접 제어하지 않기 때문에 일부러 생략했습니다.
이 첫 번째 커널에서는 block 안의 각 스레드(grid 안에 있는)가 C의 정확히 하나의 원소를 계산하도록 할당할 것입니다. 각 스레드는 자신의 좌표를 얻고, 대응되는 A의 행을 공유 차원 N을 따라 순회합니다(대부분 공식 자료에서는 여기를 K로 쓰지만, 저는 모든 커널에서 N으로 이미 통일해버렸기 때문에 일관성을 위해 그대로 쓰겠습니다). 동시에 스레드는 B의 대응 열을 따라 내려가며 곱셈 결과를 누적합니다. 루프가 끝나면 결과를 같은 좌표의 C에 다시 써 넣습니다.
미리 말해두자면: 이렇게 스레드 결과를 출력 원소에 1대1로 매핑하는 방식은 실제로 가장 효율적이지 않습니다(그렇게 짐작하셨다면 맞습니다). 이후 커널에서는 하나의 스레드가 출력의 여러 원소를 계산하게 만들겠지만, 일단은 여기서는 넘어가겠습니다.
아래는 이것이 어떻게 동작하는지에 대한 간단한 시각화이며, 단일 스레드 관점의 예시도 포함합니다.
흥미로운 점은 CUDA 프로그래밍 모델이 2D 좌표(x, y)를 지원함에도, 여기서는 block을 여전히 1D로 시작한다는 것입니다. Simon의 작업(커널 2)에서는 2D 시작 대신 1D 시작 후 재매핑을 사용하면 coalesced global memory access를 얻는 데 도움이 된다고 말합니다. 아이디어는 1D block을 시작한 뒤 threadIdx.x를 %와 /를 이용해 2D 좌표로 다시 해석하는 것입니다.
하지만 제가 두 접근법, 즉 재매핑을 사용하는 1D 시작과 일반적인 2D 시작을 모두 테스트했을 때 성능은 동일했습니다. 처음에는 Simon의 버전에서는 속도 향상이 있었다고 해서 꽤 혼란스러웠습니다. 핵심 차이는 Simon의 순진한 커널이 행렬 A를 열 방향으로 접근했다는 점입니다. 이것은 coalesced되지 않은 패턴입니다. 반면 그의 coalesced 커널은 blockIdx.x와 threadIdx.x가 행에 해당하도록 해서 A를 행 방향으로 접근합니다. 제가 한 방식과는 정반대입니다. 제 구현에서는 순진한 커널조차도 이미 A를 coalesced된 row-major 방식으로 접근하기 때문에, 두 버전 모두 자연스럽게 같은 효율을 달성합니다. 우습게도 이건 아주 기초적인 내용이었을지 모르지만, 처음에는 coalescing이 재매핑 꼼수에서 온다고 생각했기 때문에 꽤 오래 헷갈렸습니다.
즉 Simon의 속도 향상은 1D 대 2D block 레이아웃 자체가 아니라 메모리 접근 패턴을 고친 데서 옵니다. 제 순진한 커널은 이미 coalesced load를 사용하므로 시작 구성은 차이를 만들지 않습니다. 그래도 여기서는 1D 시작과 재매핑 방식을 계속 사용할 텐데, 메모리 coalescing을 이야기하기에 자연스러운 지점이 되기 때문이고, 동시에 Simon의 버전이 어떻게 non-coalesced 패턴을 만드는지도 보여줄 수 있기 때문입니다. 우선 warp가 무엇인지 정의해봅시다.
각 SM은 thread block 안의 스레드를 32개씩 묶어 warp로 만듭니다. warp는 warp 스케줄러가 한 번에 명령을 발행할 수 있는 기본 스케줄링 단위이며, 그 warp 안의 32개 스레드는 같은 명령을 lockstep으로 실행합니다. block은 먼저 row-major 순서의 1D 배열로 평탄화되고, 그다음 32개씩 연속 그룹으로 나뉩니다. warp 0은 스레드 0–31, warp 1은 32–63 식으로 실행합니다.
이 커널 코드는 다음과 같습니다.
template <const uint BLOCK_SIZE>
__global__ void sgemm_coalesced(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
int M, int N, int K, float alpha, float beta) {
// flattened IDs remapping
uint row = blockIdx.y * BLOCK_SIZE + (threadIdx.x / BLOCK_SIZE);
uint column = blockIdx.x * BLOCK_SIZE + (threadIdx.x % BLOCK_SIZE);
if (row < M && column < K) {
float cumulative_sum = 0.0f;
for (int n = 0; n < N; n++) {
cumulative_sum += A[row * N + n] * B[n * K + column];
}
C[row * K + column] = (alpha * cumulative_sum) + (beta * (C[row * K + column]));
}
}
제가 말한 재매핑은 여기서 일어납니다. 만약 block을 2차원으로 시작했다면 threadIdx.x % BLOCK_SIZE와 threadIdx.x / BLOCK_SIZE를 통해 계산하는 대신 threadIdx.x와 threadIdx.y를 직접 사용하면 되었을 것입니다.
저는 이 재매핑을 영화관 좌석 번호를 받았지만 행과 좌석 위치를 알려주지 않은 상황으로 시각화하는 것을 좋아합니다. 예를 들어 각 줄에 좌석이 6개 있고 제게 7번 좌석이 주어졌다고 합시다. 제 그림은 1부터 시작하는 번호를 쓰므로 먼저 1을 빼서 0부터 시작하는 체계로 바꿉니다. 6 = 7 - 1 이 됩니다. 한 줄의 좌석 수로 나누면 행 인덱스를 얻습니다. 6 / 6 = 1 이고, 이는 1부터 세는 기준으로 2번째 줄에 해당합니다. 나머지는 그 줄 안의 좌석 번호를 알려줍니다. 6 % 6 = 0 이고, 이것은 1부터 세는 기준으로 1번 좌석입니다. 따라서 7번 좌석은 두 번째 줄의 첫 번째 좌석입니다. 줄당 좌석 수로 나누면 몇 줄을 완전히 건너뛰는지가 나오고, 나머지는 그 줄 안에서의 좌석 위치를 알려줍니다.
uint row = blockIdx.y * BLOCK_SIZE + (threadIdx.x / BLOCK_SIZE);
uint column = blockIdx.x * BLOCK_SIZE + (threadIdx.x % BLOCK_SIZE);
32개 스레드로 이루어진 각 warp는 global memory load를 병렬로 다음처럼 실행합니다.
cumulative_sum += A[row * N + n] * B[n * K + column];
메모리 명령(global이든 shared든)은 warp 안에서 주소가 어떻게 분포하는지에 따라 재발행이 필요할 수 있습니다. 아직 shared memory를 프로그래밍 관점에서 소개하지 않았으니, 일단은 global memory에 집중합시다.
warp가 load를 실행하면 하드웨어는 32개 스레드가 연속된 메모리 위치를 접근하는지 확인합니다. 최선의 경우는 모든 스레드가 연속된 주소를 읽는 경우입니다. 그러면 하드웨어는 32개 요청을 하나의 트랜잭션으로 coalesce할 수 있습니다.
global memory는 디바이스 DRAM에 있고, DRAM은 32, 64, 128바이트 단위로 접근됩니다. 트랜잭션 수가 적을수록 효율이 높습니다. 이를 FP32 load 기준(스레드당 4바이트)으로 설명해보겠습니다. 만약 각 스레드의 4바이트 load마다 각각 32바이트 트랜잭션이 필요하다면 처리량은 8배나 떨어집니다.
예를 들어:
이제 이 coalesced 커널과 non-coalesced 커널이 어떻게 보이는지 global memory 접근 패턴을 분석해봅시다.
이 커널을 실행하면 4.2 TFLOP/s 처리량이 나오며, 이는 FP32 cuBLAS 커널 성능 대비 약 **8.2%**입니다. 하드웨어가 이론적으로 낼 수 있는 수준, 흔히 Speed of Light(SoL)라고 부르는 값과는 아직 거리가 멉니다. Speed of Light는 물리와 칩 설계만을 기준으로 했을 때 GPU가 낼 수 있는 연산 처리량의 이론적 상한을 말합니다. tensor core 워크로드에서는 이 상한이 perf = freq_clk_max * num_tc * flop_per_tc_per_clk 로 주어집니다(H100 SXM5 기준 BF16 Tensor Core 최대 989 TFLOP/s, FP32 최대 66.9 TFLOP/s).
이 수치는 보통 고정된 값처럼 제시되지만, 실제로는 전혀 상수가 아닙니다. GPU가 유지할 수 있는 실제 클럭 주파수에 따라 계속 움직이고, 이 주파수는 전력과 열 제한에 따라 달라집니다. GPU가 전력 제한에 가까워질수록 전압 조절기가 전압을 낮추고, 클럭 속도가 떨어지며, 유효 SoL도 함께 낮아집니다. 이런 동작을 power throttling이라고 합니다.
Horace He는 이것을 간단한 matmul 벤치마크로 아주 멋지게 보여주었습니다. PyTorch에서 큰 matmul은 약 258 TFLOPs를 냈지만, 같은 연산을 CUTLASS profiler 안에서 실행하면 약 288 TFLOPs가 나와 10–11% 향상처럼 보였습니다. 커널 수준의 진짜 속도 향상처럼 보였던 것이죠. 하지만 같은 CUTLASS 커널을 Python에서 바인딩해 같은 입력으로 돌리자 그 이점은 사라졌습니다. 차이는 CUTLASS profiler는 텐서를 정수로 초기화하고, PyTorch는 난수를 사용한다는 점뿐이었습니다.
이것이 중요한 이유는 칩에서 전력이 어떻게 소비되는지에 뿌리를 두고 있습니다. 정적 전력은 트랜지스터를 켜진 상태로 유지하는 데 쓰이고, 동적 전력은 트랜지스터가 상태를 바꿀 때마다 소비됩니다. 난수는 수십억 개 트랜지스터 전반에서 무질서한 비트 전환을 일으키므로 동적 전력을 증가시키고 throttling을 유발합니다. 반면 0이나 단순한 정수 패턴 같은 예측 가능한 값은 훨씬 적은 비트만 바꾸고, 동적 전력을 낮게 유지해 GPU가 더 높은 클럭을 유지하게 합니다. 즉 커널이 “더 빠르게” 보이는 이유는 코드가 더 효율적이어서가 아니라, 하드웨어가 전기적으로 덜 스트레스를 받기 때문입니다.
그래서 실제 커널은 광고된 최대 TFLOP/s에 거의 도달하지 못합니다. 이론적 SoL은 최대 클럭 주파수를 가정하지만, 실제 워크로드는 계속 전력 및 열 제약에 부딪힙니다. 진짜 상한은 전압, 클럭 속도, 온도, 심지어 입력 데이터의 무작위성에 따라서도 변합니다.
. 이는 예상된 결과입니다. 아직 초반 단계이고, 가장 명백한 병목 중 하나는 매 반복마다 global memory에 손을 뻗어야 한다는 사실이기 때문입니다. 앞서 말했듯 GMEM 접근은 대략 500사이클이 들지만, shared memory(SMEM) 접근은 대략 20~30사이클입니다. 다음 커널에서는 연산을 수행하기 전에 스레드들이 협력해서 GMEM의 값을 SMEM으로 불러오도록 하여 성능을 개선할 것입니다. 타일이 SMEM에 올라오면, 스레드는 더 이상 반복해서 GMEM에 가지 않고 그곳에서 피연산자를 가져올 수 있습니다. 이는 속도를 크게 높여주고 더 높은 처리량에 가까워지게 해줍니다.
이 커널의 논리는 다음과 같습니다. A와 B에서 가져온 타일을 저장할 shared memory 공간 sharedA와 sharedB를 각각 할당합니다. 각 타일의 원소 수는 TILE_SIZE * TILE_SIZE 가 되고, 이는 grid를 dim3 blockDim(32 * 32)로 시작할 때 block당 스레드 수와 일치합니다. 즉 이전처럼 각 스레드는 여전히 출력의 단일 원소를 계산합니다. 여기에 더해 각 스레드는 타일 반복마다 A에서 하나, B에서 하나, 총 두 값을 shared memory에 적재합니다.
이 점을 강조하는 이유는 이후 커널에서는 각 스레드가 하나보다 많은 원소를 적재하게 되고, 그때는 그 스레드가 shared memory 안 어디에 써야 하는지 결정하는 추가 인덱싱 로직이 필요하기 때문입니다. 하지만 이 커널에서는 아직 그럴 필요가 없습니다. 각 스레드가 각 행렬에서 정확히 하나의 원소만 적재하므로 ty와 tx만 알면 충분합니다.
이건 말로 설명하는 것보다 그림으로 보는 편이 훨씬 쉽기 때문에, 먼저 4 x 4 행렬 A와 B를 쓰는 작은 예시로 아이디어를 보여준 뒤 실제 커널에서 같은 논리가 어떻게 보이는지 보여드리겠습니다.
이 논리를 실제 시작 구성에 맞춰 적용하면 커널은 다음과 같습니다.
전체 코드는 다음과 같습니다.
template <const uint TILE_SIZE>
__global__ void sgemm_tiled_shared(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
int M, int N, int K, float alpha, float beta) {
// Allocate shared memory
__shared__ float sharedA[TILE_SIZE * TILE_SIZE];
__shared__ float sharedB[TILE_SIZE * TILE_SIZE];
// Identify the tile of C this thread block is responsible for (We assume tiles are same size as block)
const uint block_row = blockIdx.y;
const uint block_column = blockIdx.x;
// Calculate position of thread within tile (Remapping from 1-D to 2-D)
const uint ty = threadIdx.x / TILE_SIZE; // (0, TILE_SIZE-1)
const uint tx = threadIdx.x % TILE_SIZE; // (0, TILE_SIZE-1)
// Move pointers from A[0], B[0] and C[0] to the starting positions of the tile
A += block_row * TILE_SIZE * N; // Move pointer (block_row * TILE_SIZE) rows down
B += block_column * TILE_SIZE; // Move pointer (block_column * TILE_SIZE) columns to the right
C += (block_row * TILE_SIZE * K) + (block_column * TILE_SIZE); // Move pointer (block_row * TILE_SIZE * K) rows down then (block_column * TILE_SIZE) columns to the right
// Calculate how many tiles we have
const uint num_tiles = CEIL_DIV(N, TILE_SIZE);
float cumulative_sum = 0.0f;
// Iterate over tiles (Phase 1: Loading data)
for (int t = 0; t < num_tiles; t++) {
sharedA[ty * TILE_SIZE + tx] = A[ty * N + tx];
sharedB[ty * TILE_SIZE + tx] = B[ty * K + tx];
__syncthreads();
// Phase 2: Compute partial results iteratively
for (int i = 0; i < TILE_SIZE; i++) {
cumulative_sum += sharedA[ty * TILE_SIZE + i] * sharedB[i * TILE_SIZE + tx];
}
__syncthreads();
// Move all pointers to the starting positions of the next tile
A += TILE_SIZE; // Move right
B += TILE_SIZE * K; // Move down
}
// Write results back to C
C[ty * K + tx] = (alpha * cumulative_sum) + (beta * C[ty * K + tx]);
}
이 커널은 이전 커널보다 처리량이 대략 1.7배 향상되어 cuBLAS(FP32) 대비 13.9% 수준에 도달하지만, Nvidia의 Nsight Compute로 프로파일링해보면 몇 가지 핵심 문제가 드러납니다.
먼저 프로파일러의 Speed of Light 섹션을 보면 흥미로운 점이 있습니다. Compute throughput이 76.63%, Memory throughput이 91.13%로 나옵니다. 이 수치는 FP32 기준 하드웨어 SoL 대비 비율입니다. 처음 보면 76.63% compute throughput이 꽤 괜찮아 보이지만, 이건 비교적 단순한 GEMM 커널이므로 그렇게 해석하면 말이 되지 않습니다.
퍼센트로 병목을 보여주는 throughput breakdown을 보면 오해가 바로 풀립니다. SM: Inst Executed Pipe Lsu = 76.63% 입니다(개요에 보인 값과 같은데, 개요는 instruction breakdown 중 가장 높은 비율을 표시하기 때문입니다). Pipe LSU는 load/store 유닛입니다. 프로파일러 설명에 따르면 그 역할은 다음과 같습니다.
"LSU pipeline issues load, store, atomic, and reduction instructions to the L1TEX unit for global, local, and shared memory. It also issues special register reads (S2R), shuffles, and CTA level arrive or wait barrier instructions to the L1TEX unit."
그다음으로 높은 항목은 SM: Mio Inst Issued = 40.08%인데, 이는 memory input or output 유닛입니다. 즉 명령 발행의 거의 절반이 LSU 연산과 비메모리 연산으로 갈라져 있고, 다시 말해 메모리 쪽이 지배적이라는 뜻입니다.
그렇다면 FP32 계산을 하는 명령 수는 어떨까요? SM: Pipe FMA Cycles Active = 14.81% 입니다. 우리가 진짜 봐야 할 수치는 이것입니다. FP32 하드웨어 능력을 전혀 활용하지 못하고 있다는 뜻이며, 이 시점에서는 정확히 예상되는 결과입니다. 이 지표는 활성 SM 사이클 중 FP32 FMA 실행 파이프가 실제로 일을 하고 있었던 비율이 얼마인지를 알려줍니다. 좋은 GEMM이라면 이 수치를 매우 높게 보고 싶습니다(대략 60~80% 이상). GEMM은 거의 전부 FMA이기 때문입니다. 스포일러를 하자면, 커널 6에서 이 수치를 62%까지 끌어올리는 방법을 보게 됩니다.
따라서 개요의 compute throughput 숫자(최댓값으로 선택된 그 값)는 다소 오해를 부를 수 있었습니다. 전체 값은 SM에서 실행된 모든 명령(ALU, FMA, SFU, LSU 등)을 함께 고려하기 때문입니다. 물론 다른 유용한 지표들도 많지만 지금은 넘어가겠습니다. 프로파일링 리포트는 다운로드 가능한 .ncu-rep 형식으로 GitHub repository에 모두 올려두었으니 Nsight Compute에서 열어 자세히 보고 싶으시면 참고하시면 됩니다. 이제 메모리 throughput breakdown으로 시선을 옮겨봅시다.
상위 세 지표는 L1: Data Pipe Lsu Wavefronts = 91.13%, L1: Lsu Writeback Active = 87.07%, L1: Lsuin Requests = 76.63% 입니다. 이쯤 되면 우리가 shared memory로 오가는 엄청난 양의 load/store 요청으로 simply shared memory를 압도하고 있다는 점이 명확해집니다.
프로파일러에서 이 문제를 바라볼 수 있는 흥미로운 지표와 각도가 아직 남아 있으므로, 아래 그림에서는 프로파일러의 세 가지 다른 뷰를 주석과 함께 보여드립니다. SASS 코드는 프로파일러의 source 섹션에서 직접 보거나, 이 GoodBolt 링크에서 전체 코드를 볼 수 있습니다.
제가 roofline plot에서 지적하듯이
Arithmetic intensity는 커널에서 산술 연산 수와 메모리 연산 수의 비율입니다. , 우리는 arithmetic intensity를 높이거나, 시각적으로는 그래프에서 오른쪽으로 이동하고 싶습니다. 현대 GPU에서는 산술 대역폭과 메모리 대역폭의 비율이 매우 높기 때문에, 가장 효율적인 커널은 높은 arithmetic intensity를 가집니다. 이는 메모리 병목을 해소하면 종종 메모리 서브시스템의 작업을 연산 서브시스템으로 옮길 수 있고, 메모리 대역폭을 절약하는 동시에 산술 유닛의 부하를 늘릴 수 있음을 뜻합니다.
따라서 다음 커널에서는 각 스레드가 출력 행렬의 한 원소만 계산하게 두는 대신, 여러 원소를 계산하게 만들 것입니다. 각 스레드는 여러 결과를 자신의 레지스터에 부분적으로 누적하고, 계산이 완전히 끝난 뒤에만 레지스터에서 C로 최종 값을 저장합니다.
작성 중입니다. 모든 커널 코드는 GitHub에서 확인할 수 있습니다.
커널의 arithmetic intensity를 조금 더 짜내기 위해 이제는 각 스레드가 단일 출력 원소보다 더 많은 값을 계산하게 합니다. 아래 그림처럼, thread block이 출력 행렬의 하나의 타일을 덮고 그 타일 안에서 각 스레드는 자신만의 작은 2D 패치를 맡아 ROWS_PER_THREAD * COLS_PER_THREAD 개의 결과를 계산합니다.
이제는 shared memory 타일 안의 원소 수보다 더 적은 스레드를 시작하므로, 각 스레드는 global memory에서 shared memory로 여러 원소를 적재해야 합니다. 이를 위해 stride를 사용해 모든 타일 원소를 겹침 없이 덮습니다. A와 B의 shared memory 타일이 완전히 채워지면, 각 스레드는 자신이 필요한 A와 B 조각을 shared memory에서 레지스터로 반복해서 읽고, 부분 outer product 업데이트를 수행하며, 최종 ROWS_PER_THREAD * COLS_PER_THREAD 크기의 C 블록이 완성될 때까지 로컬 레지스터에 누적합니다.
이 커널 코드는 다음과 같습니다.
template <const uint TILE_SIZE_M, const uint TILE_SIZE_N, const uint TILE_SIZE_K, const uint ROWS_PER_THREAD>
__global__ void sgemm_1D_registertiling(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C,
int M, int N, int K, float alpha, float beta) {
// Allocate shared memory
__shared__ float sharedA[TILE_SIZE_M * TILE_SIZE_N];
__shared__ float sharedB[TILE_SIZE_N * TILE_SIZE_K];
// Identify the tile of C this thread block is responsible for
const uint block_row = blockIdx.y;
const uint block_column = blockIdx.x;
// Calculate position of thread within tile (Remapping from 1-D to 2-D)
const uint ty = threadIdx.x / TILE_SIZE_K;
const uint tx = threadIdx.x % TILE_SIZE_K;
// Move pointers from A[0], B[0] and C[0] to the starting positions of the tile
A += block_row * TILE_SIZE_M * N;
B += block_column * TILE_SIZE_K;
C += (block_row * TILE_SIZE_M * K) + (block_column * TILE_SIZE_K);
// Calculate position of thread within shared memory tile
const uint smem_ty_A = threadIdx.x / TILE_SIZE_N;
const uint smem_tx_A = threadIdx.x % TILE_SIZE_N;
const uint smem_ty_B = threadIdx.x / TILE_SIZE_K;
const uint smem_tx_B = threadIdx.x % TILE_SIZE_K;
// Calculate number of tiles
const uint num_tiles = CEIL_DIV(N, TILE_SIZE_N);
// Initialise thread-local results in registers
float thread_results[ROWS_PER_THREAD] = {0.0f};
// Iterate over tiles
for (int t = 0; t < num_tiles; t++) {
sharedA[smem_ty_A * TILE_SIZE_N + smem_tx_A] =
A[smem_ty_A * N + smem_tx_A];
sharedB[smem_ty_B * TILE_SIZE_K + smem_tx_B] =
B[smem_ty_B * K + smem_tx_B];
__syncthreads();
// Inner computation loop
for (int i = 0; i < TILE_SIZE_N; i++) {
float fixed_B = sharedB[i * TILE_SIZE_K + tx];
for (int row = 0; row < ROWS_PER_THREAD; row++) {
uint global_row_idx = ty * ROWS_PER_THREAD + row;
thread_results[row] +=
sharedA[global_row_idx * TILE_SIZE_N + i] *
fixed_B;
}
}
__syncthreads();
// Move to next tile
A += TILE_SIZE_N;
B += TILE_SIZE_N * K;
}
// Write results back to C
for (int row = 0; row < ROWS_PER_THREAD; row++) {
uint global_row_idx = ty * ROWS_PER_THREAD + row;
C[global_row_idx * K + tx] =
(alpha * thread_results[row]) +
(beta * C[global_row_idx * K + tx]);
}
}
이전 커널은 메모리 IO stall을 줄였지만, 출력 하나당 여전히 shared memory read를 너무 많이 하고 있었습니다.
이 커널에서는 각 스레드가 수직 스트립 하나가 아니라 행과 열의 타일을 계산합니다. 그 결과 다음처럼 줄었습니다.
이는 출력 하나당 SMEM load 트래픽이 4.5배 줄고 GMEM은 2배 줄어든 것이며, 동시에 스레드당 8배 더 많은 결과를 계산합니다.
이 커널을 프로파일링해보면 H100의 FP32 피크 대비 38%에 도달하는데, 이는 즉시 compute-bound가 아니라는 뜻입니다. global memory도 전혀 포화되지 않았습니다. 이 커널은 DRAM 대역폭의 2.90%, L2의 약 10.13%만 사용합니다. GEMM 워크로드로서는 둘 다 매우 낮은 수치입니다. 따라서 global bandwidth도 L2 트래픽도 병목이 아닙니다.
첫 번째 의미 있는 신호는 Speed of Light 섹션에서 나타납니다.
프로파일러는 심지어 병목 방향을 직접 가리켜줍니다.
“The kernel is utilising greater than 80% of the available compute or memory performance. To further improve performance, work will likely need to be shifted from the most utilised unit to another. Start by analysing L1 in the Memory Workload Analysis section.”
DRAM과 L2는 거의 건드리지 않는데 L1/TEX는 90%에 가까운 활용률을 보이므로, 압박이 온칩 메모리 계층에 집중되어 있다는 뜻이 분명합니다. 다시 말해 이것은 DRAM 문제가 전혀 아닙니다. 제한 요인은 L1 cache/SMEM 경로의 대역폭과 지연 시간입니다.
이 그림은 스케줄러 지표로도 강화됩니다. SM Issue Active는 55.50%인데, 이는 warp 스케줄러가 전체 사이클의 절반 조금 넘는 시간에만 명령을 발행한다는 뜻입니다. 나머지 약 45% 사이클은 주로 DRAM이나 L2가 아니라 L1/SMEM을 통한 데이터 이동을 기다리며 stall 상태로 소비됩니다.
스케줄러 통계를 보면 다음과 같습니다.
"Every scheduler is capable of issuing one instruction per cycle, but for this kernel each scheduler only issues an instruction every 1.8 cycles. This might leave hardware resources underutilized and may lead to less optimal performance."
Stall MIO throttle도 0.59이므로 이 값을 줄이고 싶습니다. 정리하면, 매우 낮은 DRAM 사용량, 매우 낮은 L2 사용량, 높은 L1/TEX 활용률(약 88%), 그리고 50%대 중반의 SM issue는 모두 같은 결론으로 모입니다. 이 커널은
L1-bound 또는 SMEM-bound입니다. compute-bound가 아니고, 당연히 GMEM-bound도 아닙니다. 병목은 레지스터, shared memory, L1/TEX 경로 사이의 온칩 데이터 이동입니다. 다음 최적화에서는 L1/SMEM 파이프라인 내부의 instruction overhead를 줄여야 하며, 이를 통해 1.8 cycles per instruction 문제를 직접 겨냥하고 warp 스케줄러가 SM Issue Active를 높일 수 있도록 해야 합니다.
제약이 있음에도 불구하고, 이전 커널 대비 1.40배 향상을 얻어 12.2 TFLOPs에서 19.1 TFLOPs로 올라갔고, cuBLAS 대비 36.8% 수준에 도달했습니다.
지금까지의 모든 이전 커널에서는 스칼라 하나당 load instruction 하나를 발행했습니다. coalescing을 최적화해왔기 때문에 실제로는 스칼라 하나당 load 하나가 아니라고 느껴질 수 있어 약간 헷갈릴 수 있습니다. 여기서 빠진 디테일은 다음 두 가지를 구분해야 한다는 점입니다.
coalescing은 첫 번째만 도와줍니다. 서로 다른 스레드가 요청한 데이터를 하드웨어가 하나의 연속된 메모리 트랜잭션으로 합칠 수 있게 해줍니다. 하지만 두 번째는 전혀 바꾸지 않습니다. 각 스칼라 load는 여전히 별도의 명령으로 나타나고, 각각 warp 스케줄러에 의해 발행되어 load/store 파이프라인을 통과해야 합니다.
이를 확인하기 위해 이전 커널의 SASS를 살펴보면, 특히 GMEM에서 SMEM으로 load하는 부분에서 실제로 루프 반복마다 스레드당 별도의 instruction issue가 생성되는 것을 볼 수 있습니다. 예를 들어 B에 대한 접근이 coalesced된다는 것은 알고 있고, 이것은 메모리 컨트롤러 수준에서는 목적을 달성합니다. 하지만 하드웨어는 여전히 스레드당 네 번의 별도 instruction issue를 필요로 합니다. coalescing은 컴파일 타임에 일어나는 것이 아닙니다. 실제 주소가 알려진 런타임에 하드웨어가 동적으로 수행합니다. 컴파일러는 행렬 포인터가 함수 인자로 넘어오기 때문에 정렬이나 레이아웃을 가정할 수 없으므로, 이것이 자연스럽습니다. 이제 이것을 warp 전체로 확장해봅시다. 커널이 이 구간을 한 번 돌 때마다 warp 스케줄러는 총 32 threads * 4 loads/thread = 128개의 별도 load instruction을 발행해야 합니다. 메모리 컨트롤러는 128개의 스칼라 요청을 몇 개 안 되는 크고 효율적인 메모리 트랜잭션으로 병합할 수 있겠지만, 파이프라인 병목은 여전히 프런트엔드에 남습니다. warp 스케줄러는 과로 상태이고, load/store 파이프라인은 반복적인 요청으로 포화되어 실제 연산 유닛이 쓸 수 있는 클럭 사이클을 빼앗고 있습니다.
이 instruction pressure를 해결하는 유일한 방법은 컴파일러가 데이터 전송을 바라보는 방식을 근본적으로 바꾸는 것입니다. 하나의 load 요청을 봤을 때 여러 스칼라를 한 번에 가져오라고 컴파일러에 알려줘야 합니다. CUDA는 이런 압박을 줄일 수 있도록 벡터화 변수를 제공합니다. float2, float4 같은 더 넓은 데이터 타입이 여기에 해당합니다. 일반 float는 32비트(4바이트)이고, float4는 128비트(16바이트)입니다. float*를 float4*로 캐스팅한 뒤 단 한 번의 load를 발행하면 128비트 global memory instruction 하나가 발생합니다. SASS에서는 이것이 LDG.E.128 또는 때로는 LDG.E.CI.128로 보입니다. 이렇게 하면 load instruction 수가 극적으로 줄고, 하드웨어 메모리 경로도 더 효율적으로 사용할 수 있습니다.
그럼 다음 커널을 살펴봅시다. 전체 구조는 동일하고, 유일한 차이는 sharedA를 전치한다는 점입니다. 따라서 GMEM에서 한 행을 읽어오면 sharedA에서는 그것이 하나의 열로 저장됩니다. 이렇게 하는 이유는 벡터화된 load를 사용하기 위해서입니다. 벡터화된 load instruction을 발행하려면 주소가 물리적으로 연속되어 있어야 하기 때문입니다.
핵심은 행렬 A의 접근 패턴입니다. 각 스레드는 계산에 필요한 ROWS_PER_THREAD 개 원소를 load해야 하는데(이는 열을 따라 이동하는 것에 해당), 기본 메모리 접근은 stride가 있습니다. 이 열 데이터를 전치하면 sharedA에 물리적으로 연속된 행으로 저장됩니다. 그러면 이후 SMEM에서 reg_m으로 이동할 때 float4 load를 발행할 수 있습니다.
sharedB도 같은 접근을 적용하지만 전치하지는 않을 것입니다. 지난 커널에서 기억하시겠지만, 각 스레드는 COLS_PER_THREAD를 reg_k로 load하고, 이 원소들은 이미 같은 행에서 서로 붙어 있으므로 전치할 필요가 없습니다.
이렇게 하면 scalar + offset 방식은 완전히 버리게 되고, 따라서 GMEM에서 SMEM으로 load할 때 stride를 두는 루프도 필요 없게 됩니다. 또한 레지스터로 load할 때는 float4를 발행할 수 있고 ROWS_PER_THREAD와 COLS_PER_THREAD가 모두 8이므로, 레지스터 load는 여전히 루프를 돌지만 이제 4칸씩 stride를 밟게 됩니다. 이 모든 것이 말로 들으면 다소 헷갈릴 수 있으니, 늘 그렇듯 시각적으로 그려보겠습니다.
이제 이 커널의 SASS를 살펴보고 이전 버전과 비교한 뒤, 프로파일러가 무엇을 말해주는지 봅시다.
우리 접근법을 사용하자 스레드당 발행되는 명령 수가 8개에서 단 2개로 줄어든 것을 볼 수 있습니다. 위 그림은 LDG.E.CI.128을 사용하는 GMEM load 단계만 보여주지만, SMEM에서 RMEM(Register Memory)으로 읽을 때도 마찬가지로 instruction 수를 줄였습니다. 여기에는 길어서 다 넣지 않겠지만, SASS에서 LDS.U.128 명령을 분명히 볼 수 있고, 이를 통해 shared memory read도 벡터화하는 데 성공했음을 확인할 수 있습니다. 더 자세히 보고 싶다면 이 커널의 전체 SASS/PTX도 볼 수 있습니다.
이 커널을 실행하면 또 한 번 약 2배 속도 향상을 얻고, 이제 37.2 TFLOP/s로 cuBLAS 성능의 **72%**에 도달합니다. 좋습니다. 점점 가까워지고 있습니다(FP32 경로만 기준입니다! Tensor Core를 쓰는 cuBLAS를 이기기까지는 아직 갈 길이 멉니다)
프로파일러를 보면 이제 커널이 GPU를 꽤 잘 활용하고 있습니다. compute throughput은 피크의 약 66%, memory throughput은 약 85%, 디바이스 FP32 roofline의 56%에 도달합니다.
참고로 cuBLAS는 약 85% 정도이고, 실제 워크로드가 하드웨어 이론 피크의 100%에 도달하는 일은 거의 없습니다.
이제 이전에 찾았던 문제들과 이 커널의 새로운 프로파일 결과를 비교해봅시다.
SM Issue Active: 55.50%에서 66.05%로 증가(연산에 더 많은 시간이 쓰입니다. 스케줄러가 약 19% 더 바빠졌습니다).SM Pipe Fma Cycles Active: 42.00%에서 56.73%로 증가(더 많은 계산이 수행됩니다).SM Inst Executed Pipe Lsu (Load Store Unit): 28.78%에서 17.09%로 감소(instruction count 감소의 증거).SM Mio Inst Issued: 14.99%에서 9.21%로 감소.Stall MIO Throttle: 0.59에서 0.02로 감소.또한 스케줄러가 명령당 1.8사이클만 발행한다는 경고도 이제 사라졌습니다. 엄청난 개선입니다. 하지만 아직 손볼 여지는 있습니다.
자세히 보면 현재 성능을 깎아먹고 더 높은 compute throughput을 막고 있는 몇 가지 미묘한 지표 차이가 있습니다.
가장 치명적인 것은 shared memory 접근에서 심한 bank conflict가 보인다는 점입니다. load에서는 약 5-way conflict, store에서는 2.6-way conflict가 나타나며, shared memory wavefront의 40% 이상이 직렬화로 낭비되고 있습니다.
Nsight Compute에서 wavefront는 shared memory 요청을 하드웨어가 한 사이클에 처리할 수 있는 단위를 뜻합니다. bank conflict가 발생하면 요청은 여러 wavefront로 나뉘고, 이들이 순차적으로 처리되어 stall을 유발합니다.
지금까지 커널에서 bank conflict를 본격적으로 고려하지 않았으므로, 지금이 이 개념을 소개하기에 적절한 시점입니다.
이를 더 잘 시각화하려고 shared memory 구성을 그려봤는데, 핵심 개념은 NVIDIA GPU(H100 포함)의 shared memory가 32개의 bank로 나뉘어 있고, 각 bank가 사이클당 4바이트 word 하나를 처리할 수 있다는 점입니다. 저는 이것을 슈퍼마켓의 32개 계산대처럼 생각하는 편입니다. 각 계산대는 사이클당 손님 한 명을 처리합니다. 여기서 “word”는 기본 저장 단위인 4바이트를 뜻합니다(예를 들어 float 하나).
bank index는 잘 알려진 modulo 방식으로 간단히 계산할 수 있습니다.
bank_index = word_index % 32
이제 32개 스레드로 이루어진 warp가 shared memory 접근을 발행할 때:
이 그림에서는 32개 lane 모두가 같은 banked word를 읽을 수 있습니다. 하드웨어가 값을 효율적으로 broadcast하므로 conflict가 발생하지 않습니다.
먼저 store conflict부터 보겠습니다. 우리 코드에서 store의 약 2.6 conflict는 전치된 타일로 sharedA를 채울 때 나타납니다.
// Populate smem using vector loads
float4 tempA = reinterpret_cast<const float4*>(&A[smem_ty_A * N + smem_tx_A*4])[0]; // [0] dereference issues one ld.global.nc.v4.f32
// Transpose A (instead of 128x8 previously for ex, now it will be 8x128)
sharedA[(smem_tx_A * 4 + 0) * TILE_SIZE_M + smem_ty_A] = tempA.x;
sharedA[(smem_tx_A * 4 + 1) * TILE_SIZE_M + smem_ty_A] = tempA.y;
sharedA[(smem_tx_A * 4 + 2) * TILE_SIZE_M + smem_ty_A] = tempA.z;
sharedA[(smem_tx_A * 4 + 3) * TILE_SIZE_M + smem_ty_A] = tempA.w;
smem_ty_A는 전치된 sharedA에서 열 방향으로 움직이고, smem_tx_A는 이 커널 구성에서는 0 또는 1입니다. 각 스칼라 store에 대한 word index는 다음과 같습니다.
word_index = (smem_tx_A*4 + q) * TILE_SIZE_M + smem_ty_A → q는 {0,1,2,3}bank = word_index % 32TILE_SIZE_M = 128이면 leading stride가 32 bank로 나누어떨어집니다. 128 % 32 = 0 이므로, bank는 stride 항이 아니라 계산에서 살아남는 offset에만 의존하게 됩니다.
결국 bank는 사실상 smem_tx_A 값(행 인덱스)이 아니라 smem_ty_A 값(열 offset)에만 의존하게 됩니다. 각 두 스레드가 같은 smem_ty_A 값을 공유하므로, 네 개의 스칼라 store 각각에서 모두 같은 bank를 겨냥하게 됩니다. 바로 이것이 프로파일러가 지적한 2-way store conflict 패턴입니다.
leading stride가 32 word의 배수일 때 이런 종류의 conflict를 피하는 흔한 방법 중 하나가 padding입니다.
// Allocate shared memory. Use padded leading strides that keep float4 alignment
constexpr uint STRIDE_A = (TILE_SIZE_M % 32u == 0u) ? (TILE_SIZE_M + 4u) : TILE_SIZE_M;
constexpr uint STRIDE_B = (TILE_SIZE_K % 32u == 0u) ? (TILE_SIZE_K + 4u) : TILE_SIZE_K;
static_assert((STRIDE_A % 4u) == 0u, "STRIDE_A must keep float4 alignment");
static_assert((STRIDE_B % 4u) == 0u, "STRIDE_B must keep float4 alignment");
leading stride를 132 word로 padding하고 sharedA를 건드리는 모든 곳(전치 store와 이후 read 모두)에서 그 stride를 사용하면, 이제 행을 나타내는 smem_tx_A가 bank에 영향을 미치게 됩니다. 이전에 충돌하던 두 lane은 16개 bank 떨어진 위치로 분리되고, x, y, z, w 네 개의 스칼라 store도 하나의 bank에 몰리지 않고 여러 bank를 순환하게 됩니다. 이를 증명하기 위해 padding 후 커널을 프로파일링해보니 store conflict가 제거된 것을 확인했습니다.
이제 더 큰 문제인 load에서의 5-way bank conflict가 남아 있습니다. 이 conflict는 주로 sharedB에서 load할 때 나타나며, 특히 다음 부분에서 드러납니다.
for (int col = 0; col < COLS_PER_THREAD; col += 4) {
uint global_smem_col_idx = tx * COLS_PER_THREAD + col;
float4 temp_shared_B =
reinterpret_cast<float4*>(&sharedB[i * TILE_SIZE_K + global_smem_col_idx])[0];
reg_k[col + 0] = temp_shared_B.x;
reg_k[col + 1] = temp_shared_B.y;
reg_k[col + 2] = temp_shared_B.z;
reg_k[col + 3] = temp_shared_B.w;
}
lane 0..15에서는 ty가 여전히 0이고 tx는 0..15를 걷습니다. 단순화를 위해 col = 0을 고정하면 각 lane의 float4 첫 번째 word에 대한 bank는 다음과 같습니다.
bank = (i* 128 + 8 * tx) % 32 = (8 * tx) % 32여기서 중요한 점은 우리는 벡터화된 float4 load를 하고 있으므로, 각 lane은 연속된 네 개 bank를 가로지른다는 것입니다. bank 시작이 0인 lane은 {0,1,2,3}, 8인 lane은 {8,9,10,11}, 16인 lane은 {16,17,18,19}, 24인 lane은 {24,25,26,27}을 건드립니다.
패턴이 네 lane마다 반복되므로, 동시에 {0..3}을 원하는 lane이 네 개, {8..11}을 원하는 lane이 네 개, 이런 식으로 겹치게 됩니다. 바로 이 때문에 이 명령에서 4-way conflict가 발생합니다.
sharedA load는 다릅니다. half warp 안에서 lane에 따라 달라지는 것은 tx지만, 주소에는 tx가 나타나지 않습니다. 하나의 half warp 내부에서는 ty가 상수입니다. 고정된 i와 row에 대해 모든 lane은 같은 주소를 계산합니다. 따라서 half warp의 16개 lane 모두가 그 단계에서 sharedA의 같은 네 word를 읽습니다. 앞에서 말했듯 이것은 broadcast될 수 있으므로 load 측면에서는 conflict-free입니다.
여기서 중요한 점은 padding이 이 load conflict를 해결하지 못한다는 것입니다. padding은 주소에서 변화하는 부분이 32 word의 배수 stride와 곱해질 때 도움이 됩니다. 그런데 위 sharedB load에서 변화하는 부분은 tx * COLS_PER_THREAD + col이고, 이 부분은 padded stride와 곱해지지 않습니다. 따라서 STRIDE_B = 132로 바꾸더라도 half warp 안의 lane들은 여전히 같은 네 bank 그룹에 몰립니다. 즉 padding은 store 측 문제는 해결했지만, sharedB load conflict는 다른 접근이 필요합니다.
padding을 적용한 벡터화 2D 레지스터 타일링 커널의 최종 코드는 다음과 같습니다.
template <const uint TILE_SIZE_M, const uint TILE_SIZE_N, const uint TILE_SIZE_K, const uint ROWS_PER_THREAD, const uint COLS_PER_THREAD>
__global__ void sgemm_vectorised(const float *__restrict__ A, const float *__restrict__ B, float *__restrict__ C,
int M, int N, int K, float alpha, float beta)
{
// Allocate shared memory. Use padded leading strides that keep float4 alignment
constexpr uint STRIDE_A = (TILE_SIZE_M % 32u == 0u) ? (TILE_SIZE_M + 4u) : TILE_SIZE_M;
constexpr uint STRIDE_B = (TILE_SIZE_K % 32u == 0u) ? (TILE_SIZE_K + 4u) : TILE_SIZE_K;
static_assert((STRIDE_A % 4u) == 0u, "STRIDE_A must keep float4 alignment");
static_assert((STRIDE_B % 4u) == 0u, "STRIDE_B must keep float4 alignment");
// Allocate shared memory
__shared__ float sharedA[STRIDE_A * TILE_SIZE_N];
__shared__ float sharedB[TILE_SIZE_N * STRIDE_B];
// Identify the tile of C this thread block is responsible for
const uint block_row = blockIdx.y;
const uint block_column = blockIdx.x;
// Calculate position of thread within tile (Remapping from 1-D to 2-D)
const uint ty = threadIdx.x / (TILE_SIZE_K / COLS_PER_THREAD);
const uint tx = threadIdx.x % (TILE_SIZE_K / COLS_PER_THREAD);
// Move pointers from A, B, C to tile starts
A += block_row * TILE_SIZE_M * N;
B += block_column * TILE_SIZE_K;
C += (block_row * TILE_SIZE_M * K) + (block_column * TILE_SIZE_K);
// Map each thread to one 4-float chunk
const uint smem_ty_A = threadIdx.x / (TILE_SIZE_N / 4);
const uint smem_tx_A = threadIdx.x % (TILE_SIZE_N / 4);
const uint smem_ty_B = threadIdx.x / (TILE_SIZE_K / 4);
const uint smem_tx_B = threadIdx.x % (TILE_SIZE_K / 4);
// Tile count
const uint num_tiles = CEIL_DIV(N, TILE_SIZE_N);
float thread_results[ROWS_PER_THREAD * COLS_PER_THREAD] = {0.0f};
float reg_m[ROWS_PER_THREAD] = {0.0f};
float reg_k[COLS_PER_THREAD] = {0.0f};
// Outer loop iterate over tiles
for (int t = 0; t < num_tiles; t++)
{
// Populate smem using vector loads
float4 tempA = reinterpret_cast<const float4 *>(&A[smem_ty_A * N + smem_tx_A * 4])[0];
sharedA[(smem_tx_A * 4 + 0) * STRIDE_A + smem_ty_A] = tempA.x;
sharedA[(smem_tx_A * 4 + 1) * STRIDE_A + smem_ty_A] = tempA.y;
sharedA[(smem_tx_A * 4 + 2) * STRIDE_A + smem_ty_A] = tempA.z;
sharedA[(smem_tx_A * 4 + 3) * STRIDE_A + smem_ty_A] = tempA.w;
float4 tempB = reinterpret_cast<const float4 *>(&B[smem_ty_B * K + smem_tx_B * 4])[0];
reinterpret_cast<float4 *>(&sharedB[smem_ty_B * STRIDE_B + smem_tx_B * 4])[0] = tempB;
__syncthreads();
// Outer loop over shared dimension N
for (int i = 0; i < TILE_SIZE_N; i++)
{
// Load regs from sharedA
for (int row = 0; row < ROWS_PER_THREAD; row += 4)
{
uint global_smem_row_idx = ty * ROWS_PER_THREAD + row;
float4 temp_shared_A = reinterpret_cast<float4 *>(&sharedA[i * STRIDE_A + global_smem_row_idx])[0];
reg_m[row + 0] = temp_shared_A.x;
reg_m[row + 1] = temp_shared_A.y;
reg_m[row + 2] = temp_shared_A.z;
reg_m[row + 3] = temp_shared_A.w;
}
// Load regs from sharedB
for (int col = 0; col < COLS_PER_THREAD; col += 4)
{
uint global_smem_col_idx = tx * COLS_PER_THREAD + col;
float4 temp_shared_B = reinterpret_cast<float4 *>(&sharedB[i * STRIDE_B + global_smem_col_idx])[0];
reg_k[col + 0] = temp_shared_B.x;
reg_k[col + 1] = temp_shared_B.y;
reg_k[col + 2] = temp_shared_B.z;
reg_k[col + 3] = temp_shared_B.w;
}
// Outer product
for (uint m = 0; m < ROWS_PER_THREAD; m++)
for (uint k = 0; k < COLS_PER_THREAD; k++)
thread_results[m * COLS_PER_THREAD + k] += reg_m[m] * reg_k[k];
}
__syncthreads();
A += TILE_SIZE_N;
B += TILE_SIZE_N * K;
}
// Write results back
for (uint row = 0; row < ROWS_PER_THREAD; row++)
for (uint col = 0; col < COLS_PER_THREAD; col += 4)
{
uint global_row_idx = ty * ROWS_PER_THREAD + row;
uint global_col_idx = tx * COLS_PER_THREAD + col;
float4 tempC = reinterpret_cast<float4 *>(&C[global_row_idx * K + global_col_idx])[0];
tempC.x = (alpha * thread_results[row * COLS_PER_THREAD + col]) + (beta * tempC.x);
tempC.y = (alpha * thread_results[row * COLS_PER_THREAD + col + 1]) + (beta * tempC.y);
tempC.z = (alpha * thread_results[row * COLS_PER_THREAD + col + 2]) + (beta * tempC.z);
tempC.w = (alpha * thread_results[row * COLS_PER_THREAD + col + 3]) + (beta * tempC.w);
reinterpret_cast<float4 *>(&C[global_row_idx * K + global_col_idx])[0] = tempC;
}
}
지금까지 우리는 두 단계의 병렬성을 활용했습니다.
(ROWS_PER_THREAD × COLS_PER_THREAD)을 전부 레지스터 안에서 계산해, 결과를 global memory로 다시 쓰기 전에 데이터 재사용을 극대화했습니다.이번 커널에서는 블록 타일링과 스레드 타일링 사이에 새로운 타일링 단계를 도입합니다. 그것이 바로 warp 타일링입니다.
warp 타일링은 최적화 계층에서 block 타일링과 thread 타일링 사이에 위치합니다. block 안의 모든 스레드가 하나의 큰 타일에 협력하는 대신, 그 타일을 더 작은 서브타일로 나누고 각 서브타일을 하나의 warp에 할당합니다. 이렇게 하면 warp가 중간 수준의 계산 단위가 됩니다. block은 여전히 C의 128 × 128 패치를 덮지만, 이것을 네 개의 64 × 64 서브타일로 나눕니다. M 방향으로 warp 두 개, K 방향으로 warp 두 개여서 block당 총 네 개 warp가 됩니다.
TILE_SIZE_M = 128
TILE_SIZE_N = 16
TILE_SIZE_K = 128
WARP_TILE_M = 64
WARP_TILE_K = 64
WARP_STEPS_K = 4
ROWS_PER_THREAD = 8
COLS_PER_THREAD = 4
NUM_THREADS = 128 // four warps per block
block은 여전히 이전 커널에서 사용한 벡터화, padding, 전치 기법을 이용해 GMEM에서 SMEM으로 데이터를 협력해서 불러옵니다. 데이터가 온칩에 올라오면 스레드들은 네 개 warp로 나뉘고, 각 warp는 warp_row와 warp_col로 식별되는 출력 행렬의 한 사분면에 대한 독점적인 소유권을 가집니다.
warp 내부의 32개 스레드는 자신들의 세부 서브 인덱스(ty, tx, 8×4 스레드 서브그리드에서 유도됨)를 사용해 자신에게 할당된 64×64 영역을 공략합니다. warp는 수직 차원은 한 번의 패스로 덮지만 WARP_STEPS_M = 1, 수평 방향은 WARP_STEPS_K=4 번 반복해야 합니다(물론 설정 가능하지만 신중해야 합니다!). 부분 결과 누적은 TILE_SIZE_N 공유 차원을 순회하는 루프(i 루프) 안에서 일어납니다. 계산 루프에서 이 설계는 스레드가 sharedA에서 reg_m 조각을 한 번 load한 뒤 네 번의 수평 단계 동안 재사용하게 해주며, 각 단계마다 sharedB에서 새로운 reg_k 데이터를 불러와 outer product(reg_m과 reg_k 사이)의 결과를 큰 thread_results 배열에 누적합니다. 이렇게 해서 모든 i 반복에 걸쳐 최종 결과가 축적됩니다.
먼저 커널의 고수준 구조를 보여주며 warp 타일링이 어떻게 새로운 계층으로 통합되는지 설명하겠습니다. 그다음에는 몇 가지 더미 파라미터를 사용해 단일 스레드의 시점에서 그 스레드의 전체 생애를 시각화할 것입니다. 즉, 계산이 어디서 일어나고, load와 store를 정확히 어디서 수행하는지 볼 것입니다.
이제 단일 스레드의 관점에서 계산 흐름을 시각화해봅시다.
이 추가 타일링 계층은 여러 이점을 제공합니다.
warp는 NVIDIA GPU에서 기본 실행 단위입니다. 각 warp에 출력의 자체 서브타일을 맡기면, 작업 분할 방식이 하드웨어가 실제로 명령을 스케줄링하는 방식과 맞아떨어집니다.
이렇게 하면 각 warp가 독립적으로 실행될 수 있습니다. 한 warp가 메모리에서 stall되더라도 다른 warp들이 계속 실행될 수 있어서 warp 스케줄러 슬롯이 가득 차고 idle cycle이 줄어듭니다.
Simon의 블로그에서
warp 타일은 각 warp의 footprint를 작게 유지하고 lane당 stride 패턴을 단순하고 반복 가능하게 유지합니다. 그래서 bank 친화적인 레이아웃을 설계하기 쉬워집니다. 스포일러를 하자면, 바로 그래서 이 커널에서는 SMEM load conflict가 나타나지 않았습니다.
각 Streaming Multiprocessor 안의 register file(RF)은 스레드별 변수를 저장합니다. Hopper에서는 이것이 여러 개의 단일 포트 bank로 나뉘어 있습니다(shared memory bank와 비슷합니다!). 한 bank는 사이클당 한 번만 접근을 처리할 수 있습니다. 같은 warp의 두 스레드가 같은 사이클에 같은 bank를 읽으려고 하면 접근이 직렬화됩니다. 이것도 bank conflict이지만 레지스터에서 발생하는 경우이며, 명령의 피연산자를 가져오는 시간을 늘립니다. 불행히도 NVIDIA의 프로파일링 도구는 이런 conflict에 대한 지표를 제공하지 않기 때문에, 이 커널에서 실제로 이것이 개선되었는지는 확인하기 어렵습니다.
RF와 실행 유닛 사이에는 Operand Collector Unit(OCU)이 있습니다. 논문: BOW Breathing Operand Windows to Exploit Bypassing in GPUs. 각 OCU는 register bank에서 source operand를 가져와 작은 버퍼에 저장하며, 128바이트 엔트리 세 개를 저장할 수 있습니다. 피연산자가 곧 다시 필요하면 메인 RF로 돌아가지 않고 이 버퍼에서 바로 제공될 수 있습니다. 이렇게 하면 bank conflict와 추가 RF 트래픽을 모두 피할 수 있습니다.
warp 타일링은 각 warp가 출력 행렬의 작고 고정된 서브타일을 다루기 때문에 도움이 됩니다. 따라서 inner loop에서 같은 레지스터를 반복적으로 재사용하는 경향이 있습니다. 이로 인해 bank conflict가 줄어들 가능성이 있고, 피연산자가 OCU 버퍼에서 직접 재사용될 가능성도 높아집니다.
다시 말하지만 이것은 추측입니다. 실제로 차이가 있는지는 저도 확신하지 못하지만, 그럴듯해 보입니다.
주요하게 바뀐 코드 부분은 다음과 같습니다.
// Iterate over the shared dimension of the SMEM tiles
for (int i = 0; i < TILE_SIZE_N; i++)
{
// Load slice at current i iteration in sharedA's register
for (int wSubRow = 0; wSubRow < WARP_STEPS_M; wSubRow++)
{
uint base_row =
(warp_row * WARP_TILE_M) +
(wSubRow * WARP_SUB_M) +
(ty * ROWS_PER_THREAD);
// Each thread loads ROWS_PER_THREAD into the register
#pragma unroll
for (int row = 0; row < ROWS_PER_THREAD; row += 4)
{
const float4 va =
reinterpret_cast<const float4*>(
&sharedA[i * STRIDE_A + base_row + row])[0];
reg_m[wSubRow * ROWS_PER_THREAD + row + 0] = va.x;
reg_m[wSubRow * ROWS_PER_THREAD + row + 1] = va.y;
reg_m[wSubRow * ROWS_PER_THREAD + row + 2] = va.z;
reg_m[wSubRow * ROWS_PER_THREAD + row + 3] = va.w;
}
for (int wSubCol = 0; wSubCol < WARP_STEPS_K; wSubCol++)
{
uint col_base =
(warp_col * WARP_TILE_K) +
(wSubCol * WARP_SUB_K) +
(tx * COLS_PER_THREAD);
// Each thread loads COLS_PER_THREAD into the register x 4 times in our case since WARP_STEPS_K = 4
#pragma unroll
for (int col = 0; col < COLS_PER_THREAD; col += 4)
{
const float4 vb =
reinterpret_cast<const float4*>(
&sharedB[i * STRIDE_B + col_base + col])[0];
reg_k[wSubCol * COLS_PER_THREAD + col + 0] = vb.x;
reg_k[wSubCol * COLS_PER_THREAD + col + 1] = vb.y;
reg_k[wSubCol * COLS_PER_THREAD + col + 2] = vb.z;
reg_k[wSubCol * COLS_PER_THREAD + col + 3] = vb.w;
}
}
// Compute outer product
for (int wSubRow = 0; wSubRow < WARP_STEPS_M; wSubRow++)
{
for (int wSubCol = 0; wSubCol < WARP_STEPS_K; wSubCol++)
{
#pragma unroll
for (int im = 0; im < ROWS_PER_THREAD; im++)
{
float fixed_temp =
reg_m[wSubRow * ROWS_PER_THREAD + im];
#pragma unroll
for (int ik = 0; ik < COLS_PER_THREAD; ik++)
{
float out =
fixed_temp * reg_k[wSubCol * COLS_PER_THREAD + ik];
int out_idx =
(wSubRow * ROWS_PER_THREAD + im) *
(WARP_STEPS_K * COLS_PER_THREAD) +
(wSubCol * COLS_PER_THREAD + ik);
thread_results[out_idx] += out;
}
}
}
}
}
}
__syncthreads();
A += TILE_SIZE_N; // Move right
B += TILE_SIZE_N * K; // Move down
저는 이 커널을 padding 전후로 테스트했습니다.
정리하자면, 이 warp-tiling 커널에서 padding은 주로 store 경로(sharedA로 전치해 쓰는 부분)를 도와주었고, 실제로 store-conflict 카운터가 약 4.0에서 약 2.5-way로 떨어진 것과 일치합니다. 이전 벡터화 커널과 달리 여기서는 load conflict가 문제가 아니었습니다. load 쪽에서는 별다른 복잡한 처리를 하지 않았는데도 두 가지가 조용히 도움을 주었습니다. COLS_PER_THREAD = 4 를 사용해 sharedB lane들을 더 많은 bank 그룹에 분산시켰고, warp 지역 서브타일 덕분에 lane 패턴이 덜 겹치게 되었습니다. 이 둘이 합쳐져 padding 전후 실행 모두에서 프로파일러가 shared-load conflict를 경고하지 않은 이유입니다.
우리를 여전히 막고 있는 것은 다른 곳입니다. 레지스터 압박 때문에 achieved occupancy가 약 18%에 머물고 있고, 이것이 “not selected” 스케줄러 stall로 나타납니다. 그리고 여전히 대체로 compute-bound 상태입니다(FMA busy가 60%대 후반). memory는 약 52%이므로, 몇 GB/s를 더 짜내는 것보다 copy와 compute를 겹치게 하거나 레지스터를 줄여 상주 warp 하나를 더 확보하는 편이 더 효과적입니다.
📝 중요한 메모: 이 커널부터는 차원 표기를 바꿔
A = MxK,B=KxN으로 씁니다. 이유는 이후 tensor core 명령이 이런 형식의 행렬을 기대하기 때문입니다. 이전 커널들의 논리는 모두 동일하고 단지 이름만 바뀌는 것입니다. 향후에는 일관성을 위해 위의 모든 코드와 그림도 바꿀 예정입니다.
서론에서 H100의 Tensor Core 구성 요소를 잠깐 언급했습니다. 이제 이것들이 어떻게 동작하고, 어떻게 활용하면 성능을 크게 끌어올릴 수 있는지 자세히 봅시다. 이 커널이 끝날 때쯤이면 성능이 급상승할 것이므로 기대해봅시다.
NVIDIA의 최근 GPU 아키텍처에서 가장 중요한 발전 중 하나는 Tensor Core의 도입과 진화입니다. 진지한 병렬 계산을 위해 고급 GPU를 구매하는 주된 이유이기도 합니다. 이것들은 항상 존재했던 것은 아니고, Volta 아키텍처(V100)에서 처음 도입되었습니다. 즉 CUDA Core용으로 최적화해온 이전 모든 커널은 Volta 이전 아키텍처에서의 최첨단이라고 볼 수 있습니다.
Tensor Core는 GPU의 계산 모델을 근본적으로 바꿉니다. 이들은 행렬 곱셈 및 누산(MMA) 을 가속하기 위해 특별히 설계된 엔진입니다.
CUDA Core가 a @ b + c 같은 단순한 스칼라 명령을 실행하는 것과 달리, Tensor Core는 D = A @ B + C 같은 전체 행렬 연산을 수행하는 단일 명령을 실행합니다. 이런 구조는 종종 CISC(Complex Instruction Set Computer)와 비교됩니다. CISC 명령 하나는 메모리에서 값 로드, 산술 계산 수행, 결과 다시 저장처럼 여러 저수준 작업을 한 번에 처리할 수 있습니다. 반면 RISC 아키텍처는 한 번에 하나의 기본 연산만 수행하는 매우 단순한 명령을 사용합니다.
전력 밀도: 이런 CISC 유사 접근은 엄청난 속도의 핵심입니다. 명령 하나당 큰 데이터 블록을 처리하면 명령 해독 같은 작업당 오버헤드가 극적으로 줄어듭니다.
예를 들어:
이후에 사용할 Warp Group Matrix Multiply and Accumulate (WGMMA) 명령은 wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 처럼 쓰입니다. 여기서 m64n64k16이 행렬 차원을 나타냅니다. 바깥 차원은 m과 n이며 앞과 뒤에 오고, 누산을 위한 공유 내부 차원 k는 가운데에 있습니다. 이 복합 명령은 행렬 A, B, 그리고 누산기 C에 대해 D = A @ B + C를 계산합니다(C는 종종 물리적으로 D와 같은 행렬입니다). 이를 곱해보면 이 명령은 64 * 16 * 64 = 65,536개의 multiply-accumulate(MAC) 연산을 수행합니다.
이를 CUDA Core만 사용한 warp 타일링 커널과 비교해봅시다.
float out = fixed_temp * reg_k[...]; // multiplication
thread_results[out_idx] += out; // addition (accumulation)
여기서는 명령당 1 MAC(FMA)만 했습니다. 같은 작업량을 완료하려면 CUDA Core는 65,536개의 FMA 명령을 실행해야 합니다. 반면 WGMMA에서는 하나의 warp group(128 threads)이 WGMMA 명령 하나로 65,536 MAC 전체를 수행합니다.
여기서 잠깐 물러서서 방금 도입한 두 개념, WGMMA와 warp-group을 짚고 넘어가야 합니다. 이 개념들은 Hopper 아키텍처 고유의 것이며 이전 GPU에는 없었습니다. 이것들이 왜 중요한지, 그리고 왜 이 커널에서 WGMMA를 사용할 것인지 이해하려면, Hopper 이전에 tensor core가 어떻게 프로그래밍되었고 그 프로그래밍 모델이 어떻게 진화해왔는지 간단히 보는 것이 도움이 됩니다.
Hopper 이전에는 보통 WMMA API(Warp Matrix Multiply Accumulate)를 통해 tensor core를 프로그래밍했습니다. 이 인터페이스는 Volta에서 도입되어 Turing과 Ampere까지 nvcuda::wmma 형태로 이어졌습니다. tensor core를 활용하기 위한 고수준 추상화를 제공했고, API가 내부 세부 사항 대부분을 처리했습니다.
이후 NVIDIA는 하부 tensor core 명령을 직접 노출했습니다. 이로써 Turing과 Ampere에서 MMA PTX 명령이 등장했습니다. 이들도 warp 수준에서 동작하며, 32개 스레드가 협력해 더 작은 MMA 연산을 수행합니다. 이 명령에 올바르게 데이터를 공급하기 위해 아키텍처에는 ldmatrix라는 특수한 warp-wide load 명령이 추가되었고, shared memory에서 필요한 packed fragment를 레지스터로 가져옵니다. 이 단계에서 전형적인 tensor core 커널은 각 warp 안에서 분명한 패턴을 가졌습니다. ldmatrix로 A와 B fragment를 SMEM에서 load하고, 하나 이상의 mma.sync 명령을 발행한 뒤, 누적된 결과를 써내는 방식입니다.
Hopper에서는 tensor 연산이 warp 수준에서 warp group 수준으로 확장되며, 128 (32*4) 개의 스레드가 하나의 훨씬 더 큰 MMA에 협력합니다. 이 명령들은 더 이상 작은 warp 크기 타일에서 동작하지 않고, 이미 shared memory에 존재해야 하는 훨씬 큰 A와 B 블록을 대상으로 합니다. 그리고 WGMMA는 이 타일들이 특정한 swizzled 레이아웃으로 배치되어 있기를 기대합니다. 그러면 자연스럽게 질문이 생깁니다.
WGMMA가 기대하는 정확한 형식으로 필요한 행렬 타일을 shared memory에 어떻게 배치하고, tensor core가 놀지 않도록 충분히 빠르게 그것을 수행할 수 있을까?
바로 여기서 서론에서 이야기한 Tensor Memory Accelerator (TMA) 가 등장합니다.
H100에서 TMA는 전용 병렬 copy 엔진으로 동작해 데이터 병목을 해결합니다.
대량 적재: 단 하나의 하드웨어 명령으로 A와 B의 전체 2D 타일을 GMEM에서 SMEM으로 옮깁니다.
비동기 전송: 결정적으로 이 전송은 백그라운드에서 실행됩니다. 따라서 Tensor Core가 현재 데이터를 처리하는 동안 TMA는 이미 다음 반복에 필요한 다음 2D 타일을 가져오고 있을 수 있습니다.
“대량 적재”라는 말은 예전에는 프로그래머의 책임이었던 많은 복잡성을 숨기고 있습니다. 이전 커널들에서는 각 스레드가 정확히 어떤 원소를 가져와야 하는지 일일이 코드로 써야 했습니다. 이제는 그런 수동 인덱싱이 하드웨어로 넘어갑니다. 더 이상 GMEM에서 SMEM으로 특정 원소를 로드하도록 스레드를 세세하게 지휘할 필요가 없습니다.
또한 이전처럼 SMEM bank conflict를 피하기 위해 수동으로 padding을 넣을 필요도 없습니다. 하드웨어가 Swizzling이라고 알려진 레이아웃 변환을 자동으로 적용합니다. 이 레이아웃을 손으로 코딩하는 것은 매우 복잡하기 때문에, 다행히도 NVIDIA가 이런 패턴을 TMA에 직접 구현해두었습니다. 우리가 알아야 할 것은 bank conflict가 essentially 공짜처럼 처리된다는 점 정도입니다. swizzling 패턴의 구체적 원리가 궁금하다면 Aleksa의 글이 매우 깊이 다루고 있습니다. 아래는 그가 swizzled copy into SMEM이 고수준에서 어떻게 보이는지 설명하기 위해 사용한 그림입니다.
TMA를 사용하려면 세 가지 주요 단계가 필요합니다.
A와 B에 대한 tensor map을 생성합니다(호스트에서).Tensor Map은 하드웨어가 해석할 수 있는 descriptor입니다. 메모리 안 텐서의 shape, layout, stride를 설명해주며, 이를 통해 TMA가 thread 수준 주소 계산 없이 전체 다차원 타일을 이동할 수 있게 합니다.
이전 커널들과 달리, 커널 인자로 const float* A 같은 raw pointer를 넘기지 않습니다. 대신 CUtensorMap descriptor에 대한 포인터를 넘깁니다. 이것은 CUDA Driver가 정의한 구조체로, 행렬의 shape, stride, swizzle pattern 같은 전체 메타데이터를 담고 있어 하드웨어가 직접 타일을 가져올 수 있게 합니다.
이 map을 만들기 위해 CUDA Driver API의 cuTensorMapEncodeTiled 함수를 사용합니다.
이것들은 CUDA 소프트웨어 플랫폼을 구성하는 여러 구성 요소들이며, 어떤 것은 호스트에서 호출해야 하고 어떤 것은 디바이스에서 사용할 수 있는지 이해하는 데 중요합니다.
스택의 가장 아래에는 CUDA Driver API가 있습니다. 가장 세밀한 GPU 제어를 제공하지만 더 장황하고 복잡합니다. tensor map 생성 같은 저수준 기능은 이 계층에서 명시적으로 노출됩니다.
Driver API 위에는 CUDA Runtime API가 있으며, 많은 기능을 감싸 더 고수준 인터페이스를 제공합니다. 예를 들어 runtime API의 cudaMalloc은 driver API의 cuMemAlloc을 감싼 얇은 래퍼입니다.
그리고 두 API 위에는 cuBLAS 같은 선형대수 라이브러리, cuDNN 같은 딥러닝 라이브러리처럼 일반 및 도메인 특화 워크로드를 위한 고도로 최적화된 CUDA 라이브러리들이 있습니다. 실제로 대부분의 코드는 runtime API를 사용하지만, TMA tensor map 같은 일부 Hopper 기능은 현재 Driver API를 통해서만 노출됩니다.
. 이 함수는 우리가 기술한 행렬 정보를 받아 TMA 엔진이 이해할 수 있는 128B 하드웨어 descriptor로 패킹합니다. TMA 하드웨어가 매우 특수화되어 있기 때문에, 이 128B 객체는 메모리에서 128B 경계에 정렬되어 있어야 하며 그렇지 않으면 하드웨어가 읽을 수조차 없습니다.
과정은 다음과 같습니다.
cudaMalloc을 사용해 디바이스에 tensor map용 메모리를 할당합니다.cudaMemcpy로 그 map을 호스트에서 디바이스로 복사합니다.아래 코드 조각은 이 단계를 다음과 같이 수행합니다.
template <const uint BlockMajorSize, const uint BlockMinorSize>
__host__ static inline CUtensorMap *
create_and_allocate_tensor_map(bf16 *tensor_ptr, uint blocks_height, uint blocks_width) {
CUtensorMap *tensor_map;
// Allocate device memory for the tensor map descriptor.
CUDA_CHECK(cudaMalloc((void **)&tensor_map, sizeof(CUtensorMap)));
// Register the tensorMap in our device memory pointers
// resources.add_device_ptr(tensor_map);
// Create on host
CUtensorMap tensor_map_host;
create_tensor_map<BlockMajorSize, BlockMinorSize>(&tensor_map_host, tensor_ptr, blocks_height, blocks_width);
// Copy descriptor to device
CUDA_CHECK(cudaMemcpy(tensor_map, &tensor_map_host, sizeof(CUtensorMap), cudaMemcpyHostToDevice));
return tensor_map;
}
그리고 텐서의 메타데이터를 인코딩해 실제 tensor map을 만드는 함수는 다음과 같습니다.
template <const uint BlockMajorSize, const uint BlockMinorSize>
void create_tensor_map(CUtensorMap *tensor_map, bf16 *tensor_ptr, uint blocks_height, uint blocks_width) {
// Starting address of memory region described by tensor (casting to void
// as the tensor map descriptor is type-agnostic.)
void *gmem_address = static_cast<void *>(tensor_ptr);
uint num_tiles_major = blocks_height;
uint num_tiles_minor = blocks_width;
// full size of the tensor in global memory (API expects the 5D supported
// tensor ranks to be defined)
uint64_t global_dim[5] = {
static_cast<uint64_t>(BlockMinorSize * num_tiles_minor),
static_cast<uint64_t>(BlockMajorSize * num_tiles_major),
1, 1, 1};
// Define the tensor strides (in bytes) along each of the tensor ranks dims - 1
uint64_t global_strides[5] = {
sizeof(bf16),
sizeof(bf16) * BlockMinorSize * num_tiles_minor,
0, 0, 0};
// Define the shape of the "box_size" -> the tile shapes a TMA ops will load
uint32_t box_dim[5] = {
static_cast<uint32_t>(BlockMinorSize),
static_cast<uint32_t>(BlockMajorSize),
1, 1, 1};
uint32_t elem_strides[5] = {1, 1, 1, 1, 1};
// Create tensor map
CU_CHECK(cuTensorMapEncodeTiled(
tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, gmem_address,
global_dim, global_strides + 1, box_dim, elem_strides,
CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
}
다음으로 WGMMA 명령을 이야기해봅시다. WGMMA 명령은 우리가 이전 모든 커널에서 했던 것처럼 raw byte address를 직접 사용하지 않습니다. 대신 행렬이 shared memory 어디에 있고 어떻게 배치되어 있는지를 하드웨어에 알려주는 packed 64-bit matrix descriptor를 받습니다.
matrix descriptor 형식은 문서에 다음처럼 설명되어 있습니다.
원래 주소는 바이트 주소입니다. 예를 들어 16384바이트 위치의 주소는 16진수로 0x4000입니다. descriptor가 raw byte address를 직접 저장한다면 비트가 금방 모자라 주소 가능한 범위를 심하게 제한하게 됩니다.
그래서 하드웨어는 중요한 성질 하나를 이용합니다. WGMMA의 SMEM 피연산자는 항상 최소 16B 정렬입니다. 즉 어떤 유효한 주소든 하위 4비트는 항상 0이고 정보가 없습니다.
따라서 descriptor는 바이트 주소 대신 16B 단위 주소를 저장합니다. 이 인코딩은 base address뿐 아니라 leading dimension과 stride offset에도 동일하게 적용됩니다. 이렇게 하면 descriptor는 훨씬 더 큰 SMEM 영역을 표현하면서도 필요한 모든 메타데이터를 단 하나의 64비트 값 안에 담을 수 있고, 하드웨어는 이것을 warp group 전체에 효율적으로 broadcast하고 decode할 수 있습니다.
이 인코딩이 어떻게 일어나는지 개념적으로 보면 다음과 같습니다.
![Image 36: matrix-desc-encoding]](https://bear-images.sfo2.cdn.digitaloceanspaces.com/testing-inf1/excalidraw-52.svg)
인코딩 로직과 descriptor 생성은 WGMMA matrix descriptor가 커널 내부에서 만들어지므로 디바이스에서 실행되어야 합니다. 우리는 make_smem_descriptor 함수에서 matrix descriptor를 구성합니다.
또한 인코딩 전에 generic pointer bf16*를 __cvta_generic_to_shared를 사용해 shared-memory address로 먼저 변환해야 합니다. 이 단계는 미묘하지만 필수적입니다. WGMMA는 C++ 추상화가 아니라, 매우 특정한 형식의 SMEM byte address를 기대하는 저수준 하드웨어 명령입니다. 일반 CUDA C++ pointer는 자신의 주소 공간(Global, Shared 등)을 명시적으로 인코딩하지 않으므로, 그것을 직접 사용하거나 인코딩할 수 없습니다. 따라서 먼저 포인터를 하드웨어가 이해하는 구체적인 SMEM 주소로 바꾼 뒤에야, 그것을 압축하고 matrix descriptor에 패킹해야 합니다. CUDA C++ 포인터는 generic하기 때문에 global memory 객체도, shared memory 객체도 가리킬 수 있지만 주소 공간을 명시적으로 담지 않습니다. 이런 추상화는 일반 C++ load/store에는 잘 맞지만, PTX 명령이나 WGMMA 같은 하드웨어 인터페이스와 직접 상호작용할 때는 깨집니다. 이런 인터페이스는 특정 메모리 공간을 명시하는 주소를 요구합니다. 이를 연결하기 위해 CUDA는 __cvta_generic_to_shared 같은 address-space conversion intrinsic을 제공하며, 이것은 generic pointer를 PTX 명령과 하드웨어 descriptor가 소비할 수 있는 shared-memory address로 바꿔줍니다. 두 함수의 코드는 다음과 같습니다.
__device__ static inline uint64_t matrix_descriptor_encode(uint64_t x) {
return ((x) & 0x3FFFF) >> 4;
}
__device__ uint64_t make_smem_desc(bf16* ptr) {
uint32_t address = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
// Initialise an empty 64 bit descriptor
uint64_t desc = 0x0000000000000000;
// bitwise OR
// sets bits [13:0] encoded matrix start address
desc |= matrix_descriptor_encode(address);
// sets bits [29:16] leading dimension byte offset
desc |= matrix_descriptor_encode(static_cast<uint64_t>(16)) << 16;
// sets bits [45: 32] stride dimension byte offset
desc |= matrix_descriptor_encode(static_cast<uint64_t>(1024)) << 32;
// sets bits [62: 63] swizzle mode
desc |= 1llu << 62;
return desc;
}
make_smem_desc 함수에서는 먼저 비어 있는 64비트 descriptor를 초기화하고, 이를 위 matrix descriptor 레이아웃에 따라 필드별로 채워 반환합니다.
먼저 matrix start address를 인코딩해 descriptor의 [13:0] 비트에 넣습니다. 이것은 shared memory 안에서 행렬의 base address를 나타내며 16바이트 단위로 인코딩됩니다. 다음으로 leading dimension byte offset을 인코딩해 [29:16] 비트에 배치합니다. 이것은 하드웨어가 행렬의 leading dimension을 따라 이동하기 위해 몇 바이트를 건너뛰어야 하는지를 나타냅니다.
세 번째 필드는 stride dimension byte offset이고, [45:32] 비트에 놓입니다. 문서 설명에 따르면 다음과 같습니다.
개념적으로 이 필드는 다음 질문에 답합니다.
“K 차원을 따라 columns 0–7에서 columns 8–15로 가려면 몇 바이트를 이동해야 하는가?”
아직 정확한 타일 모양은 소개하지 않았지만, 스포일러를 하자면 sharedA는 64 × 64 행렬이 됩니다. 각 열에는 64개의 bf16 원소가 있고, bf16 하나는 2바이트이므로 열 하나는 128바이트 폭입니다. K 방향으로 8개 열을 건너뛰려면 따라서 8 × 128 bytes = 1024 bytes 가 필요합니다. 그래서 stride dimension byte offset은 1024로 인코딩합니다.
마지막으로 마지막 필드는 swizzling mode를 지정하며 [63:62] 비트에 저장됩니다. 우리는 128-byte swizzling을 사용하므로 이 필드를 1로 설정합니다. 이 필드는 byte offset이 아니라 고정된 열거형 값이므로 인코딩 함수를 거치지 않습니다.
이제 본론으로 들어가 커널 자체를 이야기하며 지금까지의 조각들을 하나로 연결해봅시다. 다만 한 가지 주의할 점은, 지금은 WGMMA 명령을 어떻게 정의하는지 세부까지는 들어가지 않겠다는 것입니다. 이야기 흐름을 위해 일단은 블랙박스로 두고, 뒤에서 자세히 돌아오겠습니다.
커널 코드를 보기 전에 먼저 전체 커널 흐름을 고수준에서 정리하고, 그것이 시각적으로 어떻게 보이는지를 보여드리고 싶습니다.
A와 B의 전체 2D 타일을 GMEM에서 SMEM으로 불러옵니다.m64n64k16 WGMMA 명령 네 개를 발행하는 것을 의미합니다.C를 만듭니다.
보시다시피 각 tensor core MMA 동안 우리는 sharedA의 서브타일과 sharedB의 서브타일의 outer product를 취해 TILE_SIZE_M(64) × TILE_SIZE_N(64) 크기의 accumulator 타일을 갱신합니다.
피연산자 배치 규칙은 다음과 같습니다.
sharedA는 레지스터 또는 shared memory에 있을 수 있습니다.sharedB는 반드시 shared memory에 있어야 합니다.D는 반드시 레지스터에 있어야 합니다(그림의 C's tile에 해당).물론 단일 스레드가 전체 64 × 64 accumulator를 들고 있을 수는 없습니다. 대신 accumulator는 warp group 전체에 분산됩니다. 여기서는 warp group이 128개 스레드로 구성되며, 이 128개 전체가 각 m64n64k16 WGMMA 명령을 협력해 실행합니다. 한 번의 K 반복 안에서 warp group은 여러 WGMMA 명령을 연달아 발행하며, 각각 K 차원의 서로 다른 slice(2.1, 2.2 등)에 대해 동작하면서 같은 스레드별 accumulator 레지스터에 누적합니다. 이 명령들이 함께 전체 64 × 64 출력 타일을 만들어냅니다.
이 레이아웃은 우리가 수동으로 정하는 것이 아닙니다. 하드웨어에 의해 고정되어 있으며 PTX 문서의 9.7.15.5.1.1. 절에서 설명됩니다. 여기에는 m64n64k16의 register fragment layout과 accumulator 타일이 warp group 전체에 어떻게 분배되는지가 정의되어 있습니다.
코드에서는 이것이 스레드별 accumulator 레지스터를 초기화하는 부분으로 나타납니다.
// Initialise thread's accumilator
// d[4][8] = 32 floats per thread
float d[WGMMA_N / 16][8];
memset(d, 0, sizeof(d));
각 스레드는 정확히 32개의 부동소수점 accumulator 값을 소유하고, 이것들이 함께 전체 64 × 64 출력 타일 중 그 스레드가 담당하는 fragment를 이룹니다. 각 WGMMA 단계에서 이 레지스터들은 제자리에서 갱신됩니다. 다음으로 sharedA와 sharedB에 대한 쓰기를 동기화하기 위해 두 개의 SMEM barrier를 만들고 초기화합니다. 두 barrier는 모두 warp group의 128개 스레드를 기준으로 초기화됩니다.
// SMEM barriers for A and B
__shared__ barrier barA;
__shared__ barrier barB;
if (threadIdx.x == 0) {
init(&barA, blockDim.x);
init(&barB, blockDim.x);
cde::fence_proxy_async_shared_cta();
}
__syncthreads();
여기서 cuda::barrier<cuda::thread_scope_block> API(barrier는 alias로 사용) 를 쓰고 있음에도, 여전히 __syncthreads()를 한 번 발행하는 점에 주목하세요. 이유는 일종의 부트스트래핑 문제 때문입니다.
문서에 적혀 있듯, 스레드들이 이 barrier를 이용해 동기화를 시작하기 전에 barrier 자체를 커널 내부에서 먼저 초기화해야 합니다. 그런데 그러면 애초에 barrier를 초기화하기 위해 스레드들은 무엇으로 동기화할까요?
바로 이 때문에 __syncthreads()를 아직 한 번 씁니다. barrier가 초기화된 뒤 모든 스레드가 동기화되도록 하기 위해 단 한 번만 예전 방식의 동기화 primitive를 사용하는 것입니다(__syncthreads()는 초기화가 필요 없다는 것을 이전 커널에서 이미 봤습니다). 이 일회성 부트스트랩이 끝나면, 나머지 커널 동안은 cuda::barrier를 안전하게 정상 사용하면 됩니다.
이후에는 이전 모든 커널과 마찬가지로 num_blocks_k(공유 차원)를 따라 도는 바깥 루프를 시작하고, TMA를 이용한 bulk load를 발행하기 시작합니다.
barrier::arrival_token tokenA, tokenB;
for (int block_k_iter = 0; block_k_iter < num_blocks_k; block_k_iter++) {
// Async loads (Only 1 thread launches the TMA op)
if (threadIdx.x == 0) {
// Thread 0 launches async bulk tensor copy operations for both matrices
cde::cp_async_bulk_tensor_2d_global_to_shared(&sharedA[0], tensorMapA, block_k_iter * TILE_SIZE_K, num_block_m * TILE_SIZE_M, barA);
// Signal barrier and wait for both loads to complete
tokenA = cuda::device::barrier_arrive_tx(barA, 1, sizeof(sharedA));
cde::cp_async_bulk_tensor_2d_global_to_shared(&sharedB[0], tensorMapB, block_k_iter * TILE_SIZE_K, num_block_n * TILE_SIZE_N, barB);
tokenB = cuda::device::barrier_arrive_tx(barB, 1, sizeof(sharedB));
}
else {
// Other threads arrive at barrier to synchronise data loads
tokenA = barA.arrive();
tokenB = barB.arrive();
}
// All threads wait for async loads to complete
barA.wait(std::move(tokenA));
barB.wait(std::move(tokenB));
__syncthreads();
}
K 차원을 따라 반복하는 동안, block 안의 한 스레드(thread 0)만이 cp_async_bulk_tensor_2d_global_to_shared를 사용해 행렬 A와 B에 대한 TMA load를 시작하는 책임을 집니다. 이 명령들은 앞서 만든 tensor map을 기반으로 GMEM에서 SMEM으로 전체 2D 타일을 비동기 복사 큐에 넣습니다. 각 copy를 발행한 직후 thread 0은 barrier_arrive_tx를 호출하는데, 이것은 두 가지를 합니다. barrier에 자신이 도착했음을 알리고, 비동기 copy가 몇 바이트를 실제로 써 넣을 예정인지 barrier에 알려줍니다. 나머지 모든 스레드는 단지 bar.arrive()만 호출하여 데이터 전송 없이 도착 사실만 더합니다. 마지막으로 모든 스레드는 bar.wait(token)으로 barrier를 기다립니다(문서 참고). 이 대기는 block의 모든 스레드가 도착하고, TMA 엔진이 전체 타일을 SMEM에 써 넣는 작업을 끝낼 때에만 해제됩니다.
그 시점이 되면 SMEM 타일이 완전히 채워졌고 모든 스레드에게 보인다는 것이 보장되므로, WGMMA 계산 단계로 안전하게 넘어갈 수 있습니다. 바로 이 단계에서 warp group이 실제 MMA를 수행하는 WGMMA 명령을 실행합니다. 계산 단계는 다음 순서를 따릅니다.
wgmma.fence.sync.aligned를 발행합니다. 개념적으로는 warp group 전체에서 관련 레지스터와 SMEM 쓰기가 모두 완료되어 visible 상태이며, 이제 WGMMA 명령을 발행할 준비가 되었다는 의미입니다.wgmma.mma_async로 순차 발행합니다. 코드에서 각 wgmma64 호출은 사실상 하나의 wgmma.mma_async.m64n64k16 명령을 감싼 얇은 래퍼이며, 자세한 내용은 다음에 보고 지금은 블랙박스로 둡니다. 각 WGMMA 명령은 64 × 64 × 16 행렬 곱셈을 수행하고 같은 스레드별 accumulator 레지스터에 누적합니다. 네 번의 호출을 통해 사실상 K 차원의 서로 다른 slice를 따라 이동하면서 같은 64 × 64 출력 타일에 누적하고 있습니다. 이 wgmma.mma_async 명령들은 비동기이므로, 발행했다고 해서 즉시 완료를 의미하지는 않습니다. 대신 하드웨어가 나중에 실행하도록 큐에 넣습니다.wgmma.commit_group 연산을 사용해 앞서 발행한 미완료 wgmma.mma_async 연산들을 하나의 wgmma-group으로 묶어 commit합니다.wgmma.wait_group.wgmma.mma_async 연산이 실행되었고 누적된 결과를 레지스터에서 안전하게 사용할 수 있습니다. 이제 커널은 다음 K 타일로 넘어가거나 store 단계로 진행할 수 있습니다.// Compute phase using WGMMA tensor cores
warpgroup_arrive(); // asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
wgmma64<1, 1, 1, 0, 0>(d, &sharedA[0], &sharedB[0]);
wgmma64<1, 1, 1, 0, 0>(d, &sharedA[WGMMA_K], &sharedB[WGMMA_K]);
wgmma64<1, 1, 1, 0, 0>(d, &sharedA[2 * WGMMA_K], &sharedB[2 * WGMMA_K]);
wgmma64<1, 1, 1, 0, 0>(d, &sharedA[3 * WGMMA_K], &sharedB[3 * WGMMA_K]);
warpgroup_commit_batch(); // asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
warpgroup_wait<0>(); // asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory");
이제 마지막으로 남은 조각은 WGMMA 명령 자체입니다. 위에서는 일부러 전체 커널 구조에 먼저 집중했습니다. 즉 TMA로 타일을 load하는 방식, 동기화 방식, 계산 단계의 구성 같은 큰 흐름을 먼저 다뤘습니다. 계산 단계에서 wgmma64 함수가 호출되는 것은 이미 봤지만, 지금까지는 블랙박스로 취급했습니다. 이제 마침내 그 내부를 열어봅시다.
wgmma64 함수는 inline PTX 명령 wgmma.mma_async.m64n64k16.f32.bf16.bf16 를 감싼 얇은 래퍼입니다. 함수의 형태와 시그니처는 전부 하드웨어 인터페이스가 강제한 결과입니다.
template <int ScaleD, int ScaleA, int ScaleB, int TransA, int TransB>
__device__ void wgmma64(float d[4][8], bf16 *sharedA, bf16 *sharedB)
{
uint64_t desc_a = make_smem_desc(&sharedA[0]);
uint64_t desc_b = make_smem_desc(&sharedB[0]);
각 호출은 행렬 A와 B에 대한 두 개의 matrix descriptor를 만드는 것으로 시작합니다. 이 descriptor는 행렬이 SMEM 어디에 있고 어떻게 배치되어 있는지를 인코딩합니다. 앞서 이야기했듯 WGMMA는 raw pointer를 받지 않고, 대신 이 packed 64-bit descriptor를 소비합니다.
함수의 핵심은 다음 inline PTX 블록입니다.
asm volatile(
"{\n"
"wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 "
"{%0, %1, %2, %3, %4, %5, %6, %7, "
" %8, %9, %10, %11, %12, %13, %14, %15, "
" %16, %17, %18, %19, %20, %21, %22, %23, "
" %24, %25, %26, %27, %28, %29, %30, %31},""
" %32,"
" %33,"
" %34, %35, %36, %37, %38;\n"
"}\n"
: "+f"(d[0][0]), "+f"(d[0][1]), "+f"(d[0][2]), "+f"(d[0][3]), "+f"(d[0][4]), "+f"(d[0][5]),
"+f"(d[0][6]), "+f"(d[0][7]), "+f"(d[1][0]), "+f"(d[1][1]), "+f"(d[1][2]), "+f"(d[1][3]),
"+f"(d[1][4]), "+f"(d[1][5]), "+f"(d[1][6]), "+f"(d[1][7]), "+f"(d[2][0]), "+f"(d[2][1]),
"+f"(d[2][2]), "+f"(d[2][3]), "+f"(d[2][4]), "+f"(d[2][5]), "+f"(d[2][6]), "+f"(d[2][7]),
"+f"(d[3][0]), "+f"(d[3][1]), "+f"(d[3][2]), "+f"(d[3][3]), "+f"(d[3][4]), "+f"(d[3][5]),
"+f"(d[3][6]), "+f"(d[3][7])
: "l"(desc_a), "l"(desc_b), "n"(int32_t(ScaleD)), "n"(int32_t(ScaleA)),
"n"(int32_t(ScaleB)), "n"(int32_t(TransA)), "n"(int32_t(TransB)));
긴 목록 {%0 … %31} 은 호출한 스레드가 소유한 accumulator 레지스터에 해당합니다. 이들은 입력이면서 출력이기도 하기 때문에 operand 목록에서 "+f" 로 표시됩니다. 각 WGMMA 명령은 현재 값을 이 레지스터들에서 읽고, MMA를 수행한 뒤, 갱신된 결과를 같은 레지스터에 다시 씁니다.
그다음 두 operand %32 와 %33은 SMEM 안의 A와 B에 대한 matrix descriptor입니다. 마지막 operand들은 scaling 및 transpose 플래그를 나타내며, template parameter이기 때문에 컴파일 타임 상수로 직접 명령 안에 박힐 수 있습니다.
주어진 K 반복에 대한 모든 WGMMA 명령이 발행되고 완료되면 accumulator 레지스터에는 출력 타일의 최종 결과가 들어 있게 됩니다. 마지막 단계는 이 값들을 global memory로 다시 쓰는 것입니다.
먼저 store를 담당하는 코드 조각을 보여드리겠습니다. 그런 다음 문서에 정의된 accumulator D(우리 코드에서는 C에 해당)의 전체 register fragment layout을 다시 보겠습니다. 다만 warp-group 전체 레이아웃을 한 번에 해석하려고 하기보다, 단일 스레드(특히 thread 0)에 집중해서 설명할 것입니다. 한 스레드의 register fragment가 어떤 출력 원소에 대응하는지 이해하면 전체 레이아웃도 자연스럽게 보입니다.
for (int m_it = 0; m_it < TILE_SIZE_M / WGMMA_M; ++m_it) {
for (int n_it = 0; n_it < TILE_SIZE_N / WGMMA_N; ++n_it) {
for (int w = 0; w < WGMMA_N / 16; ++w) { // w = {0, 1, 2, 3}
// (16 * w) selects the base col of the 16 col block
int col = 16 * w + 2 * (tid % 4);
#define IDX(i, j) ((j + n_it * WGMMA_N) * M + ((i) + m_it * WGMMA_M))
// Apply alpha scaling to accumulator results and add beta*C
block_C[IDX(row, col)] = __float2bfloat16(alpha * d[w][0] + beta * __bfloat162float(block_C[IDX(row, col)]));
block_C[IDX(row, col + 1)] = __float2bfloat16(alpha * d[w][1] + beta * __bfloat162float(block_C[IDX(row, col + 1)]));
block_C[IDX(row + 8, col)] = __float2bfloat16(alpha * d[w][2] + beta * __bfloat162float(block_C[IDX(row + 8, col)]));
block_C[IDX(row + 8, col + 1)] = __float2bfloat16(alpha * d[w][3] + beta * __bfloat162float(block_C[IDX(row + 8, col + 1)]));
block_C[IDX(row, col + 8)] = __float2bfloat16(alpha * d[w][4] + beta * __bfloat162float(block_C[IDX(row, col + 8)]));
block_C[IDX(row, col + 9)] = __float2bfloat16(alpha * d[w][5] + beta * __bfloat162float(block_C[IDX(row, col + 9)]));
block_C[IDX(row + 8, col + 8)] = __float2bfloat16(alpha * d[w][6] + beta * __bfloat162float(block_C[IDX(row + 8, col + 8)]));
block_C[IDX(row + 8, col + 9)] = __float2bfloat16(alpha * d[w][7] + beta * __bfloat162float(block_C[IDX(row + 8, col + 9)]));
#undef IDX
}
}
}
여기서 각 스레드는 자신이 담당하는 row 인덱스를 warp와 lane에 따라 계산합니다. 이어지는 중첩 루프는 M과 N 차원의 논리 타일, 그리고 내부 fragment 구조를 순회합니다. 각 반복에서 스레드는 accumulator 레지스터의 대응 값을 GMEM의 올바른 위치에 써 넣습니다.
모든 스레드에 대한 register fragment layout은 다음처럼 생기며, 예시로 thread 0에만 집중해봅시다.
이제 우리는 코드 관점과 하부 하드웨어 실행 모델 관점 모두에서 커널 전체를 끝까지 살펴보았습니다. 모든 조각이 제자리에 놓였으니, 다음 단계는 커널을 벤치마크하고 성능 결과를 보는 것입니다.
이 커널은 280.4 TFLOP/s 처리량을 달성합니다. 이전 커널이 전체 FP32 경로에서 41.4, 혼합 정밀도에서 31.5를 냈던 것과 비교하면 엄청난 개선입니다. 완전히 다른 게임입니다.
전체 FP32 cuBLAS 벤치마크와 비교하면 성능은 이제 약 544% 입니다. 하지만 당연히 이제부터는 이 비교가 “공정”하지 않습니다. 이 cuBLAS 버전은 tensor core를 사용하지 않기 때문입니다. 다만 차이가 얼마나 큰지 보기 위해 함께 놓아본 것입니다.
이제부터의 진짜 기준은 bf16과 tensor core가 활성화된 cuBLAS입니다. 이 기준으로는 37.8% 를 달성합니다. 참고로 이전 커널들에서는 FP32 비교 상대 성능만 주로 보고했지만, tensor core가 활성화된 cuBLAS와 혼합 정밀도로도 테스트했습니다. warp 타일링은 그 기준에서 겨우 4.3% 였습니다. 즉 기술적으로 보면 4.3%에서 37.8%로 올라간 것이고, 거의 9배 향상입니다.
지금까지는 m64n64k16 하나의 WGMMA 모양만 실험했습니다. 전통적인 CUDA 커널에서는 수학적으로만 맞으면 임의의 타일 크기를 자유롭게 시도할 수 있지만, tensor core MMA 연산은 훨씬 제약이 많습니다. NVIDIA는 PTX ISA 명세에 문서화된 것처럼 A, B, accumulator C 피연산자에 대해 제한된 고정 행렬 모양만 지원합니다(PTX ISA specification, 아래 표 참고).
지원되는 각 MMA 모양은 특정 하드웨어 명령에 대응하며, 단일 WGMMA 연산이 만들어내는 출력 타일의 크기를 결정합니다. 어떤 모양은 자연스럽게 더 큰 accumulator 타일을 만들고, 어떤 모양은 더 작은 타일을 만듭니다. 이것은 warp group이 명령 하나당 수행하는 일의 양에 직접 영향을 주고, 따라서 커널 구조도 그에 맞춰 달라져야 합니다.
자연스럽게도, warp group이 명령 하나당 더 많은 출력 원소를 만들 수 있는 MMA 모양을 선택할 수 있다면, 더 높은 성능을 기대할 수 있습니다. 발행된 WGMMA 하나당 더 많은 산술 작업을 수행하기 때문입니다. 물론 이것은 레지스터 압박처럼 occupancy 감소로 이어질 수 있는 다른 요인에 막히지 않는 한에서만 성립합니다.
이 커널의 목표는 서로 다른 WGMMA 설정을 탐색하고 어떤 것들이 H100 아키텍처에 가장 잘 맞는지 평가하는 것입니다. 가능한 모든 옵션을 전부 시험하는 대신, bf16을 지원하는 몇 가지 대표적인 모양만 골라 성능에 어떤 영향을 주는지 비교해보겠습니다.
아래는 하드웨어가 지원하는 bf16 WGMMA 모양들입니다. 이 목록에서 몇 가지 후보를 골라 더 실험하고 분석할 것입니다.
N차원만 바뀝니다. 이 bf16 dense 모양들에서는M은 64로 고정이고,K도 16으로 고정입니다.
시각적으로 보면 이 커널은 이전 것과 거의 동일합니다. 핵심 차이는 더 큰 TILE_SIZE_M을 덮을 수 있도록 바깥 루프를 하나 도입했다는 것입니다. 이것을 구체적으로 이해하려면, TILE_SIZE_M을 64에서 128로 늘렸다고 상상해보면 됩니다. 그러면 이제 하나의 block 안에서 M 타일 두 개를 반복해 처리하게 됩니다.
TILE_SIZE_M = 128, TILE_SIZE_N = 128, TILE_SIZE_K = 64이고, WGMMA 모양 표에서 WGMMA_K가 16으로 고정임을 알고 있으므로, K 차원은 여전히 이전 커널과 마찬가지로 네 개의 세로 패널로 분할됩니다.
바뀌는 것은 TILE_SIZE_M 차원입니다. WGMMA 모양 표에 따르면 WGMMA_M도 64로 고정이므로, 단일 WGMMA 연산은 한 번에 64행의 출력만 만들 수 있습니다. 그런데 이제 block이 128행을 덮으므로, TILE_SIZE_M 차원을 따라 명시적으로 루프를 돌며 WGMMA를 두 번 호출해야 합니다. 한 번은 앞쪽 64행, 한 번은 뒤쪽 64행에 대해 호출해 전체 출력 타일을 모두 덮는 것입니다.
따라서 이 커널의 구조는 다음과 같습니다.
WGMMA_N = TILE_SIZE_N 을 선택해 하나의 명령으로 block 전체 너비를 덮습니다.m_it를 따라 루프를 돌며 TILE_SIZE_M을 덮습니다.k_it를 따라 루프를 돌며 TILE_SIZE_K를 덮습니다.그 외에는 TMA load, SMEM 레이아웃, store 측면에서는 논리적으로 동일하며, 주요 차이는 계산 단계가 WGMMA 명령 모양에 맞춰 어떻게 구성되느냐에 있습니다. 예를 들어 m64n128k16 모양을 쓰고 TILE_SIZE_M = 128로 늘리면 계산 단계는 다음처럼 생깁니다.
이 커널의 핵심은 따라서 다음과 같은 계산 단계입니다.
// 2. Compute phase using WGMMA tensor cores instructions
warpgroup_arrive();
// Outer loop over TILE_SIZE_M in WGMMA_M steps
// If we have two warp groups, we let each work on a different partition of TILE_SIZE_M
// @example:
#pragma unroll
for (int m_iter = 0; m_iter < rows_per_warp_group / WGMMA_M; m_iter++) {
bf16* sharedA_wgmma_tile_base = sharedA + ((warp_group_idx * rows_per_warp_group) + (m_iter * WGMMA_M)) * TILE_SIZE_K;
// Inner loop iterating over TILE_SIZE_K in WGMMA_K steps
#pragma unroll
for (int k_iter = 0; k_iter < TILE_SIZE_K / WGMMA_K; k_iter++) {
wgmma<WGMMA_N, 1, 1, 1, 0, 0>(d[m_iter], &sharedA_wgmma_tile_base[k_iter * WGMMA_K], &sharedB[k_iter * WGMMA_K]);
}
}
warpgroup_commit_batch(); // asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
warpgroup_wait<0>(); // asm volatile("wgmma.wait_group.sync.aligned %0;\n" ::"n"(N) : "memory");
}
| WGMMA_N | TFLOP/s | cuBLAS 대비 성능 % |
|---|---|---|
| 32 | 230.2 | 31.7% |
| 128 | 407.7 | 56.9% |
| 256 | 70.3 | 9.7% |
몇 가지 다른 WGMMA_N 모양을 실험해보면, 앞서 암시했던 것과 같은 동작을 확인할 수 있습니다. 128이 instruction 효율성과 메모리 재사용의 균형이 가장 좋고, 64는 그 뒤를 바짝 쫓으며 이전 커널과 비슷한 성능을 냅니다. 32는 tensor core를 충분히 활용하지 못하고, 256은 레지스터 압박을 폭발시켜 성능을 무너뜨립니다. 전체적으로 보면 이 커널은 직전 커널 대비 또 한 번 약 1.5배 향상을 제공하고, 이제 cuBLAS 성능의 56.9% 에 도달합니다. 즉 대략 절반쯤 왔습니다.
이전 커널 섹션의 시작에서 우리는 이런 질문을 던졌습니다.
“WGMMA가 기대하는 정확한 형식으로 필요한 행렬 타일을 shared memory에 어떻게 넣고, tensor core가 절대 놀지 않을 만큼 충분히 빠르게 그 작업을 할 수 있을까?”
첫 번째 질문에는 답했지만, 이 커널의 프로파일링 결과를 보면 두 번째 질문에는 아직 답하지 못했다는 것이 분명합니다.
먼저 SoL 섹션의 compute throughput breakdown을 보면 tensor pipeline이 상당히 저활용 상태임을 알 수 있습니다. SM: Pipe Tensor Cycles Active 지표는 지속 가능한 피크 비율의 고작 51.17% 에 도달할 뿐이며, 이는 kernel 실행 시간의 거의 절반 동안 tensor MMA 유닛이 idle 상태라는 뜻입니다. 동시에 어떤 메모리 시스템도 이론 피크에 가깝지 않으므로, 이 커널이 memory-bound가 아니라 scheduling-bound 라는 강한 신호가 됩니다. 왜 그런지 이해하기 위해 이제 scheduler와 warp 통계를 봅시다.
scheduler 통계를 보면 다음과 같습니다.
이는 각 scheduler에 몇 개의 warp가 상주하고는 있지만, 주어진 사이클에서 실제로 실행 준비가 된 warp는 거의 없다는 뜻입니다. 즉 scheduler는 자주 일을 발행하고 싶어 하지만, 선택할 eligible warp가 없는 상태입니다. 정의상 이것은 scheduling 및 dependency 문제입니다.
처음에는 NUM_THREADS를 256으로 늘려 block이 두 개의 warp group을 갖게 해보았습니다. 직관은 occupancy와 active warp 수를 늘려 지연 시간을 숨기자는 것이었습니다. 실제로 occupancy와 active warp는 둘 다 증가했지만, 성능은 좋아지지 않았고 오히려 약간 나빠졌습니다. 이유는 occupancy만 높다고 eligible warp가 생기지는 않기 때문입니다. 이 경우 scheduler당 eligible warp는 0.19 에서 0.23 으로 아주 조금 늘었을 뿐이라, scheduling 동작을 의미 있게 바꿀 정도는 아니었습니다.
🔎 여기에는 128-thread 커널(기준선)과 256-thread 커널(현재)을 비교한 scheduler 및 warp 통계가 있습니다. active warp는 늘었지만 warp들이 여전히 비슷한 stall 문제를 겪고 있어 개선으로 이어지지 않았음을 보여줍니다(지표를 읽으려면 확대해 주세요):
스레드 수를 늘리면 occupancy는 올라갈 수 있지만, 그 warp들이 여전히 barrier, wait, dependency chain에 막혀 있다면 scheduler가 발행할 수 있는 유용한 작업은 거의 늘어나지 않습니다. 그래서 기준이 되는 128-thread 커널의 warp state 통계를 더 자세히 봐야 합니다.
가장 큰 stall 원인은 10.73 의 Long Scoreboard입니다. 이는 warp들이 실행 및 메모리 파이프라인을 흐르는 outstanding dependency에 자주 막히고 있음을 의미합니다. 이 커널에서 이 dependency의 주된 원인은 누산 루프 자체입니다. 각 WGMMA 명령은 같은 accumulator 레지스터를 읽고 쓰므로, 연속된 WGMMA 연산은 본질적으로 서로 의존적입니다. accumulation loop 안의 이 직렬 dependency는 GEMM 알고리즘의 근본적인 성질이며 변하지 않습니다.
Barrier와 Wait stall도 상당합니다. K 차원의 각 반복은 같은 엄격한 직렬 패턴을 따릅니다. TMA load를 발행하고, 그 load가 완료되기를 기다리고, CTA를 동기화하고, WGMMA 계산을 실행하고, 마지막으로 다음 반복으로 넘어가기 전에 모든 outstanding WGMMA 작업이 빠질 때까지 기다립니다. scheduler당 eligible warp 수가 너무 적기 때문에 이런 지연을 숨길 방법이 없고, tensor pipeline은 실행 시간의 상당 부분 동안 idle 상태로 남습니다.
결과적으로 이 커널은 반복 수준에서 큰 pipeline bubble을 가집니다. WGMMA 계산 단계 동안 메모리 시스템은 idle입니다. TMA load 및 동기화 단계 동안 tensor core는 idle입니다. 이런 패턴이 128번의 K 반복 전체에서 똑같이 반복되므로, idle 시간이 누적되어 큰 효율 손실을 만들고, 이것이 프로파일링에서 본 낮은 tensor pipe 활용률을 직접 설명합니다.
이 stall이 코드의 어디서 발생하는지는 아래에 나와 있습니다.
// TMA launch on one thread
if (threadIdx.x == 0) {
cde::cp_async_bulk_tensor_2d_global_to_shared(..., barA);
tokenA = cuda::device::barrier_arrive_tx(barA, 1, sizeof(s.A));
cde::cp_async_bulk_tensor_2d_global_to_shared(..., barB);
tokenB = cuda::device::barrier_arrive_tx(barB, 1, sizeof(s.B));
}
else {
tokenA = barA.arrive();
tokenB = barB.arrive();
}
// Stall Barrier: arrival skew (other warps reach arrive/wait earlier than thread 0)
// Stall Wait: arrive_tx ties barrier completion to async copy bytes landing
barA.wait(std::move(tokenA));
barB.wait(std::move(tokenB));
// Stall Wait: waiting for TMA transaction completion (bytes written to SMEM)
// Stall Barrier: waiting for all warps to arrive at the barrier phase
__syncthreads();
// Stall Barrier: explicit CTA barrier each K-iteration (often redundant here)
for (int k_iter = 0; k_iter < TILE_SIZE_K / WGMMA_K; k_iter++) {
wgmma(...)(d[m_iter], ...);
// Stall Long Scoreboard: dependency chain on d[m_iter] registers
// each WGMMA reads+writes d, next WGMMA needs updated d
}
warpgroup_wait<0>();
// Stall Wait: explicit drain of all WGMMA work before next iteration
시각적으로 보면 각 스레드의 파이프라인은 이렇습니다.
이전 커널에서 확인한 파이프라인 직렬화를 고려하면, 다음 단계는 커널을 producer-consumer 패턴으로 재구성하는 것입니다. 이 설계에서는 하나의 warp group이 producer 역할을 맡아 주로 TMA load를 발행하고, 나머지 warp group(들)은 consumer가 되어 WGMMA를 이용한 tensor core 계산을 수행합니다.
중요한 점은 이 역할들이 병렬로 작동한다는 것입니다. consumer warp group이 현재 K 타일에 대해 WGMMA를 실행하는 동안, producer warp group은 이미 다음 타일에 대한 비동기 TMA load를 발행할 수 있습니다. 이는 TMA 전송이 tensor core pipeline과 독립적으로 동작한다는 사실을 활용하는 것입니다.
이 커널의 목표는 WGMMA 내부의 직렬 accumulation dependency를 제거하는 것이 아닙니다. 그것은 GEMM의 본질이라 제거할 수 없습니다. 대신 목표는 K 반복 사이에서 TMA load와 tensor core 계산을 겹치게 하여 load 및 동기화 지연을 숨기는 것입니다. 그렇게 함으로써 (1) pipeline bubble을 줄이고, (2) tensor core pipe 활용률을 높이며, (3) eligible warp 수를 늘려 scheduler 동작을 개선하고, (4) long scoreboard stall을 줄이고자 합니다.
구현으로 들어가기 전에, 아래 그림은 우리가 목표로 하는 실행 모델을 시각적으로 보여줍니다.
이 커널의 핵심은 다음과 같습니다.
#pragma nv_diag_suppress static_var_with_dynamic_init
__shared__ barrier full[NUM_STAGES]; // Signals data is ready
__shared__ barrier empty[NUM_STAGES]; // Signals slot is available
if (threadIdx.x == 0) {
for (int i = 0; i < NUM_STAGES; i++) {
init(&full[i], num_consumer_groups * 128 + 1); // consumers + producer thread 0
init(&empty[i], num_consumer_groups * 128 + 1);
}
cde::fence_proxy_async_shared_cta();
}
__syncthreads();
if (is_producer) {
// Producer warp group: Issues TMA loads
if (threadIdx.x == 0) {
// Fill the pipeline
for (int stage = 0; stage < NUM_STAGES && stage < num_blocks_k; stage++) {
int block_k_iter = stage;
// Wait for empty slot (initially all are empty, so this passes immediately)
empty[stage].wait(empty[stage].arrive());
// Get pointers for this stage in the flat arrays
bf16* A_stage = s.A + (stage * A_stage_size);
bf16* B_stage = s.B + (stage * B_stage_size);
// TMA loads for A and B
cde::cp_async_bulk_tensor_2d_global_to_shared(A_stage, tensorMapA, block_k_iter * TILE_SIZE_K, num_block_m * TILE_SIZE_M, full[stage]);
cde::cp_async_bulk_tensor_2d_global_to_shared(B_stage, tensorMapB, block_k_iter * TILE_SIZE_K, num_block_n * TILE_SIZE_N, full[stage]);
// Signal data is ready
barrier::arrival_token token = cuda::device::barrier_arrive_tx(full[stage], 1, A_stage_size * sizeof(bf16) + B_stage_size * sizeof(bf16));
}
// Main loop: Continue issuing loads
for (int block_k_iter = NUM_STAGES; block_k_iter < num_blocks_k; block_k_iter++) {
int stage = block_k_iter % NUM_STAGES;
// Wait for this stage to be empty before overwriting
empty[stage].wait(empty[stage].arrive());
// Get pointers for this stage in the flat arrays
bf16* A_stage = s.A + (stage * A_stage_size);
bf16* B_stage = s.B + (stage * B_stage_size);
// Issue next TMA loads
cde::cp_async_bulk_tensor_2d_global_to_shared(A_stage, tensorMapA, block_k_iter * TILE_SIZE_K, num_block_m * TILE_SIZE_M, full[stage]);
cde::cp_async_bulk_tensor_2d_global_to_shared(B_stage, tensorMapB, block_k_iter * TILE_SIZE_K, num_block_n * TILE_SIZE_N, full[stage]);
// Signal data is ready
barrier::arrival_token token = cuda::device::barrier_arrive_tx(full[stage], 1, A_stage_size * sizeof(bf16) + B_stage_size * sizeof(bf16));
}
}
} else {
// Consumer warp groups: Execute WGMMA compute
// Accumulator registers - declared inside consumer branch only so
// ptxas doesn't allocate them for the producer warp group
float d[TILE_SIZE_M / WGMMA_M / num_consumer_groups][WGMMA_N / 16][8];
memset(d, 0, sizeof(d));
// Initially signal all empty slots are available
for (int i = 0; i < NUM_STAGES; i++) {
barrier::arrival_token token = empty[i].arrive();
}
// Main compute loop
for (int block_k_iter = 0; block_k_iter < num_blocks_k; block_k_iter++) {
int stage = block_k_iter % NUM_STAGES;
// Get pointers for this stage in the flat arrays
bf16* A_stage = s.A + (stage * A_stage_size);
bf16* B_stage = s.B + (stage * B_stage_size);
// Wait for data to be ready
full[stage].arrive_and_wait();
// Compute phase using WGMMA
warpgroup_arrive();
#pragma unroll
for (int m_iter = 0; m_iter < rows_per_consumer_warp_group / WGMMA_M; m_iter++) {
bf16* sharedA_wgmma_tile_base = A_stage + ((consumer_warp_group_idx * rows_per_consumer_warp_group) + (m_iter * WGMMA_M)) * TILE_SIZE_K;
#pragma unroll
for (int k_iter = 0; k_iter < TILE_SIZE_K / WGMMA_K; k_iter++) {
wgmma<WGMMA_N, 1, 1, 1, 0, 0>(d[m_iter], &sharedA_wgmma_tile_base[k_iter * WGMMA_K], &B_stage[k_iter * WGMMA_K]);
}
}
warpgroup_commit_batch();
warpgroup_wait<0>();
// Signal this slot is now empty and can be reused
barrier::arrival_token empty_token = empty[stage].arrive();
}
}
파이프라인을 올바르게 조정하기 위해, 파이프라인 각 단계마다 두 개의 shared memory barrier 배열을 할당합니다. empty[stage]는 “이 stage buffer는 써도 된다”는 뜻이고, full[stage]는 “이 stage buffer에는 이제 유효한 A와 B 타일이 들어 있다”는 뜻입니다. thread 0은 모든 barrier를 예상 참가자 수로 초기화한 뒤, barrier가 준비되도록 모든 스레드가 한 번 동기화합니다. 그 후 producer(역시 thread 0이 주도)는 먼저 초반 몇 개 K 타일에 대해 파이프라인을 “채웁니다”. 각 stage가 비어 있기를 기다렸다가, 해당 stage의 s.A와 s.B 슬라이스에 두 개의 TMA load를 발행하고, transaction arrive로 full[stage]를 신호해 소비자들이 바이트가 도착했을 때 깨어날 수 있게 합니다. 파이프라인이 가득 차면 producer는 steady state로 들어갑니다. 각 새 K 타일마다 stage = block_k_iter % NUM_STAGES를 재사용하고, 소비자가 그 stage를 empty로 표시할 때까지 기다렸다가 다음 TMA load로 덮어쓰고 다시 full을 신호합니다. 한편 각 consumer warpgroup은 accumulator fragment d를 레지스터에 할당하고, 모든 stage를 처음에는 empty로 표시해 producer가 즉시 시작할 수 있게 한 뒤, 같은 순환 stage를 선택하는 K 루프를 수행합니다. full[stage]를 기다려 producer가 그 stage loading을 마칠 때까지 대기하고, 자신이 맡은 A 서브타일과 해당 stage의 B 타일에 대해 WGMMA 마이크로루프를 실행하고, WGMMA batch를 commit하고 기다린 뒤, 다음 사이클에서 producer가 그 stage buffer를 재사용할 수 있도록 empty[stage]를 표시합니다.
아래 시각화는 producer와 consumer가 어떻게 협력하는지 더 분명하게 보여줍니다.
이 커널을 구현한 뒤, 저는 Pranjal의 구현에서 보고된 것과 같은 성능 향상을 기대했습니다. 이전 커널은 그의 성능과 거의 정확히 일치했기 때문에 비교 자체는 공정했습니다. 그의 커널이 파이프라인 추가로 큰 향상을 보였다면, 제 구현도 비슷한 이득을 보여야 했습니다. 그러나 그렇지 않았습니다. 이 사실은 곧바로 두 가지 가능성을 시사했습니다. 제 구현이 구조적으로 달라 overlap이 제대로 일어나지 않았거나, 혹은 파이프라인 효과를 제한하는 더 근본적인 무언가가 있었던 것입니다. 저는 producer 쪽을 약간 재구성해 prefill 단계와 steady state를 별도 루프로 나누고, stage 주소 계산 구조도 바꾸었습니다. 이런 변경은 의미적으로는 맞았지만 그의 구현 구조와는 달라졌기 때문에, 문제가 어디서 비롯되는지 분명하지 않았습니다.
무엇이 잘못되었는지 찾는 데는 프로파일링이 핵심이었습니다. 여러 가지 목표 지향 실험을 통해 확인할 수 있었습니다. 프로파일링 리포트는 저장소의 h100_gemm/profiling_reports/producer-consumer-pipeline에서 볼 수 있습니다.
핵심 통찰은 구현 차이와 무관하게 출력 scaling의 효과를 독립적으로 분리해보는 것이었습니다. 저는 Pranjal의 구현을 직접 두 가지 설정으로 프로파일링했습니다.
이 실험 설계는 매우 중요했습니다. scaling이 들어간 참조 커널이 제가 보고 있는 것과 동일한 성능 하락을 보여준다면, 문제는 제 파이프라인 구조가 아니라 epilogue scaling 자체라는 뜻이 되기 때문입니다.
기준선(Scaling 없음): 기준선 커널은 잘 튜닝된 파이프라인에서 기대되는 특성을 보였습니다.
Nsight Compute는 하나의 비효율을 지적했습니다. global store에 대해 “On average only 16 of 32 bytes transmitted per sector are utilised”라는 경고였지만, 이것은 성능 병목으로 이어지지는 않았습니다.
Scaling 포함(알파/베타 scaling epilogue): epilogue에 scaling을 넣자 성능 프로파일이 근본적으로 달라졌습니다. 기준선 대비 처리량 저하:
메모리 동작:
C를 load하고 scaling 후 store하므로 예상 가능한 결과)이 42% 메모리 트래픽 증가가 특히 많은 것을 말해주었습니다. scaled epilogue는 추가 메모리 연산을 도입했을 뿐 아니라, 그것을 비효율적으로 수행하고 있었습니다.
결정적인 발견은 scaling을 넣은 참조 커널이 이제 제 구현과 거의 동일한 성능 특성을 보였다는 점입니다. 처리량 감소, 메모리 트래픽 증가, compute activity 감소(Pipe Tensor Cycles Active [%] 지표로 확인)가 두 커널 모두에서 나타났습니다. 이로써 성능 문제가 제 커널 구조나 레지스터 수 때문이라는 초기 가설은 사라졌습니다. local memory로 spilling조차 없었기 때문입니다. 성능 저하는 scaling이 들어간 참조 구현 자체에서도 재현되었습니다. 파이프라인 구조는 건전했고, 병목은 epilogue로 옮겨간 것이었습니다.
이 관찰을 더 구체적 수치로 입증하기 위해, 이 작업 내내 일관되게 사용해온 문제 크기에서 제 파이프라인 커널 자체를 두 설정으로 벤치마크했습니다. (참고로 코드는 512, 1024, 2048, 4096, 8192 크기에서 벤치마크합니다. 이 부분에 대한 피드백을 받은 적이 있어 언급해둡니다. 하지만 지금까지 보고한 수치는 8192 기준입니다.)
M=N=K= 8192에서 alpha와 beta scaling이 없고, shared memory epilogue staging도 없는 파이프라인 커널은 436.7 TFLOPs를 달성하며, 이는 cuBLAS의 58.4% 에 해당합니다. 반면 alpha와 beta scaling이 모두 포함되고 shared memory staged epilogue를 쓰는 버전은 356 TFLOPs, 즉 cuBLAS의 49.5% 에 도달합니다.
M=N=K= 4096에서는 scaling 없는 변형이 cuBLAS의 73.8% 에 도달하며, 이 크기에서는 Pranjal이 보고한 pipeline 성능을 약간 넘어섭니다. 그러나 cuBLAS가 더 큰 문제 크기에서 더 효율적으로 스케일하기 때문에, 이 이점은 8192에서는 줄어들어 상대 성능이 58.4%로 떨어집니다.
결과는 명확합니다. scaling을 도입하면 8192에서 대략 80 TFLOPs 를 잃습니다. 성능 저하는 주로 producer–consumer 파이프라인 자체 때문이 아니라, mainloop가 충분히 효율적이 되면서 read–modify–write epilogue가 지배적인 제한 요인으로 떠올랐기 때문입니다.
성능 차이는 두 epilogue 변형이 메모리와 상호작용하는 방식에서 비롯됩니다. scaling이 없을 때는 C = accumulator 형태의 write-only 경로입니다. 각 스레드는 global store 하나만 수행합니다. 접근 패턴이 완벽하게 coalesced되지는 않지만(그래서 sector utilisation 경고가 뜹니다), 커널은 여전히 compute-bound입니다. tensor core가 충분히 포화되어 있기 때문에 이런 메모리 비효율이 런타임을 지배하지 않습니다.
반면 scaling이 있으면 C = α × accumulator + β × C 형태의 read-modify-write 경로가 됩니다. epilogue는 이제 다음을 필요로 합니다.
β × Cα × accumulator + (β × C)핵심 변화는 global load 스트림이 추가된다는 점입니다. 스레드별 접근 패턴이 완전히 coalesced되지 않는다면, 그리고 프로파일러는 그렇지 않음을 확인해주고 있으므로, 우리는 이제 그 비효율 비용을 두 번 내고 있는 셈입니다. 한 번은 load 경로에서, 한 번은 store 경로에서입니다.
그러면 질문은 이렇게 바뀝니다. 왜 scaling이 있었던 이전 커널에서는 이 문제가 그렇게 두드러지지 않았을까? 이것은 설명하기가 어려웠지만, 제 추측으로는 파이프라인이 TMA load와 WGMMA compute를 성공적으로 overlap시켰기 때문입니다. 그렇게 하면서 커널이 peak compute throughput에 더 가까워졌고, 그 결과 epilogue 메모리 접근 패턴이 새로운 병목으로 노출되었습니다. epilogue 메모리 접근 패턴의 비효율은 두 경우 모두 존재했지만, 순차 커널에서는 다른 stall 뒤에 가려져 있었던 것입니다.
또 제가 궁금했던 질문 하나는 이렇습니다. epilogue가 문제이고 그것이 compute 이후에 실행된다면, tensor core 활용률은 여전히 높게 보여야 하지 않나?
핵심은 프로파일러 지표가 커널 전체 시간에 대해 평균된다는 점입니다. overlap 덕분에 mainloop가 빨라지면 epilogue가 전체 런타임에서 차지하는 비중이 더 커집니다. 그 시간 동안 tensor core는 idle이므로 overall tensor pipe activity는 떨어집니다. 순차 커널에서는 compute 단계가 더 길고 이미 stall이 있었기 때문에 비효율적인 epilogue가 덜 눈에 띄었을 뿐입니다.
이제 epilogue가 병목이라는 점을 확인했으니, 다음 질문은 이것입니다. 왜 메모리 접근 패턴이 비효율적인가?
스레드들은 WGMMA 실행에 효율적이도록 조직된 레지스터 fragment에 결과를 누적합니다. 반면 global memory는 하드웨어 소개와 이전 커널들에서 이야기했듯 최적 throughput을 위해 coalesced된 128-byte sector access를 요구합니다. 이 두 레이아웃은 자연스럽게 호환되지 않습니다.
하지만 현재 epilogue는 레지스터에서 global memory로 직접 씁니다. 이 때문에 레이아웃을 재조정하는 단계가 건너뛰어지고, 결과적으로 sector utilisation이 나빠집니다.
CUTLASS의 실제 GEMM 커널들은 레지스터에서 global memory로 직접 쓰지 않습니다. 대신 epilogue는 staging 방식으로 진행됩니다.
epilogue는 보통 다음 구조를 따릅니다.
CUTLASS는 epilogue를 위해 SMEM swizzle을 적용하는 것으로 보이지만, 저는 대신 padding을 쓰고 나중에 bank conflict가 여전히 남는지 확인해보겠습니다.
따라서 epilogue는 다음처럼 보입니다.
int tid = threadIdx.x % 128;
int lane = tid % 32;
int warp = tid / 32;
uint32_t row = warp * 16 + lane / 4;
// @note C is column-major
bf16* block_C = C + (num_block_n * TILE_SIZE_N * M) + (num_block_m * TILE_SIZE_M);
constexpr int TILE_M_PAD = TILE_SIZE_M + 8;
#define IDX_GMEM(r, c) ((c) * M + (r))
#define IDX_SMEM(r, c) ((c) * TILE_M_PAD + (r))
// Phase 1: alpha-scaled accumulators -> shared staging tile
for (int m_iter = 0; m_iter < rows_per_consumer_warp_group / WGMMA_M; m_iter++) {
int row_tile_base_C = (consumer_warp_group_idx * rows_per_consumer_warp_group) + (m_iter * WGMMA_M);
for (int w = 0; w < WGMMA_N / 16; w++) {
int col = 16 * w + 2 * (tid % 4);
s.C_epi[IDX_SMEM(row + row_tile_base_C, col)] = __float2bfloat16(alpha * d[m_iter][w][0]);
s.C_epi[IDX_SMEM(row + row_tile_base_C, col + 1)] = __float2bfloat16(alpha * d[m_iter][w][1]);
s.C_epi[IDX_SMEM(row + 8 + row_tile_base_C, col)] = __float2bfloat16(alpha * d[m_iter][w][2]);
s.C_epi[IDX_SMEM(row + 8 + row_tile_base_C, col + 1)] = __float2bfloat16(alpha * d[m_iter][w][3]);
s.C_epi[IDX_SMEM(row + row_tile_base_C, col + 8)] = __float2bfloat16(alpha * d[m_iter][w][4]);
s.C_epi[IDX_SMEM(row + row_tile_base_C, col + 9)] = __float2bfloat16(alpha * d[m_iter][w][5]);
s.C_epi[IDX_SMEM(row + 8 + row_tile_base_C, col + 8)] = __float2bfloat16(alpha * d[m_iter][w][6]);
s.C_epi[IDX_SMEM(row + 8 + row_tile_base_C, col + 9)] = __float2bfloat16(alpha * d[m_iter][w][7]);
}
}
__syncthreads();
// Phase 2: coalesced write to GMEM (alpha*D + beta*C)
int row4_in_group = lane * 4;
int group_base_row = consumer_warp_group_idx * rows_per_consumer_warp_group;
if (row4_in_group < rows_per_consumer_warp_group) {
int r0 = group_base_row + row4_in_group;
for (int c = warp; c < TILE_SIZE_N; c += 4) {
block_C[IDX_GMEM(r0 + 0, c)] = __float2bfloat16(__bfloat162float(s.C_epi[IDX_SMEM(r0 + 0, c)]) + beta * __bfloat162float(block_C[IDX_GMEM(r0 + 0, c)]));
block_C[IDX_GMEM(r0 + 1, c)] = __float2bfloat16(__bfloat162float(s.C_epi[IDX_SMEM(r0 + 1, c)]) + beta * __bfloat162float(block_C[IDX_GMEM(r0 + 1, c)]));
block_C[IDX_GMEM(r0 + 2, c)] = __float2bfloat16(__bfloat162float(s.C_epi[IDX_SMEM(r0 + 2, c)]) + beta * __bfloat162float(block_C[IDX_GMEM(r0 + 2, c)]));
block_C[IDX_GMEM(r0 + 3, c)] = __float2bfloat16(__bfloat162float(s.C_epi[IDX_SMEM(r0 + 3, c)]) + beta * __bfloat162float(block_C[IDX_GMEM(r0 + 3, c)]));
}
}
#undef IDX_GMEM
#undef IDX_SMEM
이 커널을 프로파일링하면 분명한 개선이 보입니다. Tensor core activity가 거의 32% 증가하고, L1/TEX, L2, DRAM throughput도 모두 상승하며, 커널 실행 시간은 줄고 전체 처리량은 개선됩니다. global memory load/store의 비정상 coalescing 경고도 사라집니다. 여기서 비교 기준은 초기 파이프라인 버전입니다.
하지만 producer–consumer 파이프라인이 올바르게 동작하고 staged epilogue가 메모리 접근 비효율을 해결했음에도, 이 커널은 여전히 커널 8을 넘어서지 못합니다. Tensor core 활용률도 여전히 기대보다 약간 낮습니다.
파이프라인이 의도대로 동작하는지 확인하기 위해 Nsight Systems로 커널을 프로파일링했습니다. trace를 보면 TMA 전송과 WGMMA 실행이 직렬이 아니라 시간상 overlap되어 있음을 확인할 수 있습니다. 이는 producer–consumer 메커니즘이 실제로 load latency를 숨기고 있음을 뜻합니다. 그럼에도 불구하고, 수많은 실험과 Nsight Compute 지표의 세심한 검토 끝에도 이 커널이 왜 이전 것을 능가하지 못하는지에 대해서는 아직 완전히 만족스러운 설명을 갖고 있지 못합니다. 이는 앞으로 더 조사할 여지가 남아 있다는 뜻입니다. 이 동작에 대한 통찰이 있으시다면 꼭 알려주세요.
따라서 남은 성능 격차는 파이프라인이 망가져서라기보다, 각 파이프라인 stage에서 고정 오버헤드(예: barrier 단계, warpgroup fence, stage handoff, epilogue)에 비해 얼마나 많은 유효 계산을 뽑아내느냐의 문제로 보입니다.
자연스러운 다음 단계는 더 큰 WGMMA_N을 사용하고 추가 consumer warpgroup을 도입해 출력 타일의 너비를 키우는 것입니다. 이렇게 하면 K 타일당 수행되는 계산량이 늘어나고, 파이프라인 오버헤드가 더 많은 산술 작업에 걸쳐 더 잘 상쇄됩니다. overlap이 이미 달성된 상태에서는 바로 이런 트레이드오프가 우리가 원하는 것입니다.
이 변경을 구현해, 구체적으로 WGMMA_N을 256으로 늘리고 두 개의 consumer warpgroup으로 실행하면 성능이 더 올라갑니다. M=N=K= 8192에서 이 커널은 이제 463.9 TFLOPs에 도달하며, 이는 cuBLAS의 63.1% 에 해당합니다. 4096에서는 70.5% 상대 성능을 달성합니다.
이는 우리의 직관을 확인해줍니다. memory와 compute가 올바르게 overlap되기 시작하면, 각 파이프라인 stage당 더 많은 작업을 뽑아내는 것이 효과적인 조절 수단이 됩니다. 각 stage의 arithmetic intensity를 높이고 producer와 consumer warpgroup 사이에 레지스터 자원을 더 의도적으로 분배함으로써, tensor core 파이프라인 포화에 더 가까워질 수 있습니다.
그럼에도 불구하고 아직 상당한 성능이 남아 있으며, 특히 epilogue가 도입한 성능 저하를 완화하는 부분에 더 많은 여지가 있습니다.