[cutlass] Blackwell SSD 예제의 TMEM 4중 해제와 조기 해제 버그를 Named Barrier 하나로 고치기

PR 링크: NVIDIA/cutlass#3803 상태: Merged | 변경: +12 / -2

들어가며

Blackwell(SM100)에는 Tensor Memory(TMEM)라는 새로운 on-chip 메모리가 있습니다. tcgen05.mma의 accumulator가 TMEM에 쌓이고, 소비자 warp은 tcgen05.ld로 이를 레지스터로 읽어옵니다. TMEM은 shared memory처럼 자동으로 관리되지 않습니다. CTA가 tcgen05.alloc으로 명시적으로 할당하고 tcgen05.dealloc으로 반납해야 합니다. 즉 수명(lifecycle)을 프로그래머가 직접 관리하는 자원입니다.

이 PR은 PTX 수준 동기화/레이스 검증기(sentryward)로 CUTLASS의 SM100 커널을 분석하다 발견한 두 가지 문제를 다룹니다.

1. Example 112 (Blackwell SSD): TMEM이 네 번 해제되고, 그것도 PreInter warpgroup의 마지막 TMEM 읽기보다 먼저 해제될 수 있음 → 이번 PR에서 수정
2. SM100 grouped GEMM: 이전 TMA load가 아직 읽고 있을 수 있는 tensor map을 덮어쓰는 문제 → 설계 판단이 필요해 보고만 함

이 글은 실제로 코드가 바뀐 1번을 중심으로 보고, 2번은 일반적 교훈과 함께 짧게 정리합니다.

문제 상황: warp 역할과 TMEM 수명

SSD 커널은 warp specialization 구조입니다. warp 0–3은 각자 역할(MMA 등)을 맡고, warp 4–7은 PreInter, warp 8–11은 PreIntra warpgroup입니다. warp 1이 512 column 전체를 할당하고, 나머지 TMEM 사용자는 TmemAllocBarrier에서 할당 결과를 기다립니다.

문제는 해제 쪽이었습니다. 수정 전 코드는 이렇습니다.

// Before: PreIntra 분기의 끝 (warp 8, 9, 10, 11 모두 실행)
uint32_t free_stage_ptr = storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, ArchTag::kTmemCapacityColumns);

여기에는 두 가지 버그가 겹쳐 있습니다.

버그 1: 네 번 해제

PTX ISA는 .cta_group::1일 때 "one warp from the CTA must perform the allocation and de-allocation"이라고 명시하고, taddr은 "must point to a previous Tensor Memory allocation"이어야 한다고 합니다. 그런데 이 코드는 PreIntra의 4개 warp가 모두 free를 호출합니다. 첫 번째 warp이 해제한 뒤 나머지 세 번은 이미 이 CTA 소유가 아닌 column을 해제하는 셈입니다. 그 사이 같은 SM의 다른 CTA가 그 column을 할당했다면, 남의 할당을 해제해버릴 수 있습니다.

버그 2: 마지막 읽기 전에 해제

마지막 MMA가 끝나면 MMAInter는 pipeline_inter와 pipeline_acc를 동시에 commit합니다. 이로써 PreInter와 PreIntra가 동시에 풀려납니다.

PreInter (warps 4–7)                    PreIntra (warps 8–11)
L953  consumer_wait(pipeline_inter)     epi L368  consumer_wait(pipeline_acc)
      (scheduler 때문에 지연)           epilogue: tcgen05.ld, math, TMA store
                                        L492  tcgen05.dealloc (warp 8)
                                        L492  tcgen05.dealloc (warp 9,10,11)
L971  tcgen05.ld inter accumulator  <-- 이미 해제된 TMEM을 읽음

PreInter는 마지막 TMEM load 이후 다른 warp가 기다리는 barrier에 전혀 arrive하지 않습니다. 남은 작업은 shared memory에서 나가는 TMA store(store_p)뿐이기 때문입니다. 결과적으로 PreInter의 마지막 load와 PreIntra의 해제 사이에 어떤 happens-before 관계도 없습니다. 이 load 결과가 최종 state가 되므로, 다른 CTA가 먼저 column을 가져가면 에러 없이 틀린 결과가 저장됩니다.

코드 분석

sm100_ssd_kernel_tma_warpspecialized.hpp — 해제용 barrier 선언

if (warp_category == WarpCategory::MMAInter || warp_category == WarpCategory::PreIntra || warp_category == WarpCategory::PreInter) {
  tmem_allocation_result_barrier.arrive_and_wait();
}
// After: 추가된 부분
// The PreInter and PreIntra warpgroups read TMEM until their last tile. Both arrive here before
// TMEM is freed. Reuses the allocation barrier id, which every participant has passed by then.
arch::NamedBarrier tmem_deallocation_barrier(NumEpilogueThreads + NumEpilogueThreads,
                                             cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);

참여 스레드 수는 PreInter + PreIntra 두 warpgroup(각 128 스레드)입니다. 흥미로운 점은 barrier id를 새로 만들지 않고 TmemAllocBarrier를 재사용한다는 것입니다. 하드웨어 named barrier는 16개뿐이라 귀한 자원입니다. 할당 barrier는 커널 초반에 모든 참여자가 이미 통과했으므로, 커널 끝에서 같은 id를 다른 스레드 수로 다시 써도 충돌이 없습니다. Example 77(MLA)의 해제 barrier도 같은 패턴을 씁니다.

PreIntra 분기 — 기다린 뒤, 한 warp만 해제

// Before
uint32_t free_stage_ptr = storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, ArchTag::kTmemCapacityColumns);

// After
// Wait for the PreInter warpgroup's last TMEM read, then free TMEM from a single warp.
tmem_deallocation_barrier.arrive_and_wait();
if (warp_idx == static_cast<int>(WarpCategory::PreIntra)) {
  uint32_t free_stage_ptr = storage.tmem_base_ptr;
  tmem_allocator.free(free_stage_ptr, ArchTag::kTmemCapacityColumns);
}

두 버그를 한 번에 해결합니다.

• arrive_and_wait(): PreInter가 도착할 때까지 해제를 미룹니다 (버그 2).
• warp_idx == WarpCategory::PreIntra 조건: PreIntra의 첫 warp(warp 8)만 해제합니다 (버그 1).

PreInter 분기 — 마지막 읽기 후 arrive

  for (; tile_scheduler.is_valid(); ++tile_scheduler) {
    ...
  }
  // After: 추가된 부분
  // pre_inter waits for its last TMEM load (tcgen05.wait::ld), so TMEM can now be freed.
  tmem_deallocation_barrier.arrive();

PreInter는 wait이 아닌 arrive()만 합니다. 기다릴 이유가 없는 쪽은 신호만 보내고 바로 진행하는 producer-consumer 형태의 비대칭 barrier입니다. 내부적으로 tcgen05.wait::ld로 load 완료를 이미 기다렸으므로, arrive 시점에 레지스터로의 읽기가 끝나 있음이 보장됩니다.

왜 이 두 warpgroup만인가

• PreIntra 내부 warp 간 순서는 추가로 필요 없습니다. epilogue가 매 TMEM load마다 wait하고, 이후 4개 warp 전체를 동기화하기 때문입니다.
• MMA 결과는 MMA pipeline을 통해 이 두 warpgroup이 소비하므로, 둘 중 누가 arrive하든 마지막 MMA는 이미 완료된 상태입니다. MMA warp를 barrier에 넣을 필요가 없습니다.

대안으로 SM100 GEMM 커널처럼 할당한 warp(warp 1)가 해제하도록 옮길 수도 있지만, 그 역시 같은 barrier와 warp 1의 wait이 추가로 필요합니다. 현재 방식이 변경 범위가 더 작습니다.

검증 결과

검증기 기준:

PTX	결과
main	invalid_lifecycle: tmem_multiple_dealloc
warp 8만 해제, barrier 없음	"tcgen05.load must complete before the allocation is deallocated"
이 PR	race_free

B200(CUDA 13.2), ptxas의 TMEM guardrail(-Xptxas -g-tmem-access-check) 사용 시:

Build	실패율	시간(200 iter)
main	16 / 21	0.1804 ms
이 PR	0 / 21	0.1823 ms

오버헤드는 약 1%, 레지스터 사용량은 168개로 동일합니다. 단계별 결과가 흥미로운데, "한 warp만 해제"로 반복 해제를 없애면 guardrail 실패는 사라지지만 검증기는 여전히 조기 해제 레이스를 잡아냅니다. 실하드웨어에서 재현되지 않는 버그도 형식 검증으로는 드러난다는 좋은 사례입니다. 실제로 PR 작성자도 조기 읽기는 하드웨어에서 관찰하지 못했다고 밝힙니다. PreIntra가 epilogue 전체를 도는 동안 PreInter는 load 하나만 하면 되고, 199KB shared memory 때문에 SM에 두 번째 SSD CTA가 올라갈 수도 없기 때문입니다. 하지만 "지금 운 좋게 안 터진다"는 정확성의 근거가 될 수 없습니다.

보너스: Grouped GEMM의 tensor map 재사용 문제

Persistent grouped GEMM은 group이 바뀔 때 global memory의 tensor map을 tensormap.cp_fenceproxy로 덮어씁니다. 그 전에 tma_desc_commit_group() / tma_desc_wait_group()(= cp.async.bulk.wait_group.read 0)을 호출하지만, mainloop의 TMA load는 mbarrier 기반 completion을 사용합니다. PTX ISA상 commit_group은 bulk_group 기반 completion을 쓰는 연산만 묶으므로, 이 wait은 실제로 아무 load도 기다리지 않을 수 있습니다. 즉 이전 group 마지막 tile의 load가 새 group의 주소/stride로 수행될 가능성이 있습니다. 이 패턴은 SM90~SM120의 16개 collective에 퍼져 있습니다.

제안된 해결책은 (1) group 전환 시 모든 stage 소비를 기다리기(pipeline bubble 발생), (2) SM당 tensor map 두 개를 번갈아 쓰기(double buffering)입니다. PR은 "현재 하드웨어에서 wait_group.read가 mbarrier 모드 load의 tensor map 읽기도 커버하는가"를 NVIDIA에 질문하며 설계 결정을 맡겼습니다.

왜 이게 좋은가 — 일반적 교훈

1. 명시적 수명 자원은 "마지막 사용자"를 명시적으로 기다려야 한다. TMEM, tensor map처럼 비동기 하드웨어 유닛이 접근하는 자원은 "내 warpgroup이 끝났다"가 "모두가 끝났다"를 의미하지 않습니다.
2. 단일 소유권 규칙을 코드로 강제하라. if (warp_idx == ...) 한 줄이 PTX 스펙 위반을 막습니다. warpgroup 단위로 실행되는 코드에 warp 단위 연산을 넣을 때 특히 주의해야 합니다.
3. 비대칭 barrier(arrive vs arrive_and_wait)로 비용을 최소화하라. 기다릴 필요 없는 쪽은 arrive만 하므로 오버헤드가 ~1%에 그칩니다.
4. 희소한 하드웨어 자원(named barrier id)은 수명이 겹치지 않으면 재사용하라.
5. completion mechanism을 혼동하지 말라. mbarrier 기반 TMA와 bulk_group 기반 TMA는 서로의 wait로 보호되지 않습니다.
6. 형식 검증은 하드웨어 테스트가 놓치는 레이스를 잡는다. guardrail과 검증기를 함께 쓴 이 PR의 접근은 GPU 커널 개발 전반에 적용할 만합니다.

참고 자료

• CUTLASS GitHub Repository — 이 PR이 적용된 NVIDIA CUTLASS 저장소
• PTX ISA — Parallel Thread Execution — tcgen05.alloc/dealloc, tcgen05.ld, cp.async.bulk.commit_group/wait_group, tensormap.cp_fenceproxy 명세
• CUDA C++ Programming Guide — warp, named barrier 및 TMA 관련 프로그래밍 모델
• sentryward (gpu-verifier) — 이 버그들을 찾아낸 PTX 수준 동기화/레이스 검증기

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.