[flashinfer] FlashInfer Cake GDP 커널 최적화: 공유 메모리 사용량 절감 및 레지스터 직접 로드

PR 링크: flashinfer-ai/flashinfer#7156 상태: Merged | 변경: +83 / -169

들어가며

최신 GPU 아키텍처(NVIDIA Blackwell, sm_100a/sm_103a)에서 고성능 연산을 수행할 때, 공유 메모리(Shared Memory)는 매우 귀중한 자원입니다. 이번 FlashInfer의 PR은 cake_gated_delta_product_mtp 커널 중 T=1인 split-V 변형에서 공유 메모리 사용량을 6400바이트에서 2304바이트로 획기적으로 줄이는 최적화를 다룹니다. 기존에는 recurrent state를 공유 메모리에 스테이징(staging)한 후 연산했으나, 이를 레지스터로 직접 로드하는 방식으로 변경하여 성능을 개선하고 자원 점유를 최소화했습니다.

코드 분석

1. cake_gdp_f2f0c7abbf10082a6996_kernel.cu의 핵심 변경

가장 중요한 변화는 cp.async를 이용한 공유 메모리 스테이징 로직을 제거하고, float4를 사용하여 글로벌 메모리에서 레지스터로 직접 데이터를 읽어오는 방식입니다.

Before:

#pragma unroll
for (int copy_iter = 0; copy_iter < 2; copy_iter++) {
  // ... (생략) ...
  asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16;" ::"r"(copy_dst), "l"(state + ...));
}
asm volatile("cp.async.commit_group;");
// ... (중략) ...
asm volatile("cp.async.wait_group 0;");
__syncthreads();

After:

{ 
  float4 _v4 = *reinterpret_cast<const float4*>(state + raw_row_a + (long long)k_col0);
  raw_a[0 + 0] = _v4.x; raw_a[0 + 1] = _v4.y; raw_a[0 + 2] = _v4.z; raw_a[0 + 3] = _v4.w;
}

이 변경을 통해 cp.async 관련 동기화 오버헤드와 공유 메모리 버퍼 할당이 제거되었습니다. 또한, SMEM_TOTAL을 비롯한 오프셋 정의가 2304바이트로 최적화되었습니다.

2. cake_jit.py 및 바인딩 설정

커널의 dynamic_smem_bytes 설정값을 업데이트하여 런타임 시 할당되는 공유 메모리 크기를 명시적으로 줄였습니다.

# cake_jit.py
"dynamic_smem_bytes": 2304,

왜 이게 좋은가

1. 공유 메모리 압박 완화: 공유 메모리는 GPU 내에서 매우 제한된 자원입니다. 6400바이트에서 2304바이트로의 감소는 동일한 블록 내에서 더 많은 스레드나 다른 연산이 공유 메모리를 활용할 수 있게 하여, 전체적인 GPU 점유율(Occupancy) 향상에 기여합니다.
2. 지연 시간(Latency) 감소: 공유 메모리 스테이징과 cp.async.wait_group을 통한 동기화 단계가 사라지면서, 데이터 로드 후 즉시 연산이 가능해졌습니다. 실제 벤치마크 결과 GB300에서 최대 4%, B200에서 최대 6%의 성능 향상을 보였습니다.
3. 일관성 유지: 데이터 타입과 연산 순서는 그대로 유지되므로, 출력값의 비트 단위 일치(bitwise equality)가 보장됩니다.

교훈

GPU 커널 최적화 시, 공유 메모리를 거치는 것이 항상 정답은 아닙니다. 데이터 재사용성이 낮거나 특정 아키텍처에서 레지스터 접근이 충분히 빠르다면, 공유 메모리 스테이징을 제거하는 것이 오히려 오버헤드를 줄이는 지름길이 될 수 있습니다.

참고 자료

• https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async

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