A learning journey: CUDA, Triton, Flash Attention

Umar Jamil 의 강의는 알고리즘 강의가 아니다 — 한 사람이 attention 을 바닥부터 자기 손으로 짤 때 어디에 부딪혔는가 의 회고에 가깝다. 그 회고가 의미 있는 이유는, 같은 길을 가려는 사람이 거의 똑같은 자리에 부딪히기 때문이다.

강의가 던지는 두 개의 질문.

1. CUDA 와 Triton 사이의 인지적 거리 — 학습자가 한쪽에서 다른쪽으로 넘어갈 때 무엇을 새로 배우고 무엇을 버려야 하는가.
2. Flash Attention 이 “어려운” 진짜 이유 — 알고리즘이 복잡해서가 아니라 메모리 계층과 수학적 안정성과 autograd 회계 가 동시에 얽히기 때문이다.

이 노트는 원본 transcript 가 실패한 강의를 도메인 지식과 공개된 자료(Tri Dao 의 Flash Attention 논문 1·2·3, Umar 의 다른 YouTube 비디오, OpenAI Triton 튜토리얼)로 재구성한 학습 노트다. 강의 안에서 직접 확인되지 않은 주장은 본문에 “원본 영상 확인 필요” 표시로 남겨둔다.

강의의 실질적 도착점은 — Tri Dao 의 Flash Attention 2 를 Triton 튜토리얼 형태로 다시 짠 코드를 한 줄씩 읽을 수 있는 상태다. 그 자리에 도착하기 위한 사다리가 §02 부터 §07 까지에 깔려 있다.

CUDA 를 처음 짜본 사람의 진술은 거의 일치한다 — 문법이 어려운 게 아니라 “내가 thread 한 명이라면 지금 무엇을 하고 있나” 의 1인칭 사고로 넘어가는 게 어렵다. PyTorch 의 batched 연산 사고와 정반대다.

// matmul tiled — CUDA 의 인지적 어려움이 한꺼번에 모이는 자리
__global__ void matmul_tiled(float* A, float* B,
                            float* C, int N) {
  __shared__ float As[16][16];
  __shared__ float Bs[16][16];

  int row = blockIdx.y * 16 + threadIdx.y;
  int col = blockIdx.x * 16 + threadIdx.x;

  float acc = 0.0f;
  for (int t = 0; t < N/16; ++t) {
    As[threadIdx.y][threadIdx.x] =
        A[row*N + t*16 + threadIdx.x];
    Bs[threadIdx.y][threadIdx.x] =
        B[(t*16 + threadIdx.y)*N + col];
    __syncthreads();

    for (int k = 0; k < 16; ++k)
      acc += As[threadIdx.y][k] * Bs[k][threadIdx.x];
    __syncthreads();
  }
  C[row*N + col] = acc;
}

CUDA 의 첫 단계에서 학습자가 보통 6~8주 정도를 보낸다. Umar 도 같은 시기를 거쳤음을 자기 채널의 다른 영상에서 여러 번 언급한다. 이 단계의 끝에 도착하면 — 그제야 같은 알고리즘을 다르게 짜면 다르게 빠를 수 있다 의 직관이 잡힌다.

Triton 의 가장 큰 인지적 변화는 thread 가 사라진다는 점이다. 코드 안에서 thread 단위 indexing 을 직접 쓰지 않는다. 대신 tile (작은 행렬 블록) 을 통째로 다룬다 — tl.load, tl.dot, tl.store. compiler 가 thread 분배를 알아서 한다.

# Triton matmul — tile 단위로 사고
@triton.jit
def matmul_kernel(A, B, C, M, N, K,
                  stride_am, stride_ak,
                  stride_bk, stride_bn,
                  stride_cm, stride_cn,
                  BLOCK_M: tl.constexpr,
                  BLOCK_N: tl.constexpr,
                  BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(A + offs_m[:,None]*stride_am
                       + (k+offs_k)[None,:]*stride_ak)
        b = tl.load(B + (k+offs_k)[:,None]*stride_bk
                       + offs_n[None,:]*stride_bn)
        acc += tl.dot(a, b)
    tl.store(C + offs_m[:,None]*stride_cm
                + offs_n[None,:]*stride_cn, acc)

Flash Attention 이 풀려고 하는 문제는 두 개다 — (1) attention matrix 가 N² 메모리를 먹는다는 점, (2) softmax 가 row 전체를 알아야 정규화된다는 점. 둘이 합쳐지면 — 큰 matrix 를 HBM 에 한 번 쓰고 다시 읽고, 또 쓰고 다시 읽는 패턴이 강제된다.

그런데 “SRAM 위에서 끝까지 돌게 한다” 가 실제로 만나는 첫 벽이 — softmax 의 정규화는 row 전체의 max 와 sum 을 안 다음에야 끝난다는 점이다. naive 하게는 (1) row 전체를 한 번 읽어서 max 구하고, (2) 다시 읽어서 exp(x-max) 의 sum 구하고, (3) 또 다시 읽어서 정규화한다 — 같은 row 를 세 번 읽는다. 한 tile 안에서 끝낼 수 없다. 이 자리를 푸는 게 online softmax다.

online softmax 의 아이디어는 단순하다 — 지금까지 본 max 와 지금까지의 (rescaled) sum 만 들고 다닌다. 새 tile 의 max 가 들어오면 둘 다 보정한다. 이게 numerically stable 하게 동작한다는 점이 50년 전부터 알려져 있었다 (Milakov & Gimelshein 2018 의 재정리, 더 거슬러 올라가면 LogSumExp 의 stable 계산).

이 자리에서 강의의 한 가지 미세한 점 — “naive softmax 가 fp16 에서 overflow 나는 이유” 는 max-shift 가 빠진 형태이고, 그 보정 자체가 online softmax 의 식과 같다는 점. PyTorch 의 F.softmax 도 내부적으로 max-shift 를 한다. Umar 가 강의에서 이 점을 강조했을 가능성이 높다 (원본 영상 확인 필요).

Flash Attention 의 코드를 한 줄씩 읽으면 — 알고리즘을 자체로 새로 발명한 게 아니라, 같은 attention 을 “tile 사이로 무엇을 carry 할 것인가” 의 관점에서 다시 짠 것이라는 사실이 드러난다. 이 사고가 Triton 의 tile 추상과 정확히 맞물린다.

# Flash Attention 2 forward 의 한 query block (의사 코드)
def fa2_forward_block(q_block, K, V, scale):
    # SRAM 에 살아있는 state
    m = -inf                  # running max
    l = 0                      # running sum
    O = 0                      # output

    for kv_block in kv_blocks_in(K, V):
        s = q_block @ kv_block.k.T * scale
        m_new_local = s.max(dim=-1)
        m_new = max(m, m_new_local)
        alpha = exp(m - m_new)
        l = l * alpha + exp(s - m_new).sum(dim=-1)
        O = O * alpha + exp(s - m_new) @ kv_block.v
        m = m_new

    return O / l, m, l       # (m,l) 은 backward 용

이 코드의 모든 변수는 SRAM 에 들어가는 작은 사이즈다. BLOCK_M = 64 ~ 128, head_dim ≤ 128 정도면 한 block 의 q, kv tile, m, l, O 가 모두 register/shared 에 들어간다. HBM 왕복은 query block 한 번 + 모든 kv block 의 한 번씩.

forward 가 풀려도 backward 가 어렵다. 왜냐하면 attention matrix S, P 가 forward 에서 HBM 에 안 저장됐기 때문이다 — 그게 알고리즘의 포인트였으니까. 그러면 backward 는 어떻게? 다시 계산한다.

강의 자체가 transcript 실패라 — 같은 길을 따라가려는 사람을 위해 자료 목록을 정리해둔다. Umar 의 채널이 가장 좋은 시작점이고, Tri Dao 의 두 paper 가 reference 다.

CUDA → Triton → FA 의 학습 시퀀스에서 학습자들이 똑같이 막히는 자리들. Umar 의 강의가 이 자리들을 명시적으로 짚었는지는 확실하지 않지만 (원본 확인 필요), 학습 노트로서 정리해둘 가치가 있다.

1. CUDA 의 “thread 1인칭” 사고 진입에 6주 이상 — 빨리 갈 수 없다. PMPP 책을 직접 풀어야 한다. 강의만으로는 안 됨.
2. Triton 으로 빨리 넘어가려는 유혹 — Triton 만 배우면 “왜 이 설정이 빠른가” 의 직관이 안 잡힌다. CUDA 의 SRAM/bank/coalesce 의미를 손에 들고 Triton 으로 가야 진짜 이해.
3. online softmax 의 식 한 번에 안 잡힘 — 종이 위에서 N=4 의 작은 예제로 한 번 손계산 해보는 게 직관 잡는 가장 빠른 길.
4. FA1 코드 → FA2 코드 차이를 “자세히 안 봄” — 같은 알고리즘처럼 보이지만 outer 루프가 다름. 같이 두고 diff 해야 의미 보임.
5. backward 식 직접 안 풀어봄 — autograd 가 자동이라서 안 풀어도 된다고 착각. 실제로 FA backward 의 reference 코드는 식을 알고 봐야 짤리는 자리가 많음.
6. Triton autotuner 결과를 “공식 답” 으로 받아들임 — 사실 search space 가 좁아서 항상 최적이 아님. 직접 sweep 해본 사람이 더 빠른 설정을 찾는 경우가 종종 있음.
7. “성능 측정” 할 때 워밍업, sync 빼먹음 — L001 의 timing 패턴 (CUDA Event + warmup + sync) 을 항상 쓸 것.
8. head_dim 이 큰 경우(>128) 의 분해 차이 모름 — register pressure 때문에 같은 코드가 안 컴파일되거나 occupancy 가 무너짐. 이 자리를 위해 CUTLASS / FA3 가 개입.

강의의 detail 보다 — 이 노트가 의미를 지니려면 손에 박아야 할 짧은 사실들.

손에 새기기 — 실습 시퀀스

1. PMPP 1–5장 풀이 — vector add, matmul naive, matmul tiled 를 직접. L002 와 같은 자리.
2. Triton 의 첫 matmul — OpenAI tutorial 03 을 그대로 베껴 짠 다음, BLOCK_M/BLOCK_N/BLOCK_K 를 sweep.
3. online softmax 종이 위에서 풀기 — N=8 의 row 한 줄로 max-shift, running update 손계산. step by step.
4. FA1 의 Algorithm 1 손코딩 — paper 의 의사코드를 그대로 PyTorch 로 짠다 (느려도 됨). 정답 검증 위주.
5. Triton tutorial 06 정독 — OpenAI 의 06-fused-attention.py 를 한 줄씩 주석.
6. FA1 vs FA2 코드 diff — Tri Dao 의 두 버전을 같이 두고 outer 루프 차이 직접 비교.
7. backward 식 직접 풀기 — paper appendix 따라 dQ, dK, dV 를 손으로. autograd 안 쓰고.
8. NCU 로 한 번 떠보기 — 자기가 짠 Triton FA 의 occupancy 와 SRAM 사용량을 직접 측정.

L050 의 학습 시퀀스가 시리즈의 어느 강의들과 직접 맞물리는지.

강의 자체의 transcript 가 실패해서, 본 노트는 도메인 지식과 공개된 자료로 재구성됐다. 다음에 영상을 직접 다시 본다면 확인할 자리들.

• Umar 가 강의에서 직접 짠 코드의 정확한 출발점 — Triton tutorial 06 을 그대로 시작점으로 썼는지, 아니면 자기 버전을 새로 짰는지. 원본 영상 확인 필요.
• CUDA → Triton 전환의 결정적 자리 — 강의에서 “여기서부터 CUDA 안 쓰기로 했다” 의 명시적 자리가 있었는지.
• backward 식 손유도 여부 — 강의 안에서 Algorithm 4 의 식을 실제로 화이트보드에 풀었는지, 아니면 paper 를 인용하고 넘어갔는지.
• FA1 vs FA2 의 차이 강조 — 강의 시점에 따라 FA2 가 막 나왔거나 이미 정착된 시점일 수 있음. 어느 쪽이었는지.
• FA3 / Hopper 디테일 — 강의 시점에 FA3 가 발표되었는지 — 발표되었다면 언급 여부.
• PMPP 책에 대한 추천 여부 — Umar 가 PMPP 를 명시적으로 권했는지, 아니면 다른 학습 자료를 추천했는지.
• NCU/profiler 사용 여부 — “자기 코드를 측정한” 자리가 강의 안에 들어갔는지. 측정 없이 알고리즘만 다뤘다면 § 09 의 함정 7번이 강의에서 다뤄지지 않은 것.
• Tri Dao 의 코드를 직접 인용했는지 — flash-attention repo 의 어느 파일을 화면에 띄웠는지 본인이 짠 코드와 별도였는지.