GPU 시스템 06 - Triton과 Kernel Optimization 실전 감각
Triton과 실제 커널 최적화를 연결해서 GPU kernel engineer 관점의 감각을 잡기
CUDA 다음에 Triton을 보는 이유
CUDA를 어느 정도 공부한 뒤 Triton을 보면 느낌이 다르다. 처음부터 Triton만 보면 문법이 간결해서 쉬워 보일 수 있지만, 왜 특정 block size를 잡고 왜 tile 단위로 쪼개는지까지 이해하기는 어렵다. 반면 CUDA에서 warp, shared memory, memory coalescing을 이미 봤다면 Triton은 그 개념을 더 빠르게 실험하는 도구처럼 보인다.
LLM 최적화 문맥에서 Triton이 자주 나오는 이유도 여기에 있다. 실제 모델 연산을 대상으로 fused softmax, layernorm, attention 계열 커널을 빠르게 시도해볼 수 있기 때문이다.
Triton을 볼 때 중요한 관점
Triton은 CUDA를 완전히 대체하는 마법 같은 도구가 아니다. 오히려 GPU 구조를 유지한 채 더 생산적인 인터페이스를 주는 편에 가깝다.
그래서 Triton을 공부할 때도 핵심은 그대로다.
- 어떤 tile을 읽을 것인가
- 메모리 접근이 얼마나 정돈되어 있는가
- 중간 결과를 얼마나 재사용하는가
- 어떤 연산을 fuse해서 global memory 왕복을 줄일 것인가
문법보다 이 질문이 더 중요하다.
이 지점에서 Triton을 "CUDA보다 쉬운 문법" 정도로만 보면 아쉽다. 더 중요한 것은 Triton이 GPU 구조를 완전히 숨기지 않으면서도, tile design과 fusion 아이디어를 더 빠르게 실험하게 해준다는 점이다.
왜 kernel fusion이 중요한가
딥러닝 모델에서는 연산 하나하나는 단순해 보여도, 그 사이에 중간 텐서를 global memory에 계속 읽고 쓰면서 비용이 커지는 경우가 많다. 예를 들어 layernorm, bias add, activation이 각각 따로 실행되면 중간 결과가 여러 번 메모리를 오가게 된다.
kernel fusion은 이 왕복을 줄이기 위해 여러 연산을 하나의 kernel 안에서 처리하는 전략이다. 이때 중요한 것은 단순히 합치는 것이 아니라, 합쳤을 때 메모리 흐름과 register 사용량이 실제로 더 좋아지는지 보는 것이다.
즉, fusion은 "연산 개수를 줄인다"보다 "global memory 왕복을 얼마나 줄일 수 있는가" 관점으로 보는 편이 더 정확하다. 그리고 이때 항상 register pressure와 numerical behavior를 같이 봐야 한다.
Flash Attention류 커널에서 보이는 사고방식
fast attention 계열 커널이 흥미로운 이유는 연산량보다 메모리 흐름과 tiling 전략이 훨씬 중요하게 드러나기 때문이다. 전체 attention matrix를 통째로 materialize하지 않고 tile 단위로 계산을 흘려보내는 사고방식이 핵심이다.
이런 예제를 보면 GPU kernel engineer 관점이 조금 더 분명해진다.
- 큰 수식을 그대로 계산하는 것이 목표가 아니다
- 메모리에 무엇을 남기고 무엇을 남기지 않을지를 설계해야 한다
- 계산과 메모리 이동을 같이 본다
이 관점이 중요한 이유는, FlashAttention이 단순히 유명한 트릭이 아니라 GPU kernel engineer 사고방식을 매우 잘 보여주는 사례이기 때문이다. 큰 intermediate를 그대로 materialize하지 않고, tile 단위로 계산을 흘려보내며 memory traffic를 줄이는 태도가 핵심이다.
실무에서 중요한 것은 프로파일링이다
여기까지 오면 직감만으로는 한계가 있다. 결국 Nsight 같은 도구나 프레임워크 프로파일러를 통해 실제 병목을 봐야 한다.
예를 들어:
- memory throughput이 얼마나 나오는가
- occupancy는 충분한가
- warp stall 이유는 무엇인가
- kernel fusion 이후 register pressure가 커지지는 않았는가
최적화는 추측이 아니라 측정에 기반해야 한다. GPU 쪽에서는 이 원칙이 특히 더 중요하다.
그래서 Triton을 포함한 실전 최적화는 보통 이렇게 흐른다. 먼저 profiling으로 병목을 보고, 그 다음 Triton이나 CUDA로 구조를 바꾸고, 다시 측정한다. 이 반복이 없으면 최적화는 금방 감에 의존하게 된다.
이 시리즈 이후 어디로 가면 좋을까
GPU 시스템 시리즈를 끝내면 다음 단계는 두 갈래로 이어진다.
첫째, PyTorch internals로 가서 custom CUDA operator와 fused kernel을 실제 프레임워크에 연결하는 길이 있다.
둘째, distributed LLM training으로 가서 여러 GPU와 대규모 학습 시스템으로 확장하는 길이 있다.
실제로 GPU kernel engineer와 distributed training engineer는 완전히 분리된 길이라기보다, 중간에 상당히 많이 만난다. 커널을 모르면 프레임워크 병목을 깊게 보기 어렵고, 분산 학습을 모르면 큰 모델 문맥에서 어떤 최적화가 진짜 중요한지 놓치기 쉽다.
마무리
이 시리즈의 핵심은 GPU를 더 이상 "빠른 장치" 정도로 보지 않는 것이다. GPU는 고유한 실행 모델과 메모리 계층을 가진 계산 시스템이고, 좋은 커널은 그 구조를 이해한 상태에서 나온다.
다음 시리즈에서는 PyTorch 내부 구조로 넘어가서, 이런 커널과 저수준 최적화가 실제 학습 프레임워크 안에서 어떻게 연결되는지 본다.