(3/n) Triton matmul 커널 탐험 - cross tile pipelining 적용
이번에는 persistent kernel에 cross tile pipelining 기법을 적용하려고 한다. 이전의 persistent kernel에서는 커널이 직접 타일을 스케줄링하는 것까지 구현을 했었고 타일이 넘어갈 때에 다른 설정은 하지 않았다. 그래서 (m, n) 방향의 타일이 바뀔 때마다 파이프라인이 끊겼다가 다시 시작되었다. cross tile pipelining은 타일 경계에서 파이프라인을 끊지 않고, 타일 루프와 그 안쪽의 k 루프를 하나로 합쳐(flatten), 한 타일의 마지막 k 블록 바로 다음에 다음 타일의 첫 k 블록이 이어서 들어가도록 한다. 이렇게 하면 타일이 바뀔 때마다 생기던 파이프라인 재시작 오버헤드를 줄일 수 있다.
triton에서는 간단하게 타일 루프에 flatten=True를 설정을 하면 triton 컴파일러가 TTGIR로 변환하는 단계에서 위와 같이 바깥 loop와 안쪽 loop를 합쳐 하나의 loop로 만들게 된다.
# 원래 이중 루프
for tile in range(start, num_tiles, 340):
[타일 준비: 좌표, 포인터, mask]
for k in range(K_steps):
[K 본문: load, dot]
[타일 마무리: store]
# 합친 뒤 단일 루프
outer_iters = tl.cdiv(num_tiles - start, 340)
total = outer_iters * max(K_steps, 1)
for i in range(total):
if k == 0:
[타일 준비]
[K 본문]
if k == K_steps-1:
[타일 마무리]
k = 0 if k == K_steps-1 else k + 1
tl.range에 붙이는 속성인 num_stages 값에 대해서도 얘기해보면, 이 값은 triton 컴파일러가 파이프라이닝을 하는 깊이 단계라고 알려져 있는데 이중 루프일때 triton 컴파일러는 가장 안쪽 loop에만 파이프라인을 적용한다.
flatten=True를 통해 루프가 합쳐지면 합쳐진 루프가 가장 안쪽 루프가 되므로 바깥 loop에 붙여진 num_stages 값도 비로소 의미를 갖는다.
성능 측정
| K | Swizzle persistent | Cross-tile persistent | Cross-tile 시간 변화 |
|---|---|---|---|
| 128 | 0.4319 ms | 0.4065 ms | −5.87% |
| 512 | 1.4001 ms | 1.3772 ms | −1.63% |
| 1024 | 2.6286 ms | 2.6399 ms | +0.43% — 거의 동률 |
| 2048 | 5.2736 ms | 5.4139 ms | +2.66% |
| 4096 | 10.7121 ms | 10.9906 ms | +2.60% |
측정 조건: RTX 5090, FP16, M = N = 16384, BLOCK 128×64×32, warps 4, stages 4, program 340개.
K가 작을 때는 빨라지고, 1024를 넘어가면 오히려 조금 느려졌다.
K가 작으면 타일당 K 단계가 적어서(K = 128이면 4단계) 결과 저장과 다음 타일 준비, 즉 타일 경계가 전체 시간의 큰 부분을 차지한다. cross-tile은 저장하는 동안 다음 타일의 입력을 미리 가져오므로 이 구간에서 이득이 크다.
반대로 K가 크면 경계에서 아끼는 시간은 작아지고, 합친 루프가 K 단계마다 더 하는 일이 쌓인다.
TTGIR을 비교해 보면 통합 루프는 매 단계 타일 경계 카운터와 조건문을 검사하고, 타일이 끝났을 때 누적값을 0으로 되돌리는 select를 128×64 누적값 전체에 대해 수행한다. 또 결과 저장이 주 루프 안으로 들어오면서 입력 버퍼가 살아 있는 채로 레이아웃을 변환해야 해서 shared memory도 36 -> 40 KiB로 늘었다.
결국 cross-tile pipelining은 항상 켜는 최적화라기보다, 타일 경계 비용의 비중이 클 때 효과가 있는 선택지다.