본문으로 건너뛰기

[flashinfer] [FlashInfer] Paged Attention 최적화: 동일 Stride 구조에서의 주소 계산 오버헤드 제거

PR 링크: flashinfer-ai/flashinfer#4736 상태: Merged | 변경: +1411 / -206

들어가며

LLM 추론 엔진의 핵심인 Attention 연산에서 Paged KV Cache는 메모리 효율성을 극대화하는 표준 기술입니다. 하지만 최근 FlashInfer 라이브러리에서 Paged FlashAttention 2(FA2) 구현 중, Key(K)와 Value(V)의 메모리 보폭(Stride)이 동일함에도 불구하고 각각 독립적인 주소 계산(Address Generation)을 수행하여 성능이 저하되는 회귀(Regression) 현상이 발견되었습니다.

이 PR은 K와 V의 Stride가 동일한 일반적인 케이스를 위해 컴파일 타임 특수화(Compile-time Specialization)를 도입하여 불필요한 연산을 제거하고, 성능을 이전 수준으로 회복한 과정을 담고 있습니다. 시니어 엔지니어의 관점에서 이 최적화가 왜 중요한지, 그리고 어떤 트레이드오프를 해결했는지 분석해 보겠습니다.

코드 분석: 무엇이 바뀌었는가?

1. batch_pod.cu: 유연한 Stride 처리

기존 코드에서는 K와 V의 Stride가 반드시 같아야 한다는 엄격한 체크(TVM_FFI_ICHECK_EQ)가 있었고, 하나의 Stride 포인터만 사용했습니다. 변경 후에는 K와 V의 Stride를 독립적으로 받아들이되, 이를 paged_kv_t 구조체에 전달하여 하위 커널에서 최적화 여부를 결정할 수 있게 했습니다.

Before:

// K와 V의 Stride가 무조건 같아야 함을 강제함
auto k_strides_p = paged_k_cache_p.strides();
auto v_strides_p = paged_v_cache_p.strides();
TVM_FFI_ICHECK_EQ(k_strides_p.size(), v_strides_p.size());
for (int i = 0; i < k_strides_p.size(); ++i) {
  TVM_FFI_ICHECK_EQ(k_strides_p[i], v_strides_p[i]);
}
kv_cache_strides_p = k_strides_p.data();

After:

// 독립적인 K/V Stride를 허용하도록 변경
auto k_strides_p = paged_k_cache_p.strides();
auto v_strides_p = paged_v_cache_p.strides();
TVM_FFI_ICHECK_EQ(k_strides_p.size(), v_strides_p.size());
// ... 중략 ...
paged_kv_t<DTypeKV, IdType> paged_kv(
    num_kv_heads_p, page_size_p, HEAD_DIM_VO, batch_size_p, kv_layout_p,
    static_cast<DTypeKV*>(paged_k_cache_p.data_ptr()),
    static_cast<DTypeKV*>(paged_v_cache_p.data_ptr()), 
    k_strides_p.data(), // K stride 전달
    v_strides_p.data(), // V stride 전달
    static_cast<IdType*>(paged_kv_indices_p.data_ptr()),
    // ...
);

2. batch_prefill.cu: 템플릿 특수화 도입

핵심은 SAME_KV_STRIDES라는 템플릿 파라미터의 도입입니다. 이를 통해 GPU 커널 내부에서 런타임 분기(Branch) 없이, 컴파일 타임에 V의 주소 계산을 K의 결과로 재사용할지 결정합니다.

Before:

template <uint32_t CTA_TILE_Q, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO,
          PosEncodingMode POS_ENCODING_MODE, bool USE_FP16_QK_REDUCTION, MaskMode MASK_MODE,
          typename AttentionVariant, typename Params>
cudaError_t BatchPrefillWithPagedKVCacheDispatched(...)

After:

template <bool SAME_KV_STRIDES, uint32_t CTA_TILE_Q, uint32_t HEAD_DIM_QK, uint32_t HEAD_DIM_VO,
          PosEncodingMode POS_ENCODING_MODE, bool USE_FP16_QK_REDUCTION, MaskMode MASK_MODE,
          typename AttentionVariant, typename Params>
cudaError_t BatchPrefillWithPagedKVCacheDispatched(...)

이 변경을 통해 SAME_KV_STRIDES가 true일 경우, 커널은 V 오프셋을 계산하기 위해 추가적인 레지스터를 소모하거나 산술 연산을 수행하지 않고 K 오프셋을 그대로 '복사'해서 사용합니다.

왜 이게 좋은 최적화인가?

1. 레지스터 압박(Register Pressure) 감소

리뷰어 qsang-nv의 분석에 따르면, SAME_KV_STRIDES=true 특수화를 적용했을 때 특정 설정(NUM_MMA_KV=1)에서 레지스터 사용량이 163개에서 128개로 감소했습니다. GPU 커널에서 레지스터 사용량 감소는 곧 Occupancy(점유율) 상승으로 이어지며, 이는 대규모 병렬 처리가 중요한 Attention 연산에서 직접적인 성능 향상을 가져옵니다.

2. AGU(Address Generation Unit) 부하 경감

현대 GPU 아키텍처에서 메모리 주소를 계산하는 연산은 단순해 보이지만, 수조 번 반복되는 루프 내에서는 상당한 오버헤드입니다. K와 V가 동일한 레이아웃을 가질 때 주소 계산 결과를 공유함으로써, 하드웨어의 AGU 부하를 줄이고 실제 데이터 로드 연산에 더 많은 자원을 할당할 수 있습니다.

3. 실제 성능 향상 수치

PR 작성자가 공유한 벤치마크 결과는 놀랍습니다.

Case Backend Before (ms) After (ms) 개선율
Llama-3.1-70B decode fa2_tc 0.413 0.335 ~18.8%
GPT-OSS prefill fa2 1.445 1.359 ~5.9%

단순히 주소 계산 로직 하나를 특수화했을 뿐인데, 전체 Decode 성능이 약 19% 가까이 개선되었습니다.

엔지니어링 트레이드오프: 바이너리 크기 vs 성능

이 최적화의 유일한 단점은 바이너리 팽창(Binary Bloat)입니다. 모든 조합에 대해 true/false 두 가지 버전을 컴파일해야 하므로, 오브젝트 파일 크기가 약 78% 증가하는 현상이 보고되었습니다.

FlashInfer 팀은 이를 해결하기 위해 영리한 전략을 선택했습니다:

  1. Default: 가장 흔한 케이스인 SAME_KV_STRIDES=true 버전만 기본 바이너리에 포함합니다.
  2. JIT/Lazy Loading: 드문 케이스인 independent stride 버전은 런타임에 필요할 때만 JIT(Just-In-Time) 컴파일하거나 별도의 모듈로 지연 로딩(Lazy Loading)합니다.

이 방식을 통해 배포되는 Wheel 파일의 크기는 오히려 1.2% 줄이면서도, 성능 이점은 그대로 챙길 수 있었습니다.

결론 및 교훈

  1. Host-side Dispatch의 힘: GPU 커널 내부에서 if 문으로 Stride를 체크하는 대신, 호스트에서 미리 판단하여 최적화된 커널을 실행(Dispatch)하는 것이 성능 면에서 훨씬 유리합니다.
  2. 사소한 중복의 무서움: "주소 계산 한두 번 더 하는 게 대수냐"라고 생각할 수 있지만, 수천 명의 유저가 사용하는 LLM 서빙 환경에서는 18%의 성능 차이를 만드는 결정적 요인이 됩니다.
  3. 트레이드오프 관리: 성능을 위해 복잡성을 도입할 때는 반드시 바이너리 크기나 유지보수 비용을 고려해야 하며, JIT와 같은 하이브리드 전략이 훌륭한 대안이 될 수 있습니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글