[flashinfer] FlashInfer Cake Sampling 최적화: Leader-Push Exchange와 Graph-Replay 개선

PR 링크: flashinfer-ai/flashinfer#6065 상태: Merged | 변경: +185059 / -14647

들어가며

대규모 언어 모델(LLM)의 추론 성능은 모델의 크기와 함께 계속해서 중요해지고 있습니다. 특히 샘플링 과정은 LLM 추론의 병목 현상 중 하나로, 더 빠르고 효율적인 샘플링 기법은 전체 추론 속도를 크게 향상시킬 수 있습니다. 최근 flashinfer-ai/flashinfer 레포지토리의 PR #5985는 이러한 샘플링 성능을 극적으로 개선하는 중요한 업데이트를 포함하고 있습니다. 이 PR은 "Cake Sampling"의 새로운 최적화 기법인 "Leader-Push Exchange"와 CUDA 그래프 캡처 시 발생하던 성능 저하 문제를 해결하는 "Graph-Replay Host Fix"를 도입했습니다.

본 글에서는 이 PR에서 이루어진 코드 변경사항을 상세히 분석하고, 각 최적화가 왜 성능 향상으로 이어지는지, 그리고 어떤 기술적 원리가 적용되었는지 살펴보겠습니다. 또한, 실제 코드 diff를 통해 변경 전후를 비교하며 최적화의 효과를 명확히 이해하는 것을 목표로 합니다.

코드 변경사항 분석

이번 PR의 핵심 변경사항은 크게 두 가지로 나눌 수 있습니다: Leader-Push Exchange 커널 도입과 Graph-Replay 시의 Host Rule 수정입니다.

1. Leader-Push Exchange (_lp, _cs_lp, _sp_lp) 도입

기존의 샘플링 방식은 여러 스레드 블록(CTA)이 데이터를 가져오는 "pull" 방식에 의존했습니다. 이 방식은 각 CTA가 필요한 데이터를 직접 가져와야 하므로 통신 오버헤드가 발생할 수 있습니다. 이번 PR에서는 이를 개선하기 위해 "Leader-Push Exchange"라는 새로운 방식을 도입했습니다. 이 방식은 특정 CTA(Leader)가 다른 CTA들의 데이터를 모아서 처리하는 "push" 방식을 활용합니다.

핵심 아이디어:

멀티-CTA 스트리밍(multi-CTA streaming) 환경에서, 각 CTA는 자신의 압축된 후보 리스트를 공유 메모리(shared memory)의 특정 위치(rank 0)에 직접 저장합니다. 또한, 각 CTA는 자신의 리스트 길이를 모든 CTA에 브로드캐스트합니다. 이후 단 한 번의 동기화(exchange barrier)를 통해 rank 0 CTA는 로컬에서 데이터를 취합하고, 임계값을 결정하며, 최종 샘플링을 수행합니다. 나머지 CTA들은 이 과정을 건너뛰고 종료됩니다.

이 방식은 기존의 pull 방식에서 발생하던 데이터 읽기(DSM read round) 및 동기화(fused_sync) 단계를 제거함으로써, 두 번의 클러스터 라운드를 절약합니다. 또한, 모든 CTA에게 데이터를 푸시하는 방식(push-to-all)에서 발생할 수 있는 N-copy 바이트 증폭 문제도 피할 수 있습니다.

코드 변경:

csrc/cake_sampling/cake_sampling_binding.cuh 파일에서 Stage1Variant 구조체에 leader_push 필드가 추가되었습니다. 이 필드는 새로운 Leader-Push Exchange 빌드를 나타냅니다.

--- a/csrc/cake_sampling/cake_sampling_binding.cuh
+++ b/csrc/cake_sampling/cake_sampling_binding.cuh
@@ -138,6 +138,10 @@ struct Stage1Variant {
   int32_t coarse_push;  // 1: each CTA stores its coarse histogram sums into every CTA's shared
                         // memory for the two-level select (the default build's twin taken by
                         // launch_flags bit 8 on a two-launch chain)
+  int32_t leader_push;  // 1: every CTA pushes its compacted candidate list into rank 0's receive
+                        // buffer and its length into every CTA before the single exchange barrier
+                        // (the default / sample builds' twin of a multi-CTA stream taken by
+                        // launch_flags bit 9)
 };
 
 struct Stage23Variant {
@@ -149,7 +153,7 @@ struct Stage23Variant {
 };
 
 #define CAKE_SAMPLING_STAGE1_ENTRY(symbol, cluster, ept, stream, threads, smem, fused, \
-                                   fused_block, wide, spec, slab, push)                \
+                                   fused_block, wide, spec, slab, push, leader)        \
   {reinterpret_cast<const void*>(&symbol),                                             \
    cluster,                                                                            \
    ept,                                                                                \
@@ -161,7 +165,8 @@ struct Stage23Variant {
    wide,                                                                               \
    spec,                                                                               \
    slab,                                                                               \
-   push},                                                                               \
+   push,                                                                               \
+   leader},                                                                             \
 #define CAKE_SAMPLING_STAGE23_ENTRY(symbol, threads, items, variant_flags, smem) \
   {reinterpret_cast<const void*>(&symbol), threads, items, variant_flags, smem}, 
 
@@ -180,7 +185,12 @@ struct Stage23Variant {
 // (coarse_push = 1: each CTA stores its coarse histogram sums into every CTA's shared memory so
 // the two-level select reads them locally) are a second build of the default build of the
 // streaming variants with the two-level select, taken only by launch_flags bit 8 on a two-launch
-// chain -- it wins those chains on GB300 / H100 and loses them on B200.
+// chain -- it wins those chains on GB300 / H100 and loses them on B200.  The leader-push exchange
+// (leader_push = 1: the CTAs store their candidate lists straight into rank 0's receive buffer
+// and their lengths into every CTA before one exchange barrier; rank 0 gathers locally) is a
+// third build of the default and sample builds of the multi-CTA streaming variants, taken only
+// by launch_flags bit 9 -- it is exclusive with the whole-CTA tail, the slab tail and the pushed
+// coarse sums, and the host selects it per capability, ept and row length.
 inline const Stage1Variant kStage1Table[] = {
     CAKE_SAMPLING_STAGE1_TABLE(CAKE_SAMPLING_STAGE1_ENTRY)};
 inline const Stage23Variant kStage23Table[] = {
@@ -191,11 +201,11 @@ inline std::atomic<uint64_t> kStage23Prepared[sizeof(kStage23Table) / sizeof(Sta
 
 inline const Stage1Variant* FindStage1(int32_t cluster, int32_t ept, int32_t stream,
                                        int32_t block_tail, int32_t wide, int32_t spec, int32_t slab,
-                                       int32_t push) {
+                                       int32_t push, int32_t leader) {
   for (const Stage1Variant& v : kStage1Table) {
     if (v.cluster == cluster && v.ept == ept && v.stream == stream &&
         v.fused_block_tail == block_tail && v.coarse_sample == wide && v.spec_sample == spec &&
-        v.slab_tail == slab && v.coarse_push == push)
+        v.slab_tail == slab && v.coarse_push == push && v.leader_push == leader)
       return &v;
   }
   return nullptr;
@@ -382,9 +392,15 @@ void RadixTopK(TensorView probs, TensorView topk_arr, int64_t topk_scalar, int64
          

## 참고 자료
- https://github.com/flashinfer-ai/flashinfer/blob/main/csrc/cake_sampling/cake_sampling_binding.cuh#L138
- https://github.com/flashinfer-ai/flashinfer/blob/main/csrc/cake_sampling/cake_sampling_binding.cuh#L149
- https://github.com/flashinfer-ai/flashinfer/blob/main/csrc/cake_sampling/cake_sampling_binding.cuh#L191

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