[flashinfer] FlashInfer의 CAKE 커널 최적화: 아키텍처 및 형상별 프로그램 선택 전략

PR 링크: flashinfer-ai/flashinfer#6008 상태: Merged | 변경: +20120 / -4198

들어가며

LLM 추론 성능을 극대화하기 위해서는 연산 커널의 세밀한 튜닝이 필수적입니다. 특히 fused_qk_rmsnorm_rope_append_paged_kv_cache와 같은 복합 연산은 다양한 GPU 아키텍처(SM90, SM100, SM103)와 입력 형상(shape)에 따라 최적의 실행 경로가 달라집니다. 본 PR은 CAKE 팀이 개발한 커널을 대상으로, 고정된 커널 대신 입력 조건에 따라 가장 효율적인 프로그램을 동적으로 선택하는 'Program Selection' 전략을 도입하여 성능을 크게 향상시켰습니다.

코드 분석

1. 매니페스트 기반의 정책 설정 (csrc/cake_fused_qk_rope_append/cake_fused_qk_rope_append_manifest.json)

핵심 변경 사항은 stage_policy의 도입입니다. 이전에는 아키텍처당 2개의 프로그램만 존재했으나, 이제는 20개의 프로그램 스테이지를 정의하고, 런타임에 입력 형상에 따라 적절한 커널을 선택합니다.

"stage_policy": {
  "rule": "decode = (num_rows == num_requests); hopper = arch sm_90*; ... -> WARPS_PER_ROW=w20, V_SUB_BASE=w20_v_sub_base; ...",
  "stages_by_arch": {
    "sm_100a": ["hq8_hkv1", "hq8_hkv1_vsub3_d", ...],
    "sm_90a": ["hq8_hkv1_vsub3_d", ...]
  }
}

이 규칙은 num_rows, num_requests, arch 등을 조합하여 최적의 WARPS_PER_ROW나 V_SUB_BASE 설정을 결정합니다. 이를 통해 하드웨어 리소스를 낭비 없이 활용할 수 있게 되었습니다.

2. JIT 컴파일 및 선택 로직 (flashinfer/jit/cake_fused_qk_rope_append.py)

Python 단에서는 stage_for 함수를 통해 매니페스트의 정책을 반영합니다. 각 프로그램은 JIT 컴파일 시 고유한 모듈 이름을 가지며, 한 번 컴파일된 후에는 캐싱되어 재사용됩니다.

왜 이게 좋은가

이번 최적화의 핵심은 '범용 커널'에서 '맞춤형 커널'로의 전환입니다.

• 성능 향상: H100 및 H200 GPU에서 기하 평균 기준 약 6~7%의 성능 향상을 보였으며, 특정 대규모 decode 배치(B200 B512)에서는 최대 1.30배의 속도 향상을 달성했습니다.
• 정밀도 유지: FP32 RMSNorm/RoPE 연산을 유지하면서도 BF16 연산의 효율을 극대화하여, 기존 커널과 비트 단위로 동일한 결과를 보장합니다.
• 교훈: GPU 커널 최적화 시, 단일 커널로 모든 형상을 커버하려 하기보다, 워크로드의 특성(batch size, head count, architecture)을 반영한 다중 스케줄을 준비하고 이를 런타임에 선택하는 방식이 SOTA 성능을 달성하는 데 필수적임을 보여줍니다.

참고 자료

• FlashInfer Documentation
• NVIDIA CUDA C++ Programming Guide

참고 자료

• https://flashinfer.ai/docs/index.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html

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