[flashinfer] FlashInfer cake_sampling 최적화: Speculative Sampling과 Whole-CTA Tail 도입으로 성능 향상

PR 링크: flashinfer-ai/flashinfer#5847 상태: Merged | 변경: +317916 / -81636

들어가며

대규모 언어 모델(LLM)의 추론 성능은 점점 더 중요해지고 있으며, 특히 샘플링 단계에서의 효율성은 전체 추론 속도에 큰 영향을 미칩니다. FlashInfer는 이러한 요구에 부응하기 위해 고성능 추론 엔진을 개발해왔습니다. 이번 PR(#5731)은 cake_sampling 모듈의 성능을 한 단계 더 끌어올리기 위한 일곱 번째 개선 사항을 담고 있습니다. 주요 내용은 Speculative Sampling (stage-1 twin)과 Whole-CTA Tail (for cluster-8 streams), 그리고 Capability-Keyed Host Policy의 도입입니다. 이 PR은 기존의 top_k_top_p_sampling_from_probs 함수의 정확성을 그대로 유지하면서, 특정 워크로드에서 상당한 성능 향상을 목표로 합니다.

이번 글에서는 이 PR에서 어떤 변경이 이루어졌고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 실제 코드 diff를 통해 자세히 알아보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 cake_sampling의 다양한 커널 변형을 도입하고, 호스트 단에서 최적의 커널을 선택하는 로직을 개선하는 데 있습니다. 변경 사항은 크게 다음과 같이 나눌 수 있습니다.

1. cake_sampling_binding.cuh: 새로운 커널 변형 및 선택 로직 추가

cake_sampling_binding.cuh 파일에서는 새로운 커널 변형을 위한 플래그와 Stage1Variant 구조체에 필드가 추가되었습니다. 또한, FindStage1 함수가 수정되어 새로운 변형을 고려하도록 변경되었습니다.

Before:

Stage1Variant 구조체는 fused_tail과 fused_block_tail, coarse_sample 필드를 가지고 있었습니다.

struct Stage1Variant {
  int32_t fused_tail;
  int32_t fused_block_tail;
  int32_t coarse_sample;
};

FindStage1 함수는 block_tail과 wide (coarse_sample을 의미) 인자를 받아 적합한 변형을 찾았습니다.

inline const Stage1Variant* FindStage1(int32_t cluster, int32_t ept, int32_t stream, 
                                       int32_t block_tail, int32_t wide) {
  static const Stage1Variant kTable[] = {CAKE_SAMPLING_STAGE1_TABLE(CAKE_SAMPLING_STAGE1_ENTRY)};
  for (const Stage1Variant& v : kTable) {
    if (v.cluster == cluster && v.ept == ept && v.stream == stream &&
        v.fused_block_tail == block_tail && v.coarse_sample == wide)
      return &v;
  }
  return nullptr;
}

After:

Stage1Variant 구조체에 spec_sample 필드가 추가되었습니다. 이는 Speculative Sampling을 위한 변형을 나타냅니다.

struct Stage1Variant {
  int32_t fused_tail;
  int32_t fused_block_tail;
  int32_t coarse_sample;
  int32_t spec_sample; // 1: the speculative-sample build (first chunk is the sample; launch_flags bit 6)
};

CAKE_SAMPLING_STAGE1_ENTRY 매크로에도 spec 인자가 추가되어 새로운 필드를 채울 수 있게 되었습니다.

#define CAKE_SAMPLING_STAGE1_ENTRY(symbol, cluster, ept, stream, threads, smem, fused, \
                                   fused_block, wide, spec) \
  {reinterpret_cast<const void*>(&symbol), \
   cluster, \
   ept, \
   stream, \
   threads, \
   smem, \
   fused, \
   fused_block, \
   wide, \
   spec},

FindStage1 함수는 spec 인자를 추가로 받도록 변경되었습니다. 또한, launch_flags에서 kFlagSpecSample (bit 6)이 정의되었고, 이 플래그가 spec_sample 필드와 연동됩니다. 새로운 TVM_FFI_ICHECK 문들은 spec_sample이 다른 플래그(bit 3, 4)와 배타적임을 명시합니다.

constexpr int64_t kFlagSpecSample =
    64;  // host-side build selection (speculative sample); never forwarded to the kernel

inline const Stage1Variant* FindStage1(int32_t cluster, int32_t ept, int32_t stream, 
                                       int32_t block_tail, int32_t wide, int32_t spec) {
  static const Stage1Variant kTable[] = {CAKE_SAMPLING_STAGE1_TABLE(CAKE_SAMPLING_STAGE1_ENTRY)};
  for (const Stage1Variant& v : kTable) {
    if (v.cluster == cluster && v.ept == ept && v.stream == stream &&
        v.fused_block_tail == block_tail && v.coarse_sample == wide && v.spec_sample == spec)
      return &v;
  }
  return nullptr;
}

// ... RadixTopK 함수 내에서
  const int32_t want_spec = (launch_flags & kFlagSpecSample) != 0 ? 1 : 0;
  TVM_FFI_ICHECK(!(want_spec && stream_variant == 0))
      << 

## 참고 자료
- https://pytorch.org/docs/stable/generated/torch.compile.html

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