currybab's blog

(2/n) Triton matmul 커널 탐험 - grouped ordering(CTA swizzle)도 늘 빠르게 만들지는 않는다.

triton matrix multiplication의 튜토리얼에서는 l2 cache 적중률을 위해서 grouped ordering을 사용한다고 얘기한다. 이 방법을 swizzle이라고도 표현하는데 gpu 커널에서 swizzle이라고 표현하는 여러 곳이 있어서 CTA swizzle이라고 표현하려고 한다.

구현

CTA swizzle의 실제 구현은 (tile_m, tile_n)을 정해주었던 부분을 간단하게 고치면 된다.

# SWIZZLE 1: GROUP_SIZE_M개의 M tile을 한 그룹으로 묶는다.
# num_tiles_in_group, group_id, first_tile_m을 계산한다.
num_tiles_in_group = GROUP_SIZE_M * num_n_tiles
group_id = tile_id // num_tiles_in_group
tile_id_in_group = tile_id % num_tiles_in_group
first_tile_m = group_id * GROUP_SIZE_M

# SWIZZLE 2: 마지막 그룹의 실제 M tile 수 group_size_m을 계산한다.
# tl.minimum을 사용한다. GROUP_SIZE_M보다 작을 수 있다.
group_size_m = tl.minimum(GROUP_SIZE_M, num_m_tiles - first_tile_m)

# SWIZZLE 3: 그룹 내부 linear id에서 M이 먼저 증가하도록 좌표를 만든다.
tile_m = first_tile_m + (tile_id_in_group % group_size_m)   # tile_id // num_n_tiles
tile_n = tile_id_in_group // group_size_m                   # tile_id % num_n_tiles

중간에 계산식이 좀 생겼기 때문에 스레드당 레지스터를 좀더 사용할 것으로 예측해볼 수 있으며, 실제로 naive tiled matmul에서 swizzle을 적용하였을 때에는 128 -> 134, persistent matmul에서는 174 -> 175로 소폭 증가하는 것을 확인할 수 있었다.

첫 측정 (M=N=8192)

RTX 5090: M=8192, N=8192, K=4096, GROUP_SIZE_M=8
torch: 2.3649 ms, 232.46 TFLOPS
naive tiled: 2.4443 ms, 224.91 TFLOPS
naive persistent: 2.5719 ms, 213.75 TFLOPS
swizzle tiled: 2.5502 ms, 215.57 TFLOPS
swizzle persistent: 2.6133 ms, 210.37 TFLOPS

위와 같이 CTA swizzle을 적용한 버전이 적용하지 않은 버전보다 느려진 것을 확인할 수 있었다. 원래라면 swizzling을 적용함으로써 L2 캐시 재사용이 개선되고, 그것이 연산 파이프라인의 대기 감소로 나타나 실행 시간 감소(처리율 증가)로 이어진다는 가정이였다.

확인할 지표들

원인을 파악하기 위해 swizzling이 실제로 캐시 재사용을 개선했는지, 그 개선이 실행 시간을 줄일 수 있는 상황인지 확인하기 위해 L2 hit rate, DRAM throughput, tensor 파이프라인 이용률을 따져보려고 한다.

L2 hit rate는 Memory Workload Analysis 섹션에서, DRAM throughput은 GPU Speed Of Light Throughput 섹션에서 쉽게 확인할 수 있다. tensor 파이프라인 이용률은 Compute Workload Analysis 섹션에서 ncu cli의 --print-details all 옵션을 켜서 확인할 수 있다. Pipe Utilization (Elapsed Cycles) 항목에서 Tensor (All) 부분을 확인하면 된다.

Pipe Utilization 읽는 법

Pipe Utilization (Elapsed Cycles)
Table Name : Aggregate Pipe Utilization (% of elapsed cycles)
Table Description : Aggregate pipeline utilization based on the number of elapsed (all) cycles in the workload. This takes the rates of different instructions executing on the pipeline into account. For an instruction requiring 4 cycles to complete execution, the counter is increased by 1 for 4 cycles. Use this to understand the overall pipeline utilization over the entire workload.
------------ ----------- ------------
Metric Name  Metric Unit Metric Value
------------ ----------- ------------
Tensor (All)           %        93.19
ALU                    %         3.94
FMA                    %         1.28
------------ ----------- ------------

이 표는 단순히 명령 개수를 세지 않고 명령이 파이프라인을 사용하는 사이클을 반영한다. 설명에 나온 예처럼, 어떤 명령이 4사이클에 걸쳐 처리된다면 명령 1개로만 세는 대신 그 4사이클의 활동을 반영하는 방식으로 측정된다. 이 지표를 기준으로 Tensor 파이프라인이 전체 실행 구간에서 높은 활동 수준을 유지했다고 이해하면 된다. 파이프라인은 서로 겹쳐서 동작할 수 있고, 각각의 처리 능력을 기준으로 측정되어서 세 값을 합쳐 100%로 해석하지는 않는다. 또한 Tensor (All)이 93.19%라고 해서 최대 TFLOPS의 93.19%는 아니다. 현재 사용 중인 명령들의 파이프라인 활동을 나타내는 값이다.

Compute Workload Analysis 섹션 헤더에서 볼 수 있는 SM Busy 값은 여러 SM 파이프라인의 처리율 중 가장 높은 값을 반영하는 요약 지표로 이 커널에서는 SM Busy = Tensor 이용률로 볼 수 있다.

첫 측정 결과 분석

이제 확인하는 방법을 모두 알게 되었으니 ncu로 값을 확인해보면 다음과 같았다.

구현 L2 적중률 DRAM 처리율 Tensor 파이프라인 이용률
naive tiled 97.94% 5.36% 96.37%
swizzle tiled 96.79% 8.54% 96.66%
naive persistent 97.71% 5.93% 93.09%
swizzle persistent 96.38% 9.25% 93.19%

swizzle을 적용하였을때 L2 적중률이 떨어졌고, 캐시 적중률 하락과 함께 DRAM 처리율은 증가했다. 또한 Tensor 파이프라인 이용률은 그대로 였다. 자원 사용량과 주요 연산 지표가 유사하다는 점을 고려하면, 캐시 재사용 악화가 소폭의 성능 저하에 기여했을 가능성이 높다.

문제 크기를 키워보면 (M=N=16384)

그렇다면 swizzle은 이 커널에서 효과가 없는 최적화일까? 이를 확인하기 위해 K=4096과 커널 설정은 유지하고, M과 N을 각각 16384로 늘려보았다.

RTX 5090: M=16384, N=16384, K=4096, GROUP_SIZE_M=8
torch: 9.8971 ms, 222.19 TFLOPS
naive tiled: 11.6640 ms, 188.53 TFLOPS
naive persistent: 11.7115 ms, 187.77 TFLOPS
swizzle tiled: 10.4398 ms, 210.64 TFLOPS
swizzle persistent: 10.6309 ms, 206.85 TFLOPS

이번에는 결과가 반대로 나타났다. 해당 벤치마크 실행에서 swizzle을 적용한 tiled 버전은 약 10.5%, persistent 버전은 약 9.2% 실행 시간이 감소했다. 앞선 작은 문제에서는 이득이 없었던 최적화가 문제 크기를 늘리자 효과를 보이기 시작한 것이다.

왜 큰 문제에서만 효과가 났을까

같은 조건에서 ncu로 캐시와 연산 지표를 별도 측정했다.

구현 L2 적중률 DRAM 처리율 Tensor 파이프라인 이용률
naive tiled 66.89% 90.04% 97.42%
swizzle tiled 94.58% 15.15% 97.88%
naive persistent 79.19% 55.96% 94.35%
swizzle persistent 94.80% 14.10% 94.63%

가장 크게 달라진 것은 L2 적중률이었다. 기존 row-major 구현의 적중률은 tiled에서 약 67%, persistent에서 약 79%까지 떨어졌지만, swizzle을 적용한 두 버전은 약 95%를 유지했다. 또한 row-major 구현에서는 swizzle 버전에 비해 높은 DRAM 처리율을 보였는데 작은 문제와 달리, 이번에는 캐시 재사용 개선과 실행 시간 감소가 함께 관측되었다.

재사용 거리와 캐시 용량

이 차이는 캐시 용량과 데이터 재사용 거리를 통해 설명해볼 수 있다. 재사용 거리는 같은 데이터를 다시 접근하기까지 사이에 접근하는 서로 다른 데이터의 양을 의미한다. 재사용 거리가 커지면 필요한 데이터가 다시 사용되기 전에 캐시에서 밀려날 가능성도 커진다.

FP16 입력에서 행렬 크기를 계산하면 다음과 같다.

데이터 M=N=8192 M=N=16384
A 전체 64 MiB 128 MiB
B 전체 64 MiB 128 MiB
C 전체 128 MiB 512 MiB
RTX 5090 L2 용량 96 MiB 96 MiB

A, B의 접근과 C 저장이 모두 캐시에 영향을 주며, 여러 CTA가 동시에 서로 다른 데이터를 사용한다. 중요한 것은 재사용할 데이터가 다음 접근 시점까지 캐시에 남아 있을 수 있는가이다.

타일 배정 순서 비교

이를 타일 배정 순서와 연결해보자. 현재 BLOCK_N=64이므로 N 방향의 출력 타일 수는 다음과 같다.

N=8192:  num_n_tiles = 8192 / 64  = 128
N=16384: num_n_tiles = 16384 / 64 = 256

row-major에서 출력 타일의 번호는 다음과 같다.

tile_id = tile_m * num_n_tiles + tile_n

따라서 같은 B 패널을 사용하는 C(m,n)과 C(m+1,n)의 tile_id 차이는 num_n_tiles이다. N을 두 배로 늘리면 이 간격도 128에서 256으로 증가한다. 같은 B 패널을 다시 사용하는 작업이 배정 순서상 더 멀어지는 것이다.

반면 GROUP_SIZE_M=8인 grouped ordering에서는 같은 그룹 안의 C(m,n)과 C(m+1,n)을 연속된 tile_id에 배정한다.

재사용 관계 Row-major Grouped ordering, G=8
같은 B를 사용하는 인접 M 타일의 tile_id 간격 128 또는 256 1
같은 A를 사용하는 인접 N 타일의 tile_id 간격 1 8

즉 swizzle은 A의 재사용 간격을 늘리는 대신 B의 재사용 간격을 줄이는 선택이다. 여기서 tile_id 간격은 실제 재사용 거리 자체가 아니라, 배정 순서상 지역성을 설명하는 지표다. 실제 접근 순서는 GPU 스케줄링과 각 CTA의 진행 속도에 따라 달라진다.

데이터 크기도 함께 고려해야 한다. 출력 타일 하나가 K 전체를 계산할 때 사용하는 입력 패널은 다음과 같다.

A 패널: BLOCK_M × K × 2 bytes
      = 128 × 4096 × 2 = 1 MiB

B 패널: K × BLOCK_N × 2 bytes
      = 4096 × 64 × 2 = 0.5 MiB

row-major 방식에서는 하나의 A 패널을 사용하는 출력 타일 행을 계산할 때, B 패널 전체를 순회한다.

N=8192:
128 × 0.5 MiB + 1 MiB = 65 MiB

N=16384:
256 × 0.5 MiB + 1 MiB = 129 MiB

RTX 5090의 L2 cahce 사이즈가 96 MiB임을 고려했을 때, N=8192에서는 이 입력 범위를 캐시에 유지할 용량상의 여지가 있지만, N=16384에서는 전체를 동시에 유지할 수 없다. 이는 문제 크기를 늘렸을 때 row-major의 L2 적중률이 크게 하락한 결과를 설명하는 단서다.

grouped ordering은 A 패널 8개를 반복해서 사용하면서 B 패널을 하나씩 순회한다. 각 B 패널은 가까이 배정된 출력 타일 8개에서 사용한 뒤 다음 패널로 넘어간다. 이 그룹 안에서 B를 재사용하기 위해 B 전체를 캐시에 유지할 필요가 없는 것이다. 같은 B 패널을 사용하는 출력 타일 8개의 고유 입력량은 다음과 같다.

서로 다른 A 패널 8개 + 공통 B 패널 1개
= 8 × 1 MiB + 0.5 MiB
= 8.5 MiB

8.5 MiB의 입력을 사용하는 타일 묶음 안에서 같은 B 패널을 8개 출력 타일에서 사용할 기회를 만든다.

정리

작은 문제에서는 기존 순서에서도 L2 적중률이 약 98%로 높아 추가 개선의 여지가 작았다. 반면 문제 크기를 늘리자 기존 순서의 캐시 재사용이 악화되었고, grouped ordering이 재사용 간격을 줄이는 효과가 드러난 것으로 해석할 수 있다.

swizzle을 적용하지 않은 두 구현 사이의 차이도 흥미롭다. naive persistent는 naive tiled보다 L2 적중률이 높고 DRAM 처리율이 낮았다. 같은 row-major 좌표 변환을 사용하더라도, program 수와 각 program에 타일을 배정하는 방식이 달라지면 실제 캐시 재사용 양상도 달라질 수 있는 것이다. 그러나 실행 시간은 거의 동일했고, Tensor 파이프라인 이용률은 persistent 쪽이 소폭 낮았다. 캐시 효율 개선이 반드시 실행 시간 개선으로 이어지는 것은 아니었다.

반대로 swizzle 적용 전후를 비교하면 Tensor 파이프라인 이용률은 모두 높은 수준을 유지하면서도 실행 시간은 감소했다. 따라서 Tensor 이용률이 높다는 사실만으로 메모리 최적화의 여지가 없다고 판단해서도 안 된다. 각 지표는 캐시·메모리·연산 경로의 서로 다른 측면을 보여주므로, 실제 실행 시간과 함께 해석해야 한다.

이번 실험에서 CTA swizzle의 효과는 문제 크기에 따라 달라졌다. 작은 문제에서는 효과가 없거나 소폭 느려졌지만, 큰 문제에서는 L2 적중률을 크게 개선하며 실행 시간을 줄였다. grouped ordering은 항상 적용해야 하는 정답이 아니라, 행렬 크기와 타일 구성, GPU 캐시 용량, CTA 간 데이터 재사용 관계에 맞춰 검증하고 선택할 최적화라고 정리할 수 있다.

#blog #gpu #matmul #swizzle #triton