[flashinfer] FlashInfer CAKE FMHA 백엔드 최적화: 코드 중복 제거 및 성능 향상

PR 링크: flashinfer-ai/flashinfer#6704 상태: Merged | 변경: +5179 / -202273

들어가며

최근 LLM(거대 언어 모델)의 발전 속도는 눈부십니다. 특히 Transformer 아키텍처의 핵심인 어텐션 메커니즘은 모델의 성능을 좌우하는 중요한 요소입니다. 하지만 어텐션 연산은 계산 비용이 매우 높아 효율적인 구현이 필수적입니다. FlashInfer는 이러한 요구에 부응하여 고성능 어텐션 구현을 제공하는 라이브러리입니다. 이번 기술 블로그 글에서는 FlashInfer의 CAKE FMHA(Fused Multi-Head Attention) 백엔드에서 이루어진 대규모 코드 최적화 PR을 심층적으로 분석합니다. 이 PR은 코드 중복을 제거하고, 커널 로딩 방식을 개선하며, 새로운 테스트 케이스를 추가하여 전반적인 코드 품질과 성능을 향상시키는 것을 목표로 합니다.

이 PR은 특히 csrc/cake_fmha 디렉토리 내의 CUDA 커널 코드들을 대폭 정리하는 데 중점을 둡니다. 여러 커밋을 통해 유사한 기능을 수행하는 커널들을 하나의 소스로 통합하고, 아키텍처별 복제본을 줄이며, 호스트 코드의 중복을 제거합니다. 이러한 변경은 코드베이스의 유지보수성을 높이고, 잠재적인 버그를 줄이며, 컴파일 시간을 단축하는 효과를 가져옵니다.

코드 분석: 파일별 변경 사항

이 PR은 총 14개의 커밋으로 구성되어 있으며, 각 커밋은 CAKE FMHA 백엔드의 특정 부분을 개선합니다. 주요 변경 사항을 파일별로 나누어 살펴보겠습니다.

1. LSE(Log-Sum-Exp) 버퍼 문제 해결 및 제네릭 멤버 활용 (flashinfer/jit/cake_fmha.py, flashinfer/cake_fmha.py, tests/attention/test_cake_fmha.py)

첫 번째 커밋은 BF16 페이징 디코드 백엔드에서 발생하던 LSE 버퍼 관련 버그를 수정합니다. 이전에는 특정 exact shapes (b4, b128)에서 LSE 버퍼가 제대로 쓰이지 않았습니다. 이 문제는 return_lse 옵션이 추가되면서 해결되었습니다. 이제 return_lse=True일 경우, 해당 exact shape 멤버 대신 동일한 플래그를 가진 제네릭 멤버를 사용하도록 JIT 셀렉터가 변경되었습니다. 만약 제네릭 멤버가 없다면, compat_v1 모듈로 폴백합니다.

Before:

# 기존 로직에서는 exact shape 멤버를 직접 사용
# LSE 버퍼가 제대로 쓰이지 않는 문제가 있었음

After:

# return_lse=True 시, exact shape 멤버 대신 제네릭 멤버를 사용하도록 변경
# JIT 셀렉터 로직 수정
if return_lse:
    # ...
    return _resolve_decode_native_bf16_selector(..., return_lse=True)

테스트 케이스도 추가되어 return_lse=True 시의 동작을 검증합니다. 이 변경은 LSE 계산이 필요한 경우에도 최적화된 경로를 사용할 수 있도록 보장합니다.

2. decode_native_bf16 제네릭 커널 클론 통합 (csrc/cake_fmha/cuda/decode_native_bf16/default.cu, csrc/cake_fmha/registry.json)

이 커밋은 decode_native_bf16의 11개 제네릭 멤버를 단일 CUDA 소스 파일(default.cu)로 통합합니다. 기존에는 HAS_SINK, HAS_WINDOW, RETAIN_KV_L2, USE_SCALE_PTR 네 가지 플래그 조합에 따라 11개의 거의 동일한 CUDA 소스 파일이 존재했습니다. 이제 이 플래그들은 CAKE_FMHA_* 매크로를 통해 동적으로 제어되며, 단일 default.cu 파일 내의 #if 블록에서 처리됩니다.

Before:

// csrc/cake_fmha/cuda/decode_native_bf16/has_sink0_has_window0_retain_kv_l20_use_scale_ptr0.cu
#define HAS_SINK 0
#define HAS_WINDOW 0
#define RETAIN_KV_L2 0
#define USE_SCALE_PTR 0
// ... 커널 코드 ...

// csrc/cake_fmha/cuda/decode_native_bf16/has_sink1_has_window0_retain_kv_l20_use_scale_ptr0.cu
#define HAS_SINK 1
#define HAS_WINDOW 0
#define RETAIN_KV_L2 0
#define USE_SCALE_PTR 0
// ... 커널 코드 ...

After:

// csrc/cake_fmha/cuda/decode_native_bf16/default.cu
// ... 헤더 ...
#define HAS_SINK CAKE_FMHA_HAS_SINK
#define HAS_WINDOW CAKE_FMHA_HAS_WINDOW
#define RETAIN_KV_L2 CAKE_FMHA_RETAIN_KV_L2
#define USE_SCALE_PTR CAKE_FMHA_USE_SCALE_PTR

// ... 플래그에 따른 #if 블록 ...
#if HAS_SINK
// ... HAS_SINK=1 일 때의 코드 ...
#else
// ... HAS_SINK=0 일 때의 코드 ...
#endif
// ... 다른 플래그들에 대한 #if 블록 ...

이 변경으로 인해 decode_native_bf16 관련 소스 파일 수가 11개에서 1개로 줄어들었습니다. 이는 코드 중복을 크게 줄이고, 컴파일 시간을 단축하며, 버그 수정 및 유지보수를 용이하게 합니다.

3. FP16 디코드 커널 클론 통합 (csrc/cake_fmha/cuda/decode_native_fp16_nhd/default.cu, csrc/cake_fmha/cuda/decode_native_fp16_hd512/default.cu)

커밋 3은 커밋 2와 유사한 방식으로 decode_native_fp16_nhd (16개 멤버)와 decode_native_fp16_hd512 (8개 멤버)의 제네릭 커널 클론들을 각각 단일 default.cu 파일로 통합합니다. 사용되는 플래그는 HAS_SINK, HAS_WINDOW, RETAIN_KV_L2, USE_SCALE_PTR 등입니다. 이로써 각 패밀리별 소스 파일 수가 크게 감소했습니다.

4. 컨텍스트(Prefill) 커널 클론 통합 (csrc/cake_fmha/cuda/context_fp8/default.cu, csrc/cake_fmha/cuda/context_bf16/default.cu)

커밋 4는 컨텍스트(prefill) 단계에서 사용되는 context_fp8와 context_bf16 커널들의 제네릭 멤버들을 통합합니다. context_fp8는 양 아키텍처(sm_100a, sm_103a)에서 단일 default.cu를 공유합니다. context_bf16은 아키텍처별로 별도의 default.cu를 가지지만, 이 역시 커밋 7에서 하나의 파일로 통합됩니다. 통합에 사용되는 플래그는 ENABLE_SINK, IS_CAUSAL, RETURN_LSE입니다.

Before (예시: context_bf16 sm_100a):

// csrc/cake_fmha/cuda/context_bf16/sm_100a/enable_sink0_is_causal1_return_lse1.cu
// ...
#define IS_CAUSAL 1
#define ENABLE_SINK 0
#define RETURN_LSE 1
// ... 커널 코드 ...

After (예시: context_bf16 통합):

// csrc/cake_fmha/cuda/context_bf16/default.cu
// ...
#define IS_CAUSAL CAKE_FMHA_CONTEXT_IS_CAUSAL
#define ENABLE_SINK CAKE_FMHA_CONTEXT_ENABLE_SINK
#define RETURN_LSE CAKE_FMHA_CONTEXT_RETURN_LSE

#if __CUDA_ARCH__ == 1030
// ... sm_103a 전용 코드 ...
#else
// ... sm_100a 전용 코드 ...
#endif
// ...

이 커밋은 context_fp8와 context_bf16의 제네릭 멤버들을 통합하여 코드 중복을 줄이고 유지보수성을 향상시킵니다.

5. 양자화 디코드 커널 클론 통합 (csrc/cake_fmha/cuda/decode_quant_bf16q/default.cu, csrc/cake_fmha/cuda/decode_quant_nvfp4/default.cu, csrc/cake_fmha/cuda/decode_quant_fp8/default.cu)

커밋 5는 양자화된 디코드 컴포넌트(decode_quant_bf16q, decode_quant_nvfp4, decode_quant_fp8)의 페이지 크기별 클론들을 통합합니다. 각 컴포넌트는 페이지 크기(16 또는 32)에 따라 두 개의 멤버를 가졌으며, 이들은 페이지 테이블 로드 부분만 달랐습니다. 이제 이들은 단일 default.cu 파일에서 CAKE_FMHA_PAGE_SIZE 매크로를 통해 처리됩니다.

6. 소형 M (Small-M) HD256 디코드 커널 클론 통합 (csrc/cake_fmha/cuda/decode_native_bf16_hd256_smallm/default.cu)

커밋 6은 6개의 소형 M BF16 디코드 컴포넌트들을 단일 default.cu 파일로 통합합니다. 이 컴포넌트들은 행(row) 수(32 또는 64)와 페이지 크기(16, 32, 64) 상수에 따라 클론되었습니다. 이제 이 상수들은 CAKE_FMHA_SMALLM_N_ROWS와 CAKE_FMHA_SMALLM_PAGE_SIZE 매크로를 통해 동적으로 설정됩니다.

7. context_bf16 아키텍처별 커널 소스 통합 (csrc/cake_fmha/cuda/context_bf16/default.cu)

커밋 7은 커밋 4에서 생성된 context_bf16의 두 아키텍처별 default.cu 파일을 하나의 파일로 병합합니다. 두 파일은 커널 본문 내의 일부 코드(스코어 타일의 행 최대값 계산 방식)만 달랐습니다. 이 변경은 __CUDA_ARCH__ 가드 아래에 아키텍처별 코드를 배치하여 하나의 소스 파일로 관리합니다.

Before:

// sm_100a/default.cu (일부)
// ...
// sm_103a/default.cu (일부)
// ...

After:

// csrc/cake_fmha/cuda/context_bf16/default.cu
// ...
#if __CUDA_ARCH__ == 1030
// ... sm_103a 전용 코드 ...
#else
// ... sm_100a 전용 코드 ...
#endif
// ...

8. 밸런스드 디코드 커널 클론 통합 (csrc/cake_fmha/cuda/decode_balanced_*)

커밋 8은 14개의 decode_balanced_* 컴포넌트 내의 클론들을 5개의 공유 default.cu 파일로 통합합니다. 이 컴포넌트들은 데이터 타입, 행 수, 페이지 크기, 최대 그룹 수 등의 인라인 상수 값에 따라 클론되었습니다. 이제 이 값들은 CAKE_FMHA_BALANCED_* 매크로를 통해 동적으로 설정됩니다. decode_balanced_fp8는 클론이 아니므로 변경되지 않습니다.

9. DCP(Decode Context-Parallel) 정적 프로그램 및 밸런스드 프로그램 클론 통합 (csrc/cake_fmha/cuda/dcp_spec/)

커밋 9는 DCP 백엔드의 정적 프로그램들과 밸런스드 패밀리의 두 프로그램을 각각 하나의 소스 파일로 통합합니다. 정적 프로그램들은 해시 이름으로 관리되었으며, 리텐션(retain) 및 분할(split) 클래스에 따라 클론되었습니다. 이제 이들은 CAKE_FMHA_DCP_* 매크로를 통해 제어됩니다. 밸런스드 패밀리의 두 커널(fold0, fold1)도 하나의 kernel.cu로 통합됩니다.

10. 정의되지 않은 DCP 런치 선언, 읽지 않는 레지스트리 메타데이터 및 오래된 문서 카운트 제거 (csrc/cake_fmha/include/cake_fmha.h, csrc/cake_fmha/cuda/dcp_spec/registry.json, docs/api/cake_fmha.rst)

이 커밋은 호스트 코드에서 사용되지 않는 코드들을 정리합니다. 정의되지 않은 3개의 DCP 런치 선언과 DCP 레지스트리에서 읽지 않는 메타데이터가 제거되었습니다. 또한, 문서의 오래된 카운트 정보도 수정되었습니다. 이 변경은 순수 호스트 코드에만 영향을 미칩니다.

11. decode_quant_fp8 아키텍처별 커널 소스 통합 (csrc/cake_fmha/cuda/decode_quant_fp8/default.cu)

커밋 11은 decode_quant_fp8 컴포넌트의 sm_100a와 sm_103a용 소스 파일을 커밋 7의 context_bf16과 유사한 방식으로 하나의 파일로 병합합니다. TMA 로드 헬퍼 및 커널 본문의 아키텍처별 차이는 __CUDA_ARCH__ 가드 아래에 배치됩니다.

12. b256 싱크 멤버 호스트 바인딩 통합 (csrc/cake_fmha/bindings/cake_fmha_decode_native_bf16_sink_peer_clc_binding.cu)

커밋 12는 decode_native_bf16의 b256 싱크 멤버에 대한 두 개의 아키텍처별 호스트 바인딩 파일을 하나의 파일로 통합합니다. 두 파일의 유일한 차이는 동적 공유 메모리 크기 리터럴이었습니다. 이제 이 크기는 CAKE_FMHA_SINK_PEER_CLC_DYNAMIC_SMEM_BYTES 매크로를 통해 로더로부터 전달받습니다.

Before:

// sm_100a_binding.cu
config.dynamicSmemBytes = 178176;

// sm_103a_binding.cu
config.dynamicSmemBytes = 210944;

After:

// binding.cu
#ifndef CAKE_FMHA_SINK_PEER_CLC_DYNAMIC_SMEM_BYTES
#error ...
#endif
// ...
config.dynamicSmemBytes = CAKE_FMHA_SINK_PEER_CLC_DYNAMIC_SMEM_BYTES;

13. FP8 전체 블록 디코드 경로 GPU 테스트 추가 (tests/attention/test_cake_fmha.py)

커밋 13은 decode_quant_fp8 경로에 대한 GPU 정확성 테스트를 추가합니다. 이 경로는 이전에 호스트 전용 테스트만 가지고 있었습니다. 새로운 테스트는 cake_batch_decode_with_kv_cache를 실행하고 FP32 참조 및 FlashInfer의 페이징 디코드 래퍼와 비교합니다. 현재 이 경로는 버그가 있는 것으로 확인되어 테스트는 예상 실패(expected failure)로 표시됩니다.

14. 호스트 헬퍼 코드 공유화 (csrc/cake_fmha/include/cake_fmha_*.inc)

마지막 커밋은 17개의 JIT 어댑터에 중복되어 있던 호스트 헬퍼 함수들을 세 개의 공유 include 파일(cake_fmha_host_helpers.inc, cake_fmha_tma_device_arena.inc, cake_fmha_tma_device_slot_uses.inc)로 추출합니다. 이 변경은 코드 중복을 제거하고 유지보수성을 향상시키지만, 런타임 동작에는 영향을 주지 않습니다. 단, 오류 메시지에서 파일 이름이 변경되는 부수적인 효과가 있습니다.

왜 이게 좋은가?

이 PR은 여러 측면에서 상당한 개선을 가져왔습니다.

1. 코드 중복 제거 및 유지보수성 향상: 가장 큰 변화는 유사한 기능을 수행하는 수많은 CUDA 커널 클론들을 단일 소스로 통합한 것입니다. 예를 들어, decode_native_bf16의 11개 제네릭 멤버가 1개의 default.cu 파일로 통합되었습니다. 이는 총 수만 라인의 코드 중복을 제거했으며, 버그 수정이나 기능 개선 시 여러 파일을 수정해야 하는 번거로움을 없애고, 실수로 인한 불일치 발생 가능성을 줄입니다. 이는 장기적으로 프로젝트의 유지보수 비용을 크게 절감합니다.
2. 컴파일 시간 단축: 코드 중복이 줄어들면 컴파일해야 할 소스 파일의 수가 감소합니다. 특히 대규모 CUDA 프로젝트에서는 이러한 변화가 전체 빌드 시간을 눈에 띄게 단축시킬 수 있습니다. 이는 개발 워크플로우의 효율성을 높입니다.
3. 성능 최적화 및 버그 수정: 첫 번째 커밋에서 LSE 버퍼 관련 버그를 수정하고, return_lse=True 옵션에서도 최적화된 경로를 사용할 수 있도록 개선했습니다. 또한, 커널 통합 과정에서 불필요한 코드나 비효율적인 부분이 정리되었을 가능성이 있습니다. PR 설명에 따르면, 이 PR의 커널들은 당시 SOTA(State-of-the-Art) 베이스라인을 능가하는 성능을 보였다고 합니다.
4. 테스트 커버리지 확대: FP8 전체 블록 디코드 경로에 대한 GPU 테스트가 추가되었습니다. 비록 현재는 버그로 인해 실패하지만(expected failure), 이는 해당 경로의 정확성을 검증하고 향후 버그 수정의 기반을 마련합니다. 또한, 코드 변경의 정확성을 검증하기 위한 다양한 단위 테스트 및 통합 테스트가 추가되거나 수정되었습니다.
5. 코드베이스의 명확성 증대: 커널 클론들을 통합하고, 사용되지 않는 코드를 제거하며, 호스트 헬퍼 함수들을 공유화함으로써 코드베이스의 전반적인 구조가 더 명확해지고 이해하기 쉬워졌습니다. 이는 신규 개발자가 프로젝트에 기여하는 데 드는 진입 장벽을 낮춥니다.

성능 수치:

PR 설명에 따르면, 이 PR에 포함된 커널들은

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html

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