(4/n) Triton matmul 커널 탐험 - TMA, Warp Specialization, Epilogue subtiling

30 Sep, 2026

Tensor Memory Accelerator

Hopper 아키텍처부터 도입된 하드웨어로 global memory와 shared memory 사이에서 다차원 타일을 비동기로 옮기는데 사용한다. 스레드가 개별 load/store를 발행하는 대신 descriptor 기반으로 명령 하나로 타일 전체를 비동기 전송하므로, 주소 계산과 경계 처리를 전용 하드웨어가 대신해 스레드·레지스터 부담이 줄고 연산과 전송을 겹치기 쉽다. 사용하려면 가장 안쪽 차원이 연속이고, 나머지 차원의 stride는 16바이트 배수, 시작 주소는 16바이트 정렬이어야 한다.

Triton에서는 host에서 텐서의 base, shape, stride, block_shape을 담은 TensorDescriptor를 만들고, 기존에 포인터 텐서를 넘기던 자리에 이 descriptor를 커널 인자로 넘긴다. 커널 안에서는 포인터 연산 대신 타일의 시작 좌표만 지정해 desc.load([m, k]), desc.store([m, n], value)로 접근한다. 범위 밖 처리는 하드웨어가 해 주므로 mask를 직접 만들 필요가 없다(load는 0으로 채움).

# @host
a_desc = TensorDescriptor(
    base=a, shape=[M, K], strides=list(a.stride()),
    block_shape=[TMA_BLOCK_M, TMA_BLOCK_K],
)

# @device
a_tile = a_desc.load([start_m, start_k])
c_desc.store([store_m, store_n], acc.to(tl.float16))

Warp Specialization

같은 thread block 안에서 warp별로 역할을 나누는 기법이다. CUDA 기준으로는 데이터를 로드하는 producer warp와 MMA를 수행하는 consumer warp로 분리하고, shared memory의 다단계 버퍼와 mbarrier로 동기화한다.

기존 방식 vs Warp specialization

Triton에서는 루프를 tl.range(..., warp_specialize=True)로 바꾸면 컴파일러가 load 파트와 MMA 파트로 자동 분할한다.

Epilogue subtiling

epilogue에서 accumulator를 BM × BN 전체가 아니라 BM × BN/2 같은 조각으로 나눠 변환하고 store하는 기법이다.

기존에는 출력 타일 전체 크기의 shared memory가 epilogue에 필요해서, 그만큼 mainloop의 stage 수를 줄여야 했다. Subtiling은 조각 단위로 같은 버퍼를 재사용하므로 epilogue shared memory가 절반으로 줄고, 남는 공간으로 stage를 더 늘려 메모리 지연을 더 잘 숨길 수 있다. 레지스터 압력도 줄고, 한 조각의 store와 다음 조각의 변환을 겹칠 수 있다.

triton에서는 다음과 같이 구현한다.

if EPILOGUE_SUBTILE:
    #EPI 1: acc를 N 방향의 왼쪽/오른쪽 절반으로 나눈다.
    acc_left, acc_right = acc.reshape(BLOCK_M, 2, BLOCK_N // 2).permute(0, 2, 1).split()
    # EPI 2: 두 조각을 FP16으로 바꿔 c_desc.store로 저장한다.
    c_desc.store([store_m, store_n], acc_left.to(tl.float16))
    c_desc.store([store_m, store_n + BLOCK_N // 2], acc_right.to(tl.float16))
else:
    c_desc.store([off_m, off_n], acc.to(tl.float16))

성능 측정 (B200 @modal)

그동안 내가 만들어온 triton kernel 성능

• M/N/K 모두 4096, FP16
• F/W/E: flatten / warp_specialize / epilogue_subtile
• warps: 요청값 → 컴파일된 실제 값
• P: programs/SM. 전체는 모든 output tile만큼 launch

공식 커널 성능

• 테스트 조건은 거의 같음. 다만 B가 [N, K]임.
• Raw tiled와 swizzle tiled는 launch stages가 3이고, 코드 내부 K-loop에는 num_stages=4가 명시되어 있음.
• 모두 BM/BN/BK=128/256/64, stages=4, GROUP_SIZE_M=8을 선택함.

그래도 동일한 성능이라고 볼수 있을정도 까지는 되지 않았나? 싶어서 만족했다.

블록 크기를 더 늘리면,

성능이 더 좋아짐을 알수 있었다. 아마 TMA box는 차원당 최대 256이기 때문에 더 크게 하기는 어려울 것이다.

triton matmul 커널 탐험을 우선 마무리하면서 다음 방향에 대한 고민....

triton matrix multiplication을 공부하면서 느끼고 있는것중에 가장 신기한 건 compiler가 너무 많은 부분을 해줘서 생각보다는 쉽다는 것이였다. flatten, warp_specialize 같은건 정말 공짜로 구현할 수 있는 너무 사기스럽고... 올해 초에 cuda 할떄는 이런 부분이 어렵게 느껴져서 결국 직접하지는 못했었는데 말이다. 그래서 지금은 컴파일러가 마법을 부리는 부분이 좀더 관심이 가고 알아가고 싶다.

#blog #gpu #matmul #tma #triton