[cutlass] Blackwell FMHA 예제의 잠재 UB 두 가지 수정: warpgroup 단위 setmaxnreg와 TMEM 해제 전 tcgen05.wait::ld

PR 링크: NVIDIA/cutlass#3798 상태: Merged | 변경: +24 / -9

들어가며

CUTLASS의 example 77은 Blackwell(SM100)용 FMHA, MLA forward, generation 커널을 담고 있습니다. 이 커널들은 warp specialization으로 짜여 있습니다. 각 warp가 Softmax, Correction, MMA, Load, Epilogue 같은 역할을 하나씩 맡고, setmaxnreg로 warpgroup마다 레지스터를 재분배합니다. 또 S(=Q·Kᵀ) 결과와 row statistics를 TMEM(Tensor Memory) 에 두고, 여러 warp가 pipeline barrier로 소유권을 넘겨받으며 사용합니다.

이 PR은 PTX 출력을 PTX ISA 9.4 스펙과 대조하다 발견한 두 가지 잠재적(latent) 정의되지 않은 동작(UB) 을 고칩니다.

1. warpgroup 3의 setmaxnreg 위반: warpgroup 3의 네 warp(MMA, Load, Epilogue, Empty)가 각자 다른 분기에서 setmaxnreg를 호출했습니다. 게다가 Empty warp만 24, 나머지는 32(MLA는 48)로 개수도 달랐습니다.
2. TMEM statistics 해제 전 tcgen05.wait::ld 누락: correction warp가 tcgen05.ld로 읽은 statistics 컬럼을 load 완료를 기다리지 않고 consumer_release했습니다. 이 컬럼은 S0/S1과 aliasing되어 있어서, 다음 producer가 덮어쓴 값을 읽을 수 있었습니다.

작성자는 실제 하드웨어 실패는 확인하지 못했다고 밝혔습니다. 이 PR은 성능 최적화가 아니라 스펙 준수와 정합성(correctness) 을 위한 수정입니다.

코드 분석

1. sm100_fmha_fwd_kernel_tma_warpspecialized.hpp: setmaxnreg를 한 번만

Before: 역할 분기마다 따로 레지스터를 설정했습니다.

static const int NumRegsEmpty = 24;
...
else if (role == WarpRole::MMA) {
  warpgroup_reg_set<NumRegsOther>();
  ...
}
else if (role == WarpRole::Load) {
  warpgroup_reg_set<NumRegsOther>();
  ...
}
else if (role == WarpRole::Epilogue) {
  warpgroup_reg_set<NumRegsOther>();
  ...
}
else if (role == WarpRole::Empty) {
  warpgroup_reg_set<NumRegsEmpty>();
  /* no-op, donate regs and exit */
}

PTX ISA §9.7.21.5는 "같은 setmaxnreg 명령을 warpgroup의 모든 warp가 실행해야 한다"고 규정합니다. .aligned 한정자에 대해서도 "조건부 코드에서는 warpgroup 내 모든 스레드가 조건을 동일하게 평가할 때만 써야 한다"고 명시합니다. 기존 코드는 두 규칙을 모두 어깁니다. 개수가 24와 32로 다르고, 개수가 같은 세 호출도 warp마다 갈리는 분기 아래의 서로 다른 명령입니다. 실제로 ptxas는 이를 네 개의 USETMAXREG.DEALLOC으로 남겼습니다.

After: 역할 dispatch 전에 warpgroup-uniform 조건으로 한 번만 실행합니다.

// All warps of a warpgroup must execute the same setmaxnreg, so warpgroup 3
// (MMA, Load, Epilogue, Empty) sets its register count once, before the role dispatch.
if (warp_idx / 4 == 3) {
  warpgroup_reg_set<NumRegsOther>();
}

warp_idx는 canonical_warp_idx_sync()에서 오므로 같은 warpgroup의 네 warp는 조건을 똑같이 평가합니다. PTX에서는 and.b32 ..., -4; setp.ne.s32 ..., 12로 감싼 setmaxnreg.dec 32 하나로 내려갑니다.

이 수정이 왜 중요한지는 레지스터 풀 계산으로 보입니다. Softmax warp의 setmaxnreg.inc는 "CTA 레지스터 풀에 충분한 레지스터가 생길 때까지 블록"됩니다.

• launch 시 512 threads × 128 regs = 65,536
• softmax 두 warpgroup이 추가로 필요한 양: 2×128×64 = 16,384
• correction warpgroup이 반납하는 양: 128×32 = 4,096
• 나머지 12,288은 warpgroup 3의 반납에 의존합니다.

warpgroup 3의 동작이 UB라면 이 반납도 보장되지 않습니다. 최악의 경우 softmax가 무한 대기할 수 있습니다. 수정 후 예산은 정확히 맞아떨어집니다.

2×128×192 + 128×96 + 128×32 = 65,536   (FMHA)
2×128×184 + 128×96 + 128×48 = 65,536   (MLA)

대가로 Empty warp는 24가 아닌 32(또는 48)개로 내려갑니다. 하지만 이 warp는 바로 종료하고, 기존 24는 오히려 256개(MLA는 768개)를 놀리고 있었습니다. 레지스터 수는 warpgroup 단위 개념이므로 24를 유지하려면 role layout 자체를 바꿔야 합니다.

2. 세 개의 mainloop correction(): 해제 전 tcgen05.wait::ld

대상은 fwd, MLA fwd, gen mainloop 세 파일이며, 변경 패턴은 동일합니다. 루프 내부 예시는 다음과 같습니다.

Before

  correction_rescale(scale, uint32_t(TmemAllocation::O0));
}

pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state);
++pipeline_s1_c_consumer_state;

After

  correction_rescale(scale, uint32_t(TmemAllocation::O0));
}

// the stats loads (tcgen05.ld) must complete before their TMEM columns are released
cutlass::arch::fence_view_async_tmem_load();
pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state);
++pipeline_s1_c_consumer_state;

루프 이후의 최종 해제도 같은 방식으로 고쳤습니다.

copy(tiled_tmem_loadv, tTMEM_LOADVtS1, tTMEM_LOADVrS);

cutlass::arch::fence_view_async_tmem_load();   // 추가
pipeline_s1_c.consumer_release(pipeline_s1_c_consumer_state);

fence_view_async_tmem_load()는 tcgen05.wait::ld를 내보냅니다. tcgen05.ld는 비동기라서 명령이 issue된 시점과 데이터가 레지스터에 도착한 시점이 다릅니다. §9.7.18.6.4.4는 "비동기 tcgen05 명령을 issue한 producer는 consumer와 동기화하기 전에 명시적으로 완료를 기다려야 한다"고 요구합니다. 예시 1도 tcgen05.ld → tcgen05.wait::ld → ... → mbarrier.arrive 순서입니다.

세부 사항은 다음과 같습니다.

• 루프 내부 해제: 그 사이에 correction_rescale이 실행되면 tcgen05.st와 tcgen05.wait::st가 간접적으로 순서를 보장합니다. 하지만 scale이 모두 1이어서 rescale을 건너뛰면 아무것도 순서를 보장하지 않습니다.
• 최종 해제: load 직후 바로 release합니다. 이 경우는 PersistentTileScheduler나 CausalPersistentTileScheduler에서 특히 위험합니다. release 직후 다음 타일의 softmax와 MMA warp가 S0/S1에 쓰기 시작하기 때문입니다. SASS에서도 최종 LDTM.x2 바로 뒤에 SYNCS.ARRIVE가 붙어 있었습니다.

왜 이게 좋은가

비용

작성자가 ptxas -v로 head-dim 64 fp16 커널 두 개를 비교한 결과입니다.

• IndividualTileScheduler: 128 regs, spill store/load 8/12 bytes로 변화 없음
• PersistentTileScheduler: 128 regs 유지, spill load만 44 → 48 bytes

루프 내부에서는 load한 값을 release 전에 이미 소비하므로 wait 비용이 사실상 0입니다. 최종 해제에서만 load latency만큼 기다립니다. release를 첫 사용 이후로 옮기면 이 비용도 없앨 수 있지만, softmax 진행 타이밍이 바뀌므로 PR은 최소 변경을 택했습니다. 실행 시간은 측정하지 않았습니다.

일반적 교훈

1. "지금 돌아간다"는 "정확하다"가 아니다. 현재 SASS가 우연히 안전하더라도, 컴파일러 버전이나 스케줄링이 바뀌면 UB는 언제든 드러날 수 있습니다. PTX를 스펙과 기계적으로 대조하는 검증은 유효한 방법입니다.
2. warpgroup 단위 명령은 warpgroup 단위로 호출하라. setmaxnreg, wgmma 같은 .aligned 명령은 role 분기 밖에서 uniform 조건으로 실행해야 합니다. SM90 CUTLASS의 producer warpgroup이 이미 이 패턴을 씁니다.
3. 비동기 load와 버퍼 소유권 이전 사이에는 반드시 완료 대기를 넣어라. 공유 버퍼를 aliasing하는 파이프라인에서 release는 "이제 써도 된다"는 신호입니다. 그 전에 내가 읽기를 끝냈다는 보장이 있어야 anti-dependency(WAR) hazard를 막을 수 있습니다.
4. 레지스터 예산은 산수로 검증하라. inc와 dec의 합이 launch 예산과 정확히 맞는지 확인하면, 데드락 위험과 낭비를 함께 찾을 수 있습니다.

PR은 범위 밖의 관련 이슈도 정리해 두었습니다. gen 커널은 warpgroup당 하나의 count로는 예산(69,632 > 65,536)에 맞지 않습니다. 2bwd 커널의 kEmpty 값도 의심스럽습니다. 이런 부분까지 함께 남겨 두어 유지보수자가 판단하기 쉽습니다.

참고 자료

• CUTLASS GitHub Repository — example 77 Blackwell FMHA가 포함된 저장소
• PTX ISA Documentation — setmaxnreg(§9.7.21.5) 및 tcgen05.ld/tcgen05.wait(§9.7.18) 규정을 담은 공식 문서
• CUDA C++ Programming Guide — warp/warpgroup 실행 모델 공식 가이드

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