본문으로 건너뛰기

[flashinfer] FlashInfer의 CAKE GEMM 커널 최적화: 코드 중복 제거를 통한 성능 및 유지보수성 향상

PR 링크: flashinfer-ai/flashinfer#6938 상태: Merged | 변경: +2650 / -23297

들어가며

고성능 딥러닝 추론 라이브러리인 FlashInfer는 GPU의 잠재력을 최대한 활용하여 복잡한 연산을 효율적으로 수행합니다. 특히, dense-projection GEMM과 같은 핵심 연산은 모델의 성능에 지대한 영향을 미칩니다. 이러한 연산은 종종 CAKE와 같은 코드 생성(code generation) 프레임워크를 통해 최적화된 CUDA 커널로 구현됩니다. 하지만 코드 생성 방식은 때때로 많은 수의 유사하거나 중복된 파일을 만들어내어 코드 베이스의 복잡성을 증가시키고, 유지보수를 어렵게 하며, 심지어 JIT(Just-In-Time) 컴파일 시간에 영향을 줄 수 있습니다.

오늘 분석할 FlashInfer의 PR은 cake_dense_projection_gemm 패키지의 이러한 문제를 해결하기 위한 'deslop'(de-sloppification) 작업의 일환입니다. 이 PR은 87개의 프로그램과 104개의 빌드(SM100 47개, SM107 57개)를 포함하는 방대한 코드 베이스에서 중복된 호스트 바인딩(host binding)과 커널 소스를 통합하고, 사용되지 않는 코드를 제거하며, 테스트 템플릿 설정을 자동화하여 코드 품질과 유지보수성을 대폭 향상시키는 것을 목표로 합니다. 이 최적화는 런타임 성능 저하 없이 코드 베이스를 간결하게 만들었으며, 특히 JIT 컴파일 과정에서 이점을 제공합니다.

코드 변경사항 분석

이 PR은 총 4개의 커밋으로 구성되어 있으며, 각 커밋은 cake_dense_projection_gemm 패키지의 특정 영역에서 중복을 제거하고 코드를 정리하는 데 초점을 맞춥니다. 핵심 변경사항은 build.json 파일에서 소스 파일 경로와 컴파일 플래그가 어떻게 변경되는지를 통해 확인할 수 있습니다.

Commit 1: 81개의 중복 바인딩을 6개의 공유 바인딩으로 통합

문제: 기존에는 6가지 런처(launcher) 형태에 대해 87개의 호스트 바인딩 중 81개가 사실상 바이트 단위로 동일한 중복 코드였습니다. 커널 심볼, 호스트-shim 네임스페이스, TMA 인코더 이름, 동적 공유 메모리 리터럴 등 몇몇 값만 달랐습니다. 이는 코드 중복을 야기하고 유지보수를 어렵게 했습니다.

해결: 각 런처 형태별로 하나의 대표 바인딩만 남기고, 나머지 중복 바인딩은 삭제했습니다. 삭제된 바인딩의 build.json 엔트리는 이제 공유 바인딩을 두 번째 소스로 지정하고, 이전에 개별 바인딩에 있던 고유한 값들을 nvcc_flags의 컴파일 정의(-D 매크로)로 전달합니다. 공유 바인딩 파일 내부에서는 #ifndef 블록과 매크로를 사용하여 이러한 값들을 기본값으로 설정하거나 컴파일 정의를 통해 오버라이드할 수 있도록 했습니다.

Before (build.json 예시 - cake_dense_projection_gemm_087ba4f157914949c26a_sm_107a):

--- a/flashinfer/experimental/dense_projection_gemm/build.json
+++ b/flashinfer/experimental/dense_projection_gemm/build.json
@@ -105,16 +105,24 @@
       "gencode": "-gencode=arch=compute_107a,code=sm_107a",
       "min_cuda": "13.0",
       "name": "cake_dense_projection_gemm_087ba4f157914949c26a_sm_107a",
-      "nvcc_flags": [],
+      "nvcc_flags": [
+        "-DCAKE_PROGRAM_KERNEL=kernel_cake_dense_projection_gemm_087ba4f157914949c26a",
+        "-DCAKE_PROGRAM_SHIM=cake_host_shim_ef21bfa86e13f21d",
+        "-DCAKE_SMEM_TOTAL=230528",
+        "-DCAKE_ENCODE_A=EncodeTma_A_257c8b44c7da",
+        "-DCAKE_ENCODE_B=EncodeTma_B_0bd184b3b05e",
+        "-DCAKE_ENCODE_OUT32=EncodeTma_OUT32_dea947d1de88",
+        "-DCAKE_ENCODE_OUT16=EncodeTma_OUT16_5c144af57a50"
+      ],
       "nvcc_flags_if": [],
       "path": "aot/cake_dense_projection_gemm_087ba4f157914949c26a_sm_107a.so",
       "sha256": {
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_087ba4f157914949c26a_binding.cu": "151cd32a46c720f7e1de459749f8350f7d116b129d72d26da26cd0c7bddb9d22",
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_087ba4f157914949c26a_kernel.cu": "56f044ab992d763e6a9a08aa7385c1d92379dd69cf3304a2237f4b6fd1509eec"
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_0618ae0aec1af120afd0_binding.cu": "aec16c604de519ff37ba55a264ea96941ec8566baacce79ea85bd1ad2457d9a3",
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_087ba4f157914949c26a_kernel.cu": "f53a41537f6d48c923cb57bbca29b5ee0b08ac03d01281ed63ef1913c5976cf1"
       },
       "sources": [
         "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_087ba4f157914949c26a_kernel.cu",
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_087ba4f157914949c26a_binding.cu"
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_0618ae0aec1af120afd0_binding.cu"
       ]
     },

After (build.json 예시 - cake_dense_projection_gemm_087ba4f157914949c26a_sm_107a): 위 diff에서 볼 수 있듯이, sources 배열의 두 번째 항목이 cake_dense_projection_gemm_087ba4f157914949c26a_binding.cu에서 cake_dense_projection_gemm_0618ae0aec1af120afd0_binding.cu (공유 바인딩)로 변경되었습니다. 또한, nvcc_flags에 -DCAKE_PROGRAM_KERNEL, -DCAKE_PROGRAM_SHIM 등 해당 프로그램에 특화된 정의들이 추가되어, 공유 바인딩 파일이 이 정의들을 사용하여 올바른 코드를 생성하도록 합니다.

Commit 2: 10개의 클론 커널 소스를 7개의 바디로 통합

문제: 7개의 커널 소스 그룹(총 17개의 translation unit)이 서로 클론(clone) 관계였습니다. 이들은 공유 메모리 레이아웃 상수, 배리어 오프셋, 프롤로그/클린업 구문 등 몇몇 라인 범위에서만 차이가 있었습니다. 이 역시 코드 중복을 심화시키고 변경 사항 적용을 복잡하게 만들었습니다.

해결: 각 그룹에서 사전순으로 첫 번째 멤버의 소스를 '바디(body)'로 지정했습니다. 차이가 나는 라인 범위는 #if / #elif / #else 조건부 컴파일 지시문 뒤에 배치하고, 멤버의 컴파일 정의를 키로 사용했습니다. 커널 심볼은 매크로로 정의하고, #ifndef를 사용하여 기본값을 설정했습니다. build.json에서는 삭제된 10개의 소스 대신 공유 바디를 첫 번째 소스로 지정하고, 멤버별 정의를 nvcc_flags에 추가했습니다.

Before (build.json 예시 - cake_dense_projection_gemm_0c72e864a4afb5347da6_sm_100a):

--- a/flashinfer/experimental/dense_projection_gemm/build.json
+++ b/flashinfer/experimental/dense_projection_gemm/build.json
@@ -150,11 +166,11 @@
       "nvcc_flags_if": [],
       "path": "aot/cake_dense_projection_gemm_0c72e864a4afb5347da6_sm_100a.so",
       "sha256": {
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_tpl_nn_n128_m64_t_pro7__f32_kernel.cu": "679e529d8f4bd2d5d3ca9b1d1e8eb4bdb4aba728317eca48a2d58f9e4559b27a",
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_0c72e864a4afb5347da6_tpl_binding.cu": "76f55a39843108d2df7508bd852f4a130588577398c4a7115cf163a260d28af7"
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu": "551a16aed76be1869a1fc74d2350a51b5312993739b520d2735d6797110b7052",
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_0c72e864a4afb5347da6_tpl_binding.cu": "b984412fdddda0871f446f22e04184b266e8840c3bdb1cd23b5047e9008a34a4"
       },
       "sources": [
-        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_tpl_nn_n128_m64_t_pro7__f32_kernel.cu",
+        "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu",
         "csrc/cake_dense_projection_gemm/cake_dense_projection_gemm_0c72e864a4afb5347da6_tpl_binding.cu"
       ]
     },

After (build.json 예시 - cake_dense_projection_gemm_0c72e864a4afb5347da6_sm_100a): sources 배열의 첫 번째 항목이 cake_dense_projection_gemm_tpl_nn_n128_m64_t_pro7__f32_kernel.cu에서 cake_dense_projection_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu로 변경되었습니다. 이는 특정 커널의 pro7 버전이 hen 버전으로 통합되었음을 나타냅니다. nvcc_flags에 -DCAKE_DPG_KN_N256_TMA1_PRO7__F32_SELECT=1와 같은 정의가 추가되어, 공유 커널 바디가 조건부 컴파일을 통해 올바른 코드를 생성하도록 합니다.

Commit 3: 사용되지 않는 아키텍처별 코드 및 체크 함수 제거

문제: 특정 아키텍처(예: SM107)를 대상으로 하는 커널 소스에 SM100 전용 공유 메모리 주소 암(_CUDA_ARCH__ == 1000)이 포함되어 있거나, 모든 빌드에서 도달할 수 없는 pre-SM90 #error 가드, 그리고 어떤 런처도 호출하지 않는 check_dense_leading_dims 함수와 같은 도달 불가능한 코드가 존재했습니다. 이는 불필요한 코드 복잡성을 증가시키고 잠재적인 오류의 원인이 될 수 있습니다.

해결: 도달 불가능한 코드를 제거했습니다. 예를 들어, SM107 전용 커널에서 SM100 관련 #if 블록을 삭제하여 #else 바디가 무조건적인 텍스트가 되도록 했습니다. csrc/cake_ffi_checks.h에서 사용되지 않는 check_dense_leading_dims 함수도 제거했습니다. 이 변경은 build.json의 sha256 다이제스트 업데이트를 통해 반영됩니다.

코드 예시 (개념적 변화):

// Before (conceptual)
#if __CUDA_ARCH__ == 1000 // SM100 specific code
    // ... SM100 specific shared memory logic ...
#else // SM107 specific code
    // ... SM107 specific shared memory logic ...
#endif

// After (conceptual, for SM107-only kernel)
// ... SM107 specific shared memory logic ... // Now unconditional

Commit 4: 테스트의 exported-template 세트 레지스트리에서 파생

문제: tests/experimental/test_cake_dense_projection_gemm.py 파일에 EXPORTED_TEMPLATES라는 84개의 엔트리를 가진 리터럴(literal) 목록이 있었습니다. 이 목록은 레지스트리의 내용을 그대로 반영했지만, 수동으로 관리되어 레지스트리 변경 시 불일치가 발생할 위험이 있었습니다.

해결: 리터럴 목록을 제거하고, 레지스트리에서 dense_proj_gemm_* 템플릿을 동적으로 파생하도록 변경했습니다. 이는 레지스트리의 커널 테이블에서 모든 아키텍처에 걸쳐 템플릿을 수집하여 EXPORTED_TEMPLATES를 구성합니다. assert len(EXPORTED_TEMPLATES) == 84와 같은 어설션으로 예상되는 템플릿 수를 고정하여, 레지스트리 변경 시 테스트가 실패하도록 하여 불일치를 조기에 감지할 수 있게 했습니다. 이 변경은 테스트 코드에만 영향을 미치며, 런타임 동작에는 영향을 주지 않습니다.

왜 이게 좋은가

이 PR은 FlashInfer의 cake_dense_projection_gemm 패키지에 여러 가지 중요한 이점을 가져다줍니다.

  1. 코드 중복 제거 및 코드 베이스 간소화:

    • 81개의 중복 바인딩을 6개로, 10개의 클론 커널 소스를 7개의 바디로 통합함으로써 총 23,297줄의 코드를 제거하고 2,650줄을 추가하여 전체 코드 베이스를 대폭 줄였습니다. 이는 코드의 양을 줄여 가독성을 높이고, 개발자가 핵심 로직에 더 집중할 수 있게 합니다.

  2. 유지보수성 향상:

    • 중복된 코드가 사라지면서, 특정 로직을 변경해야 할 때 여러 파일을 수정할 필요 없이 공유 바인딩이나 커널 바디 하나만 수정하면 됩니다. 이는 버그 발생 가능성을 줄이고, 새로운 기능 추가나 최적화 작업을 훨씬 용이하게 만듭니다.
    • #if / #elif / #else와 컴파일 정의를 활용한 조건부 컴파일은 다양한 변형을 하나의 소스 파일에서 관리할 수 있게 하여 코드의 일관성을 유지합니다.

  3. JIT 컴파일 및 캐시 효율성 증대:

    • 가장 중요한 성능 이점 중 하나는 JIT 컴파일 시간 단축입니다. 기존에는 97개의 리포인팅된 바인딩 빌드와 14개의 리포인팅된 커널 빌드가 각각 고유한 소스 파일을 가지고 있었으므로, 소스 파일이 변경될 때마다 JIT 캐시가 해당 모듈들을 다시 빌드해야 했습니다. 이 PR 이후에는 더 적은 수의 고유한 소스 파일만 존재하므로, JIT 캐시가 재빌드해야 하는 모듈의 수가 줄어들어 컴파일 시간이 단축됩니다. PR 설명에 따르면, 기존 JIT 캐시는 이러한 모듈들을 한 번만 재빌드하면 됩니다.

  4. 런타임 성능 유지 및 검증:

    • 이 PR의 모든 변경사항은 행동 변화 없이 기계적으로 이루어졌습니다. PR 설명에 따르면, 모든 도달 가능한 입력에 대해 런처, 인자, 라우트, 익스포트, 모듈 이름 또는 테스트 결과가 변경되지 않았습니다. 특히, Device gate 검증 결과 104개의 모든 (프로그램, 아키텍처) 빌드에서 SASS 및 cubin 다이제스트가 동일하며, 리소스 사용량도 동일함을 확인했습니다. 이는 코드 최적화가 런타임 성능 저하 없이 이루어졌음을 명확히 보여줍니다.
    • CAKE 팀이 개발한 커널은 병합 당시 관련 SOTA(State-Of-The-Art) 기준보다 대부분의 성능 벤치마크 형태에서 더 나은 성능을 보였으며, 그렇지 않은 경우에도 동등한 성능을 유지했습니다. 이 PR은 이러한 성능 우위를 그대로 보존합니다.

  5. 테스트 견고성 강화:

    • 테스트 템플릿 목록을 레지스트리에서 동적으로 파생함으로써, 수동 업데이트의 필요성을 없애고 레지스트리와 테스트 간의 불일치 가능성을 제거했습니다. 이는 테스트의 신뢰성을 높이고, 개발자가 새로운 커널을 추가할 때 테스트 목록을 수동으로 업데이트하는 번거로움을 덜어줍니다.

일반적인 교훈: 이 PR은 코드 생성 시스템을 사용하는 프로젝트에서 흔히 발생하는 문제와 그 해결책을 잘 보여줍니다. 특히 다음과 같은 교훈을 얻을 수 있습니다.

  • 중복 제거의 중요성: 코드 중복은 유지보수 비용을 증가시키고 버그 발생 가능성을 높입니다. 코드 생성 시스템에서도 중복 제거는 필수적입니다.
  • 조건부 컴파일의 현명한 활용: #if, #ifndef와 같은 전처리기 지시문은 다양한 변형을 하나의 소스 파일에서 효율적으로 관리하는 강력한 도구입니다.
  • 철저한 검증: 코드 변경이 런타임 동작이나 성능에 영향을 미치지 않음을 보장하기 위한 엄격한 검증(SASS 비교, 테스트 결과 일치 등)은 필수적입니다.
  • JIT 컴파일러의 특성 이해: JIT 컴파일러의 동작 방식을 이해하고, 소스 파일 수나 변경 범위가 컴파일 시간에 미치는 영향을 고려하여 최적화를 수행해야 합니다.

마치며

FlashInfer의 cake_dense_projection_gemm 패키지에 적용된 이번 최적화는 단순히 코드 줄 수를 줄이는 것을 넘어, 프로젝트의 장기적인 건강성과 효율성을 크게 향상시키는 중요한 작업입니다. 코드 중복을 제거하고, 유지보수성을 높이며, JIT 컴파일 효율성을 개선하는 동시에 기존의 뛰어난 런타임 성능을 완벽하게 보존했다는 점에서 모범적인 기술 최적화 사례라고 할 수 있습니다. 이러한 노력은 FlashInfer가 고성능 딥러닝 추론 분야에서 선두를 유지하는 데 기여할 것입니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글