Mosaic GPU

강의의 출발점은 역할 변화다. PyTorch 가 처음 태어났을 때 ML 라이브러리의 목표는 generality + acceptable performance 였다 — 30% 손실은 받아들일 만 했다. 지금은 거꾸로다. transformer 가 거의 모든 모델이 됐고, scaling 이 비용의 대부분을 결정하면서 peak performance 가 먼저, 사용성은 그 다음이 됐다. Mosaic GPU 는 그 반전된 역할 위에서 다시 그린 DSL.

Adam 이 강의 도입에서 명시한 디자인 좌표는 세 개로 압축된다.

1. compiler 가 못 따라잡는 자리 — Ampere → Hopper → Blackwell 로 갈수록 프로그래밍 모델이 매 세대 다르다. 일반 컴파일러가 이 변화를 따라잡으려면 시간이 걸리고, 그 시간 동안 사용자는 peak 의 60–70% 에서 머문다.
2. C++ template metaprogramming 의 천장 — CUTLASS / ThunderKittens 가 가는 길. Adam 이 본 한계는 “error message 가 너무 거칠다, 디버깅이 사실상 불가능에 가까운 자리들이 있다”. 같은 metaprogramming 이 Python tracing 으로 풀리면 훨씬 깔끔하다.
3. JAX 사용자에게 Triton 이 비어 있는 자리 — Triton 은 PyTorch 와 함께 자라며 거의 표준이 됐다. JAX 사용자는 같은 진입점이 없었다. Mosaic GPU 가 그 자리.

강의의 실무적 산출은 두 개다 — Pallas 라는 frontend 와 Mosaic GPU 라는 backend, 그 둘의 결합으로 짠 FlashAttention 3 of 150 lines (Hopper 위에서 tensor core 의 70% 이상 활용). 이 결과 한 줄이 강의 중간에 박혀 있고, 나머지는 그 결과를 어떻게 끌어냈는지의 디자인 reasoning 이다.

Adam 이 도입에서 깐 사실 — Triton 의 큰 성공 요인이 integration triviality 였다는 것. PyTorch 에서 그냥 파이썬 함수처럼 import 하고 부른다. Mosaic GPU 도 같은 원칙: JAX 에서, 그리고 PyTorch 에서도 copy 없이 부를 수 있어야 한다.

강의에서 Adam 이 한 번 더 분명히 했다 — "We just want to have a DSL for fast Hopper, and at this point also Blackwell, and we also want to program them from Python." 즉 GPU 의 새 세대를 따라잡을 수 있는 빠르게 진화 가능한 backend 가 그 빈자리.

Mark 가 강의 중 끼어들어 한 질문 — "CUTLASS 는 C++ template 으로, 다른 프로젝트는 string 으로, Triton 과 Mosaic 는 MLIR 로 간다. 차이가 뭔가?" Adam 의 답이 강의의 첫 깊은 자료. 4가지 metaprogramming 패턴을 비교한다.

Adam 의 결론 — GPU 커널은 metaprogramming 이 사실상 모든 곳에 등장한다. tile shape 계산, register tile 분할, warp 별 작업 분배 모두 compile-time 결정이다. Python control flow 가 자연스럽게 unroll/specialize 되는 tracing 모델이 가장 적합하다는 입장.

강의에서 가장 많이 헷갈렸던 자리. Pallas 는 JAX 의 “커널 짜는 frontend” 다. 그리고 그 frontend 는 세 backend 로 lowering 된다.

Adam 이 두 번 강조한 작은 사실 — Pallas 의 “thread” 는 CUDA 의 thread 가 아니다. Pallas 의 thread 한 개는, Mosaic GPU backend 에서는 warp group 한 개에 매핑된다. 이게 Mosaic GPU 의 추상화 레벨 결정의 핵심이다.

# Pallas 의 단일 thread = Mosaic GPU 의 warp group
@pl.kernel
def add_kernel(x_ref, y_ref, z_ref):
    # x_ref, y_ref 는 reference (mutable, shaped)
    # [...] 는 dereference
    z_ref[...] = x_ref[...] + y_ref[...]

jax_add = pl.pallas_call(
    add_kernel,
    out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype),
    grid=(num_blocks,),     # 1D 그리드
)
# 결과: 일반 JAX 함수. jit 가능, vmap 가능, autodiff 가능.

Mosaic GPU 의 디자인 원칙은 추상이 아니다. Hopper 의 4가지 핵심 feature 각각에서 무엇을 자동화하고 무엇을 노출했는지를 보면 분명해진다.

이 패턴이 Mosaic GPU 코드를 짧게 만든다. 강의에서 Adam 이 비교한 — “같은 일을 CUTLASS 로 하면 metaprogramming + template 으로 백 줄 단위가 되는데, Pallas 위에서 한 줄 추가로 끝난다” 는 자리들이 cluster + multicast 같은 부분.

강의 후반부에 깐 첫 큰 예제. matmul 을 Pallas Mosaic GPU 로 짜는 시퀀스를 코드 한 덩이씩 보여준다 — block_spec 으로 tile 을 정의하고, plgpu.emit_pipeline 으로 producer-consumer 자동 생성, plgpu.wgmma 로 tensor core 부른다.

# Mosaic GPU 의 matmul kernel — 강의 시퀀스 정리
def matmul_kernel(a_ref, b_ref, o_ref):
    # a_ref, b_ref 는 SMEM tile (per pipeline stage)
    def compute(_, a_smem, b_smem):
        acc = plgpu.wgmma(zeros, a_smem, b_smem)        # tensor core MMA
        return acc

    acc = plgpu.emit_pipeline(
        compute,
        grid=(K // BLOCK_K,),                            # K-차원 reduction
        in_specs=[
            plgpu.GPUBlockSpec((BLOCK_M, BLOCK_K),
                lambda k: (0, k),
                transforms=plgpu.TilingTransform((64, 32))),
            plgpu.GPUBlockSpec((BLOCK_K, BLOCK_N),
                lambda k: (k, 0),
                transforms=plgpu.TilingTransform((32, 64))),
        ],
    )
    plgpu.copy_smem_to_gmem(acc, o_ref)                  # 결과 store

Adam 이 강조한 메타 메시지 — “정말 새로운 개념은 거의 없다. block_spec 은 Triton 의 block pointer 와 의미적으로 같고, transforms 는 SMEM swizzle 의 직접 노출이고, emit_pipeline 은 다단계 prefetching 의 helper 화다.” 친숙한 추상의 재배열이지 새로운 추상의 발명이 아니라는 입장.

강의의 클라이맥스. FlashAttention 3 의 핵심 — 두 consumer warp group이 WGMMA 와 softmax 를 critical section 으로 분리해서, Hopper 의 tensor core 회로와 ALU 회로가 동시에 돌게 만든다 — 가 Mosaic GPU 위에서 자연스럽게 표현된다.

강의 끝에서 Adam 이 강조 — “우리가 제공하는 건 attention-specific 한 helper 가 아니다. emit_pipeline_warp_specialized 는 일반 pipeline 의 helper다. 다른 패턴 (decode, prefill, splash attention) 모두 같은 helper 로 풀린다.”

청중에서 가장 자주 나온 질문. Adam 의 답이 분명했다 — "Triton 은 block-level 추상에 집중. Mosaic 는 warp-group level. 둘 다 가치 있는 자리".

Adam 의 입장 — "Triton 의 block-level 도 한 단계 높다고 본다. 하지만 그 위치에서 정말 잘 작동한다. CUDA thread 단위는 dense linear algebra 에선 너무 fine-grained. 우리는 그 사이의 자리". 즉 Mosaic 가 Triton 을 대체한다고 말하지 않는다.

Adam 이 도입에서 한 번, 끝에서 한 번 강조 — "early PyTorch bindings 가 있다. Mosaic GPU 커널을 PyTorch 에서 부를 때 copy 가 일어나지 않는다".

실용적 함의가 크다. PyTorch 의 학습 루프 안에서 — JAX 를 import 할 필요 없이 — FlashAttention 3 급의 Mosaic GPU 커널을 한 줄로 갈아 끼울 수 있다는 뜻. JAX 가 없는 환경에서도 Mosaic GPU 의 출력물 (Hopper 자원을 다 끌어쓰는 PTX) 를 활용할 수 있는 길.

강의를 다시 열었을 때 가장 빨리 복원해야 하는 사실들과 — 직접 손으로 봐야 할 코드 자료들.

손에 새기기 — 실습 시퀀스

1. Pallas add kernel 부터 — JAX 환경에서 pl.pallas_call 로 elementwise add 커널을 짠다. x_ref[...] dereference 의미가 손에 잡힐 때까지.
2. 1D grid 와 program_id — grid=(num_blocks,) 로 grid 를 만들고 pl.program_id(0) 로 자기 위치 조회. 각 thread 가 입력의 일부 슬라이스만 처리하도록.
3. Mosaic GPU backend 명시 — plgpu.kernel 또는 compiler_params 로 Mosaic GPU backend 활성화. 같은 코드가 Triton vs Mosaic 양쪽으로 lowering 되는지 확인.
4. WGMMA matmul — plgpu.wgmma 와 GPUBlockSpec + TilingTransform 으로 minimal matmul. 베이스라인은 jnp.matmul.
5. emit_pipeline — 같은 matmul 을 K-차원 producer-consumer 로 묶기. num_stages 를 1, 2, 3 으로 sweep — speedup 곡선이 어떻게 변하는지.
6. FlashAttention 3 코드 읽기 — JAX repo 의 attention 예제. 150줄 안에서 어떤 줄이 producer/consumer 인지, barrier 가 어디 있는지 표시.
7. PyTorch 인터롭 — Mosaic 커널을 PyTorch tensor 로 호출. torch.cuda.synchronize() 로 정확한 latency 측정.

L042 의 추상 결정 (Pallas frontend, warp-group 단위, manual knob 디자인) 이 시리즈 안에서 어떻게 다시 호출되는지 묶어둔다.

학습 노트로 정리하면서 의도적으로 비워둔 자리들 — 강의 안에서 부분적으로만 등장한 주제, 또는 후속 자료가 더 깊게 다룰 주제.

• Blackwell 대응 시점 — 강의 시점 기준 Mosaic GPU 가 Hopper 위주. Blackwell (sm_100) 의 새로운 Tensor Memory Accelerator 와 5th gen tensor core 를 어떻게 노출할지는 이후 자료 확인 필요.
• Triton 과 직접 성능 비교 — 강의에서 “Triton 도 빨라요” 정도로만 언급. 같은 attention 패턴에서 두 backend 가 정확히 어디서 갈리는지의 microbenchmark 는 외부 자료.
• warp-specialized helper 의 일반성 — strict producer-consumer 외의 패턴 (예: ring all-reduce in compute) 으로 확장 가능한지 미확인.
• TMA descriptor reshuffle 의 한계 — 어떤 transpose/swizzle 까지 free 인지, 어디서 추가 kernel 이 필요한지 — TMA 명세에 따라 다름. 직접 PTX 검증.
• PyTorch 인터롭의 실제 성능 — “no copy” 라고 하지만 stream synchronization 비용은 따로 측정 필요. 학습 루프 안에서 graph break 없는 경로인지 확인 필요.
• LLM kernel 자동생성과의 관계 — 청중 질문 중 하나 — “LLM 이 well-known language 로 코드 생성 가능한 시대에 DSL 의 의미는?”. Adam 의 답변은 강의 끝에 짧게. 노트로 따로 정리 필요.