[flashinfer] FlashInfer BatchAttention: K/V 스트라이드 최적화를 통한 성능 향상

2026년 9월 29일수정: 2026년 9월 29일

PR 링크: flashinfer-ai/flashinfer#5329 상태: Merged | 변경: +1668 / -214

들어가며

최근 대규모 언어 모델(LLM)의 발전과 함께 추론 성능 최적화는 매우 중요한 과제가 되었습니다. 특히 Transformer 아키텍처의 핵심 연산인 어텐션 메커니즘은 계산량과 메모리 대역폭 요구량이 높아 병목 현상의 주된 원인이 됩니다. FlashInfer는 이러한 문제를 해결하기 위해 고안된 고성능 추론 라이브러리로, 다양한 최적화 기법을 통해 LLM 추론 속도를 크게 향상시키고 있습니다.

이번 글에서는 FlashInfer의 BatchAttention 연산에서 K/V 캐시의 메모리 스트라이드(stride)가 일치하는 특별한 경우를 위한 최적화 PR을 분석합니다. 기존에는 K/V 캐시의 스트라이드가 다르든 같든 동일한 연산 경로를 거쳤지만, 이 PR은 K/V 스트라이드가 동일할 때 이를 특화하여 불필요한 오버헤드를 줄이고 성능을 개선하는 방안을 제시합니다. 이를 통해 특히 B200 GPU에서 평균 1.0456배의 속도 향상을 달성했으며, 이는 LLM 추론 시 K/V 캐시 접근 패턴의 중요성을 다시 한번 강조합니다.

코드 분석

이번 PR의 핵심은 BatchAttention 커널에서 K/V 캐시의 메모리 레이아웃, 특히 스트라이드 정보를 활용하여 연산 경로를 분기하는 것입니다. K/V 캐시의 스트라이드가 동일한 경우, 오프셋 계산 및 프리페치 로직을 공유하여 효율성을 높입니다. 변경 사항은 주로 C++ CUDA 코드와 Python API 레벨에서 이루어졌습니다.

1. `csrc/batch_attention.cu` 및 `csrc/batch_attention_paged_kernel_inst.jinja`

가장 핵심적인 변경은 BatchPagedAttentionPersistent 커널의 템플릿 인자에 EQUAL_KV_STRIDES라는 새로운 불리언(boolean) 타입의 파라미터가 추가된 것입니다. 이 파라미터는 K와 V 캐시 텐서의 스트라이드가 동일한지 여부를 나타냅니다.

Before:

 template <uint32_t CTA_TILE_Q_1, uint32_t CTA_TILE_Q_2, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO,
-          MaskMode MASK_MODE, typename AttentionVariant, typename Params>
+          MaskMode MASK_MODE, bool EQUAL_KV_STRIDES, typename AttentionVariant, typename Params>
   cudaError_t BatchPagedAttentionPersistent(
       const Params params_1, const Params params_2,
       const uint32_t num_blks_x, const uint32_t num_blks_y,

After:

 cudaError_t status;
 #if BATCH_ATTENTION_KV_STRIDE_MODE == BATCH_ATTENTION_KV_STRIDE_MODE_RUNTIME
         if (equal_kv_strides) {
-          status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, AttentionVariant>(
+          status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE, true, AttentionVariant>(
               params[0], params[1], plan_info.num_blks_x, plan_info.num_blks_y, stream);
         } else {
-          status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE>(
+          status = BatchPagedAttentionPersistent<128, 16, HEAD_DIM_QK, HEAD_DIM_VO, MASK_MODE,
+                                                 false, AttentionVariant>(
               params[0], params[1], plan_info.num_blks_x, plan_info.num_blks_y, stream);
         }
 #elif BATCH_ATTENTION_KV_STRIDE_MODE == BATCH_ATTENTION_KV_STRIDE_MODE_EQUAL

csrc/batch_attention.cu 파일 내에서 BatchPagedAttentionRun 함수는 K와 V 캐시 텐서의 스트라이드를 비교하여 equal_kv_strides 변수를 설정합니다. 이 변수를 기반으로 BatchPagedAttentionPersistent 커널을 호출할 때 true 또는 false를 EQUAL_KV_STRIDES 템플릿 인자로 전달합니다. 또한, csrc/batch_attention_customize_config.jinja 파일에는 BATCH_ATTENTION_KV_STRIDE_MODE 매크로가 추가되어 컴파일 타임에 어떤 스트라이드 모드를 사용할지 결정할 수 있게 합니다 (RUNTIME, EQUAL, INDEPENDENT).

csrc/batch_attention_paged_kernel_inst.jinja 파일에서는 EQUAL_KV_STRIDES 템플릿 인자에 대한 커널 인스턴스화가 추가되었습니다. 이를 통해 컴파일러는 두 가지 경로 (스트라이드가 같거나 다른 경우)에 대해 최적화된 커널을 생성할 수 있습니다.

2. `flashinfer/aot.py` 및 `flashinfer/attention/_core.py`

Python API 레벨에서는 이 C++ 변경 사항을 활용하기 위한 로직이 추가되었습니다.

flashinfer/aot.py에서는 _gen_batch_attention_primary_module 함수가 _gen_batch_attention_independent_module 함수와 함께 사용됩니다. 이는 각각 동일 스트라이드와 비동일 스트라이드 경우에 대한 커널을 생성하는 역할을 합니다.
flashinfer/attention/_core.py의 BatchAttention 클래스에는 _LazyBatchAttentionIndependentModule이라는 새로운 내부 클래스가 도입되었습니다. 이 클래스는 K/V 스트라이드가 다를 때 사용될 독립적인 모듈을 지연 로딩(lazy loading)하는 역할을 합니다.

Before (기존 plan 메서드 일부):

        self.module = get_holistic_attention_module(*get_module_args)

After (수정된 plan 메서드 일부):

        self.module = get_holistic_attention_module(*get_module_args)
        self._independent_module = get_holistic_attention_independent_module(
            *get_module_args
        )

plan 메서드에서 이제 두 가지 모듈 (self.module - 기본/동일 스트라이드용, self._independent_module - 비동일 스트라이드용)을 모두 준비합니다. run 메서드에서는 K와 V 캐시의 스트라이드를 비교하여 적절한 모듈을 선택하여 실행합니다.

Before (run 메서드 일부):

        self.module.run(
            self.float_workspace_buffer,
            self.int_workspace_buffer,
            self._plan_info,
            q, kv_cache, out, lse
        )

After (run 메서드 일부):

        module = self.module
        if k_cache.stride()[:3] != v_cache.stride()[:3]:
            module = self._independent_module.get(q.device)

        module.run(
            self.float_workspace_buffer,
            self.int_workspace_buffer,
            self._plan_info,
            q, kv_cache, out, lse
        )

또한, prewarm_paged_kv_stride_variant 함수가 추가되어 CUDA 그래프 캡처 전에 독립적인 스트라이드 변형 모듈을 미리 로드할 수 있게 되었습니다. 이는 JIT 컴파일이 비활성화된 환경에서 특히 유용합니다.

왜 이게 좋은가?

이 PR의 주요 목표는 K/V 캐시 스트라이드가 일치하는 경우, 즉 K와 V 텐서가 동일한 메모리 레이아웃을 가질 때 BatchAttention 연산의 성능을 향상시키는 것입니다. 이러한 최적화는 다음과 같은 이유로 중요합니다:

불필요한 오버헤드 제거: 기존에는 K/V 스트라이드가 다르든 같든 동일한 커널 로직을 사용했습니다. 스트라이드가 같은 경우, K와 V에 대해 별도의 오프셋 계산이나 메모리 접근 패턴을 최적화할 필요가 없습니다. 이 PR은 이러한 중복 계산을 제거하고, 더 간결하고 효율적인 커널 경로를 사용함으로써 성능을 개선합니다.
하드웨어 특화 최적화: 최신 GPU 아키텍처는 메모리 접근 패턴에 매우 민감합니다. K/V 스트라이드가 일치하는 경우는 실제 LLM 추론 시나리오에서 흔하게 발생할 수 있으며, 이 경우를 특화하면 GPU의 메모리 서브시스템을 더 효율적으로 활용할 수 있습니다. B200 GPU에서 관찰된 1.0456배의 기하 평균 속도 향상은 이러한 특화의 효과를 명확히 보여줍니다.
코드 복잡성 관리: 별도의 특화된 커널을 사용함으로써, 일반적인 경우(비동일 스트라이드)와 특별한 경우(동일 스트라이드)를 분리하여 코드의 가독성과 유지보수성을 높일 수 있습니다. 또한, _LazyPagedKVStrideModule과 같은 지연 로딩 메커니즘은 필요한 모듈만 컴파일하고 로드하여 컴파일 시간을 줄이고 메모리 사용량을 최적화합니다.

성능 결과:

PR 설명에 따르면, B200 GPU에서 30개의 테스트 케이스에 대해 기하 평균 1.0456배의 속도 향상을 보였습니다. H100 GPU에서는 상대적으로 적은 성능 향상이 관찰되었는데, 이는 H100 아키텍처가 이미 높은 성능을 제공하거나 해당 워크로드에서 K/V 스트라이드 일치로 인한 이득이 상대적으로 적기 때문일 수 있습니다. 리뷰 댓글에서도 H100에서는 런타임 이득이 미미하다고 언급되었습니다.

일반적 교훈:

데이터 레이아웃의 중요성: 딥러닝 모델의 성능은 데이터가 메모리에 어떻게 저장되고 접근되는지에 크게 좌우됩니다. 특히 K/V 캐시와 같이 반복적으로 접근되는 데이터 구조의 메모리 레이아웃 최적화는 추론 성능에 직접적인 영향을 미칩니다.
특수화(Specialization)의 힘: 일반적인 경우를 처리하는 단일 커널보다, 특정 조건(예: 동일한 스트라이드)을 만족하는 경우를 위한 특화된 커널이 훨씬 더 높은 성능을 낼 수 있습니다. 이는 GPU 커널 최적화의 고전적인 접근 방식입니다.
지연 로딩 및 AOT 컴파일: 모든 가능한 커널 조합을 미리 컴파일하는 것은 비효율적입니다. 필요할 때만 커널을 컴파일하고 로드하는 지연 로딩(lazy loading) 또는 사전 컴파일(Ahead-of-Time, AOT) 방식은 라이브러리 크기를 줄이고 초기 로딩 시간을 단축하는 데 도움이 됩니다.

리뷰 댓글 분석

리뷰 댓글에서는 주로 두 가지 중요한 이슈가 제기되었습니다.

SM 7.0/7.1 (Ampere 이전) GPU에서의 too many blocks in cooperative launch 오류: qsang-nv 리뷰어는 SM 7.0/7.1 아키텍처에서 BatchAttention 커널이 과도한 블록을 할당하여 발생하는 오류를 지적했습니다. 이는 커널이 GPU의 동시 실행 블록 수를 고려하지 않고 고정된 수의 블록을 계획하기 때문입니다. 이 문제는 기존에도 존재했던 제약 사항이었으나, 새로운 테스트 케이스들이 이 제약 사항을 우회하지 않고 직접적으로 해당 API를 호출하면서 문제가 드러났습니다. 이 PR은 이 문제를 해결하기 위해 pytest.mark.xfail과 같은 마커를 사용하여 해당 GPU 아키텍처에서의 테스트 실패를 예상하거나, 또는 커널 계획 로직을 수정하여 이 문제를 해결해야 할 필요가 있습니다.
prewarm_paged_kv_stride_variant()의 디바이스별 로딩 문제: qsang-nv 리뷰어는 prewarm_paged_kv_stride_variant 함수가 호출된 디바이스에 맞는 커널을 로드하지 못하고, 대신 _HeterogeneousAOTModule을 반환하거나 JIT가 비활성화된 경우 MissingJITCacheError가 발생하는 문제를 지적했습니다. 이는 prewarm 함수가 실제 커널 로딩 및 빌드 로직을 제대로 호출하지 못했기 때문입니다. 특히 JIT가 비활성화된 환경에서는 AOT 캐시가 준비되지 않으면 런타임에 컴파일 오류가 발생할 수 있습니다. 이 문제는 flashinfer/attention/_core.py의 _LazyBatchAttentionIndependentModule 클래스와 plan 및 run 메서드의 로직을 수정하여 해결되었습니다. PR의 후속 커밋에서 prewarm 함수가 torch.cuda.device(self.device) 컨텍스트 내에서 호출되도록 수정되었고, run 메서드에서는 k_cache.stride()[:3] != v_cache.stride()[:3] 조건을 만족할 때 self._independent_module.get(q.device)를 명시적으로 호출하여 올바른 디바이스의 모듈을 가져오도록 변경되었습니다.

이러한 리뷰 피드백은 코드의 견고성과 다양한 하드웨어 환경에서의 호환성을 보장하는 데 중요한 역할을 했습니다.

결론

이번 FlashInfer PR은 BatchAttention 연산에서 K/V 캐시 스트라이드가 일치하는 경우를 특화하여 성능을 향상시키는 중요한 최적화를 도입했습니다. B200 GPU에서의 실질적인 속도 향상은 이러한 세밀한 최적화가 LLM 추론 성능에 미치는 긍정적인 영향을 보여줍니다. 또한, 리뷰 과정에서 발견된 다양한 하드웨어 및 JIT 환경에서의 문제점들을 해결함으로써 라이브러리의 안정성과 범용성을 더욱 높였습니다. K/V 캐시 접근 패턴의 중요성을 다시 한번 확인시켜주는 사례이며, 앞으로도 FlashInfer가 제공할 더욱 정교한 최적화 기법들을 기대하게 합니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [flashinfer] FlashInfer: GB200/GB300 NVL72 환경에서의 LatentMoE 추론 최적화
현재글 : [flashinfer] FlashInfer BatchAttention: K/V 스트라이드 최적화를 통한 성능 향상
다음글 [flashinfer] FlashInfer Kimi-K3 TP12 LatentMoE Tail 최적화: 8개 토큰까지 융합 및 FP64 게이팅 도입

댓글

관련 포스트

PR Analysis 의 다른글