본문으로 건너뛰기

[flashinfer] FlashInfer: FP8 64-head 디코딩 성능 최적화 및 Programmatic Dependent Launch 도입

PR 링크: flashinfer-ai/flashinfer#8543 상태: Merged | 변경: +189 / -97

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기뿐만 아니라 추론 속도와 효율성 또한 중요하게 만들고 있습니다. 특히, GPU 하드웨어의 발전과 함께 FP8과 같은 저정밀도 데이터 타입을 활용하여 메모리 대역폭과 연산 성능을 극대화하려는 노력이 활발히 이루어지고 있습니다.

FlashInfer는 이러한 추세에 발맞춰 LLM 추론을 위한 고성능 커널 라이브러리를 제공합니다. 이번 PR은 FlashInfer의 cake_dsv4 백엔드에서 특정 FP8 64-head 디코딩 시나리오의 성능을 개선하고, Programmatic Dependent Launch (PDL) 기능을 도입하여 GPU 활용도를 더욱 높이는 것을 목표로 합니다.

이 글에서는 해당 PR의 코드 변경 사항을 분석하고, 어떤 부분이 어떻게 성능 향상에 기여하는지, 그리고 이러한 최적화가 가지는 일반적인 교훈은 무엇인지 살펴보겠습니다.

코드 변경 분석

이번 PR의 핵심 변경 사항은 cake_dsv4 백엔드에서 특정 조건의 FP8 64-head 디코딩 연산을 더 효율적인 커널로 라우팅하고, Programmatic Dependent Launch (PDL) 기능을 도입하는 것입니다.

1. 커널 라우팅 변경 (cake_dsv4_cbf3168c696847fd3f69_binding.cu)

기존에는 특정 FP8 64-head 디코딩 연산 (쿼리 토큰 수가 12개 이하이고 sparse width가 260~380 사이인 경우)이 fp8_h64_prefill_source_persistent_m64_multi_tile 커널에서 실행되었습니다. 이 PR에서는 해당 연산을 fp8_lowhead_h64 커널로 라우팅하도록 변경했습니다. 이 새로운 커널은 H64 클러스터 디코딩에 최적화되어 있으며, PDL 기능을 지원합니다.

PR 설명에 따르면, 이 변경으로 인해 NVIDIA B200 (SM100)에서는 3개의 shape에서 1.047x ~ 1.065x의 속도 향상이, GB300 (SM103)에서는 3개의 shape에서 1.027x ~ 1.034x의 속도 향상이 관찰되었습니다. 대부분의 다른 shape에서는 코드 변경이 없어 성능 변화가 없었습니다.

주요 변경 코드는 다음과 같습니다:

Before:

void Run(TensorView a_tmap_q, TensorView a_tmap_swa_kv, TensorView a_tmap_compressed_kv, TensorView a_O, TensorView a_partial_lse, TensorView a_swa_indices, TensorView a_compressed_indices, TensorView a_sparse_topk_lens, TensorView a_sinks, TensorView a_bmm1_scale, TensorView a_bmm2_scale, int64_t a_num_heads, int64_t a_swa_index_stride, int64_t a_compressed_index_stride, int64_t a_sparse_topk_lens_offset, int64_t a_num_query_tokens, int64_t a_sparse_topk, int64_t a_has_sinks, int64_t a_total_work_items, TensorView a_seq_lens, TensorView a_cum_seq_lens_q, int64_t a_ragged_query, int64_t a_max_q_len, int64_t a_batch_size, int64_t grid_x, int64_t grid_y, int64_t grid_z) {

After:

void Run(TensorView a_tmap_q, TensorView a_tmap_swa_kv, TensorView a_tmap_compressed_kv, TensorView a_O, TensorView a_partial_lse, TensorView a_swa_indices, TensorView a_compressed_indices, TensorView a_sparse_topk_lens, TensorView a_sinks, TensorView a_bmm1_scale, TensorView a_bmm2_scale, int64_t a_num_heads, int64_t a_swa_index_stride, int64_t a_compressed_index_stride, int64_t a_sparse_topk_lens_offset, int64_t a_num_query_tokens, int64_t a_sparse_topk, int64_t a_has_sinks, int64_t a_total_work_items, TensorView a_seq_lens, TensorView a_cum_seq_lens_q, int64_t a_ragged_query, int64_t a_max_q_len, int64_t a_batch_size, int64_t a_enable_pdl, int64_t grid_x, int64_t grid_y, int64_t grid_z) {
  TVM_FFI_CHECK(a_enable_pdl == 0LL || a_enable_pdl == 1LL, ValueError)
      << KERNEL_NAME << ": 'enable_pdl' = " << a_enable_pdl << " is not 0 or 1";

a_enable_pdl 인자가 추가된 것을 볼 수 있습니다. 이 인자는 PDL 기능의 활성화 여부를 제어합니다.

2. Programmatic Dependent Launch (PDL) 도입 (cake_dsv4_cbf3168c696847fd3f69_binding.cu 및 cake_dsv4_cbf3168c696847fd3f69_kernel.cu)

이 PR은 cudaLaunchKernelExC API를 사용하여 PDL 기능을 도입했습니다. PDL은 CUDA 런타임이 커널 실행 간의 종속성을 관리하도록 하여, GPU가 더 효율적으로 작업을 스케줄링하고 유휴 시간을 줄일 수 있게 합니다.

cake_dsv4_cbf3168c696847fd3f69_binding.cu 파일에서 cudaLaunchKernelExC 호출 시 cudaLaunchAttributeProgrammaticStreamSerialization 속성을 설정하는 부분이 추가되었습니다.

Before (부분):

  cudaLaunchAttribute attrs[2] = {};
  attrs[0].id = cudaLaunchAttributeClusterDimension;
  attrs[0].val.clusterDim.x = 2;
  attrs[0].val.clusterDim.y = 1;
  attrs[0].val.clusterDim.z = 1;
  attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference;
  attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread;
  config.attrs = attrs;
  config.numAttrs = 2;

After (부분):

  cudaLaunchAttribute attrs[3] = {};
  attrs[0].id = cudaLaunchAttributeClusterDimension;
  attrs[0].val.clusterDim.x = 2;
  attrs[0].val.clusterDim.y = 1;
  attrs[0].val.clusterDim.z = 1;
  attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference;
  attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread;
  int n = 2;
  if (a_enable_pdl) {
    attrs[n].id = cudaLaunchAttributeProgrammaticStreamSerialization;
    attrs[n].val.programmaticStreamSerializationAllowed = 1;
    ++n;
  }
  config.attrs = attrs;
  config.numAttrs = n;

a_enable_pdl 플래그가 활성화되면 cudaLaunchAttributeProgrammaticStreamSerialization 속성이 설정되어 PDL 기능이 활성화됩니다. 이는 GPU가 커널 간의 의존성을 더 지능적으로 관리하여 잠재적인 병렬성을 극대화하는 데 도움을 줍니다.

또한, 커널 내부에서도 griddepcontrol.wait 및 griddepcontrol.launch_dependents와 같은 PDL 관련 명령어가 추가되어, 커널 실행 흐름이 PDL과 동기화되도록 합니다.

Before (커널 내부):

    // Kernel setup ops

After (커널 내부):

    // Prologue
    // ... (기존 코드)

    // Mbarrier init (45 barriers at smem_raw[0..360))

    if (warp == 12) {
        uint32_t leader = elect_sync();
    }

    // ... (기존 코드)

    const unsigned int taddr = tmem_addr_storage[0];
    const unsigned int tmem_tmem_scratch = taddr;

    // Kernel post-init ops
    asm volatile("griddepcontrol.wait;" ::: "memory");
    asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
    if (warp >= 8 && warp <= 11) {
        asm volatile("setmaxnreg.inc.sync.aligned.u32 136;");
    }

3. 기타 변경 사항

  • shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff; -> shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xFFFFFFFF;: shfl 명령어의 마스크 비트가 0xffffffff에서 0xFFFFFFFF로 변경되었습니다. 이는 큰 차이는 없지만, 코드의 일관성을 높입니다.
  • float2 _reg_reduce_max2_7 = {-CAKE_INF, -CAKE_INF}; -> float2 _reg_reduce_max2_7 = {-CUDART_INF_F, -CUDART_INF_F};: -CAKE_INF가 CUDA의 표준 매크로인 -CUDART_INF_F로 변경되었습니다. 이는 코드의 가독성과 표준 준수를 향상시킵니다.
  • float delta = __fmaf_rn(row_max_val, softmax_scale_log2, (-new_max) * softmax_scale_log2); -> float delta = __fmaf_rn(row_max_val, softmax_scale_log2, -new_max * softmax_scale_log2);: -new_max 앞에 괄호가 제거되었습니다. 연산 결과에는 영향을 미치지 않지만, 코드 스타일을 다듬었습니다.

왜 이게 좋은가?

이 PR은 다음과 같은 이유로 좋은 최적화 및 개선이라고 할 수 있습니다.

  1. 성능 향상: 특정 FP8 64-head 디코딩 시나리오에서 1.027x ~ 1.151x의 유의미한 속도 향상을 달성했습니다. 이는 LLM 추론 시 지연 시간을 줄이고 처리량을 높이는 데 직접적으로 기여합니다.
  2. GPU 활용도 증대: Programmatic Dependent Launch (PDL) 도입은 CUDA 런타임이 커널 실행 간의 종속성을 더 효율적으로 관리하도록 하여 GPU의 잠재적인 병렬성을 최대한 활용합니다. 이는 GPU의 유휴 시간을 줄이고 전반적인 컴퓨팅 효율성을 높입니다.
  3. FP8 및 저정밀도 연산 최적화: LLM에서 FP8과 같은 저정밀도 데이터 타입의 사용이 증가함에 따라, 이러한 데이터 타입을 위한 특화된 커널 최적화는 매우 중요합니다. 이 PR은 이러한 추세를 반영하여 FP8 연산의 성능을 개선했습니다.
  4. 코드 가독성 및 유지보수성 향상: 표준 매크로 사용, 불필요한 괄호 제거 등 사소하지만 코드의 가독성과 유지보수성을 향상시키는 변경 사항들이 포함되었습니다.

성능 수치: PR 설명에 따르면, PDL 활성화 시 다음과 같은 성능 향상이 관찰되었습니다.

GPU points FlashInfer main (us) this PR (us) speedup
NVIDIA B200 (SM100, sm_100a) 64 8.83 to 9.28 7.75 to 8.33 1.098x to 1.151x
NVIDIA GB300 (SM103, sm_103a) 64 8.42 to 8.64 7.53 to 7.86 1.090x to 1.137x

이는 PDL 기능이 실제 워크로드에서 상당한 성능 향상을 가져올 수 있음을 보여줍니다.

일반적인 교훈:

  • 특정 워크로드에 대한 최적화: 모든 경우에 적용되는 범용 최적화보다는, 특정 조건(예: FP8, 64-head, 낮은 쿼리 토큰 수)에 맞는 커널로 라우팅하는 것이 더 큰 성능 향상을 가져올 수 있습니다.
  • 최신 하드웨어 기능 활용: PDL과 같은 최신 CUDA 기능을 적극적으로 도입하여 GPU의 잠재 성능을 최대한 끌어내는 것이 중요합니다.
  • 정밀한 프로파일링 및 벤치마킹: 성능 개선 효과를 정확히 측정하고 검증하기 위해 상세한 벤치마킹 데이터와 다양한 시나리오 분석이 필수적입니다.

리뷰 댓글 분석

리뷰어 yyihuang은 이 PR이 cake_dsv4의 호스트 파일과 테스트 모듈을 변경하는 후속 PR과 연관되어 있음을 인지하고, 해당 PR의 측정 결과가 변경되지 않도록 테스트를 요청했습니다. 이는 코드 변경이 특정 시나리오에 국한됨을 확인하고, 다른 부분에 미치는 영향을 최소화하려는 좋은 관행입니다. 또한, "Agreed, thanks."라는 긍정적인 피드백은 변경 사항이 타당함을 시사합니다.

결론

이번 PR은 FlashInfer의 cake_dsv4 백엔드에서 FP8 64-head 디코딩 성능을 최적화하고, Programmatic Dependent Launch (PDL) 기능을 도입하여 GPU 활용도를 높이는 중요한 개선을 이루었습니다. 이러한 최적화는 LLM 추론의 속도와 효율성을 향상시키는 데 크게 기여할 것입니다. 특히, 특정 하드웨어 아키텍처와 데이터 타입에 맞는 커널 라우팅 및 최신 CUDA 기능을 활용하는 전략은 고성능 컴퓨팅 분야에서 지속적으로 중요한 접근 방식이 될 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글