본문으로 건너뛰기

[flashinfer] Hopper FP8 GEMM+AllReduce 융합 커널의 Grid Lockstep Throttle: CTA가 앞서 달리지 못하게 해서 최대 1.32배 빨라지기

PR 링크: flashinfer-ai/flashinfer#9164 상태: Merged | 변경: +4466 / -11

들어가며

flashinfer.comm.fp8_rowwise_gemm_allreduce의 backend="cake"는 row-wise scaled FP8(e4m3) GEMM과 2/4/8 rank all-reduce를 하나의 fused kernel로 실행하고, 모든 rank에 BF16 출력을 돌려줍니다. 이전 PR들(#8905, #9101)로 sm_90a(H100) 지원과 per-call six-reducer 레이아웃이 들어왔지만, 출력 타일 수가 매우 큰 호출에서는 성능이 크게 무너졌습니다. 예를 들어 2-rank mlp N=6144, M=131072에서 이전 버전은 35.77 ms, torch._scaled_mm + NCCL은 22.02 ms였습니다.

원인은 persistent kernel의 고질적 문제인 CTA 간 진행 속도 불균형(drift)입니다. 각 CTA가 타일을 차례로 가져가 처리하는데, 타일 수가 수만 개 수준이 되면 일부 CTA는 몇 라운드씩 앞서 나가고 일부는 뒤처집니다. GEMM과 cross-rank reduction이 한 커널에 엮여 있으면, 앞선 CTA는 아직 다른 rank/CTA가 채우지 않은 데이터를 기다리게 되고, raster group 단위로 맞춰 둔 L2 locality도 깨집니다.

이 PR은 sm_90a에만 grid lockstep throttle을 도입합니다. 핵심 규칙은 단순합니다.

각 CTA의 producer는, 어떤 CTA도 grid 전체가 시작한 타일 수 기준으로 2 persistent round 이상 앞서지 않을 때만 새 타일을 시작한다.

코드 분석

1. 새로 생성된 throttled 프로그램의 binding (hopper/..._92002e9c22c1d8ea9829_binding.cu)

이 PR은 Cake source exporter가 생성한 두 개의 새 sm_90a 프로그램(four-reducer 레이아웃용, six-reducer 레이아웃용)을 추가합니다. 커널 시그니처를 보면 fused GEMM+AllReduce의 구조가 그대로 드러납니다.

extern "C" __global__ void kernel_cake_fp8_rowwise_gemm_allreduce_92002e9c22c1d8ea9829(
    const __grid_constant__ CUtensorMap A, const __grid_constant__ CUtensorMap W,
    const __grid_constant__ CUtensorMap OUT, __nv_bfloat16* __restrict__ out_ptr,
    float* __restrict__ sa, float* __restrict__ sw, unsigned int* __restrict__ ctl,
    unsigned int* __restrict__ S_mc, int M, int N, int K, int grid_n, int tiles, int K_tiles,
    int group_m, int C, int C_cap, unsigned long long S_bytes, unsigned int flag_words,
    int32_t pg_world, int32_t pg_rank, CakePeerPointerTable<unsigned int> pg_flags,
    uint8_t* __restrict__ S, CakePeerPointerTable<uint8_t> S_peers,
    unsigned int* __restrict__ flags, CakePeerPointerTable<unsigned int> flags_peers);
  • A, W, OUT은 __grid_constant__ CUtensorMap으로 전달되어 TMA가 직접 사용합니다.
  • sa, sw는 per-row / per-column FP32 scale입니다.
  • ctl은 grid 차원 제어 워드로, lockstep throttle이 "grid가 시작한 타일 수"를 세는 카운터가 여기에 위치합니다.
  • tiles, group_m이 throttle 적용 조건(타일 수, raster group 크기)과 직결되는 인자입니다.
  • S_peers, flags_peers는 다른 rank의 버퍼/플래그 포인터 테이블로, cross-rank reduction에 쓰입니다.

호스트 측에서는 peer table ABI가 디바이스와 정확히 일치하도록 컴파일 타임에 검증합니다.

template <int Capacity>
struct alignas(16) CakePeerPointerTableHost {
  uint64_t ptrs[Capacity];
};
static_assert(sizeof(CakePeerPointerTableHost<8>) == 64,
              "eight-entry peer table CUDA ABI must be 64 bytes");
static_assert(alignof(CakePeerPointerTableHost<8>) == 16,
              "peer table CUDA ABI must be 16-byte aligned");

커널 파라미터로 by-value 구조체를 넘길 때 host/device 레이아웃이 어긋나면 디버깅이 매우 어렵습니다. static_assert로 이를 빌드 시점에 막는 것은 좋은 방어적 설계입니다.

TMA descriptor 인코딩도 동일 패턴으로 생성됩니다. FP8 입력은 UINT8, 128x128 box, 128B swizzle, L2 256B promotion을 사용합니다.

uint32_t box_dim[2] = {128u, 128u};
...
CUresult r =
    cuTensorMapEncodeTiled(&tm, CU_TENSOR_MAP_DATA_TYPE_UINT8, 2, const_cast<void*>(tensor_base),
                           global_dim, global_strides, box_dim, elem_strides,
                           CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B,
                           CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);

출력 OUT은 BF16, 64x128 box, L2_PROMOTION_NONE입니다. 출력은 재사용되지 않으므로 L2 promotion을 끄는 선택이 자연스럽습니다. 또한 encoder_context_device를 static thread_local로 캐싱해 primary context 바인딩(cudaSetDevice)을 스레드·디바이스당 한 번만 수행합니다.

2. 디바이스 측 throttle 동작 (PR 설명 기준)

커널 본문은 diff에서 잘렸지만, PR 설명은 동작을 명확히 정의합니다.

  • producer warp가 새 타일을 시작하기 전, grid 전역 "begun tiles" 카운트를 확인합니다.
  • 자신이 다음에 시작할 타일이 grid 진행보다 2 persistent round 이상 앞서면 대기합니다.
  • 이 대기는 bounded이며, cross-rank wait와 같은 cycle budget을 넘으면 trap합니다. 즉 hang 대신 명시적 실패를 택합니다.

2 라운드의 여유를 두는 것이 포인트입니다. 완전한 lockstep(0 라운드)은 매 타일마다 grid barrier에 준하는 비용을 치르지만, 2 라운드 슬랙은 pipeline을 채운 상태로 유지하면서 극단적 drift만 잘라냅니다.

3. 디스패치 규칙 (flashinfer/jit/cake_fp8_rowwise_gemm_allreduce.py)

throttled 프로그램은 다음 조건을 모두 만족할 때만 선택됩니다.

  • raster group 크기 > 1
  • 출력 타일 수 ≥ 12288
  • full-capacity(기본) variant일 것 — 8-rank의 4096-row capacity 호출은 제외

144개 측정 케이스 중 24개만 throttled 프로그램을 타고, 나머지 120개는 이전과 완전히 동일한 프로그램을 실행합니다. 회귀 위험을 구조적으로 최소화한 셈입니다.

리뷰에서 나온 논의

리뷰어 @yzh119는 six-reducer + throttle 조합(_r6_ls2)이 4/8 rank에서 도달 가능한지 물었습니다. 작성자는 throttle이 raster group ≥ 2, 즉 K_l >= 4096을 요구하는 반면 TP4/TP8 six-reducer 규칙은 K_l <= 1536을 요구하므로 도달 불가능함을 보였습니다.

대신 리뷰는 실제로 도달 가능한 구멍을 찾아냈습니다. TP8 + 4096-row capacity + N >= 49152 + K_l >= 4096 호출은 throttle 규칙을 만족하지만 unthrottled g4 프로그램으로 떨어졌습니다. 수정(c0f9a695)은 다음과 같습니다.

  • LEVER_DISPATCH["sm_90a"]["LOCKSTEP"]["variants"]에 full-capacity variant만 등록
  • call_program이 capacity의 variant를 lever_variant에 전달
  • 측정 데이터가 없는 main_ws8_p2p_g4_ls2 route 제거

또한 CPU 테스트가 "호출이 도달 가능한 모든 lever key는 ROUTES에 존재해야 하며, ROUTES 밖의 key는 디스패치 규칙과 모순되어야 한다"를 증명하도록 강화되었습니다. 측정되지 않은 조합은 추가하지 않는다는 원칙도 인상적입니다.

왜 이게 좋은가

성능 수치

  • 전체 144 케이스: 11 FASTER, 111 PARITY, 0 SLOWER, 22 INCONCLUSIVE
  • FASTER 11건은 전부 throttled 케이스, median 1.0855x
  • 4-rank mlp N=6144, M=131072: 17.28 → 13.06 ms (1.325x), torch 대비 1.166x
  • 4-rank mlp N=8192, M=65536: 12.33 → 10.74 ms (1.146x)
  • 8-rank mlp N=8192, M=131072: 18.05 → 15.45 ms (1.168x), torch를 역전(1.022x)
  • 2-rank mlp N=3072, M=131072: 7.27 → 6.48 ms (1.121x)

INCONCLUSIVE로 분류된 2-rank 대형 케이스에서도 절대 시간은 35.77 → 26.89 ms, 63.61 → 47.51 ms로 크게 개선되었습니다. 다만 이 영역은 여전히 torch보다 느려 남은 과제입니다.

일반적 교훈

  1. Persistent kernel에서는 "빨리 가는 CTA"가 오히려 독이 될 수 있다. 동기화 지점이나 공유 캐시 locality가 있으면 drift를 제한하는 것이 throughput을 높입니다.
  2. 완전 동기화 대신 bounded slack. 2 라운드라는 여유는 barrier 비용과 drift 피해 사이의 실용적 절충입니다.
  3. Wait는 반드시 bounded. 무한 spin 대신 cycle budget 후 trap하면 multi-rank 환경의 hang을 진단 가능한 오류로 바꿉니다.
  4. 선택적 디스패치로 회귀 표면 최소화. 효과가 측정된 영역에만 새 프로그램을 적용하고 나머지는 기존 경로를 유지합니다.
  5. 엄격한 벤치마크 방법론. 교차 실행 순서, SM clock floor 검증, 세션 간 편차 기준으로 INCONCLUSIVE를 분리해 노이즈를 성능 향상으로 오인하지 않습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글