[flashinfer] FlashInfer, CUDA 커널 최적화를 통한 LLM 추론 속도 향상: Recurrence-Piece Persistent M128 도입

2026년 8월 26일수정: 2026년 8월 26일

PR 링크: flashinfer-ai/flashinfer#4728 상태: Merged | 변경: +4074 / -9

들어가며

최근 대규모 언어 모델(LLM)의 발전은 놀라운 속도로 이루어지고 있지만, 이러한 모델을 실제 서비스에 적용하기 위해서는 추론 속도와 효율성이 중요한 과제로 남아있습니다. 특히, LLM의 핵심 연산 중 하나인 어텐션 메커니즘은 계산량이 많아 병목 현상을 일으키기 쉽습니다. NVIDIA의 FlashInfer 라이브러리는 이러한 문제를 해결하기 위해 CUDA 커널 수준의 최적화를 통해 LLM 추론 성능을 극대화하는 데 집중하고 있습니다. 이번 글에서는 FlashInfer의 최신 Pull Request(PR)에서 제안된 perf(cake_kda): add recurrence-piece persistent M128 prefill 변경 사항을 분석하여, 어떻게 CUDA 커널 최적화를 통해 LLM 추론 성능을 획기적으로 개선했는지 살펴보겠습니다.

이 PR은 특히 recurrent_kda 함수의 prefill 단계에서 cake 백엔드를 사용할 때, 특정 조건 하에서 recurrent chain의 일부만 처리하는 'piece-persistent' 기법을 도입하여 성능을 향상시키는 것을 목표로 합니다. 이는 기존의 frozen #4605 대비 최대 2.7배의 속도 향상을 가져왔으며, 이는 LLM 추론의 효율성을 한 단계 끌어올릴 수 있는 중요한 발전입니다.

코드 분석: Recurrence-Piece Persistent M128 도입

이번 PR의 핵심은 csrc/kda/cake_flashkda_bf16_piece_persistent_m128.cu 파일에 새롭게 추가된 CUDA 커널입니다. 이 커널은 기존의 recurrent_kda prefill 경로를 최적화하여, 특정 하드웨어 및 워크로드에서 더 나은 성능을 제공합니다.

1. 새로운 CUDA 커널 파일: `cake_flashkda_bf16_piece_persistent_m128.cu`

이 파일은 piece-persistent M128 특화 커널을 정의합니다. 기존의 커널과 달리, 이 커널은 recurrent chain의 일부만 처리하고 중간 상태를 효율적으로 관리하는 데 중점을 둡니다. 이는 다음과 같은 최적화 기법을 활용합니다:

Piece-Persistent Execution: Recurrence chain의 모든 부분을 한 번에 처리하는 대신, 필요한 'piece'만 처리하고 중간 상태를 유지(persistent)하여 재사용합니다. 이는 불필요한 계산을 줄이고 메모리 접근을 최적화합니다.
M128 Specialization: 특정 하드웨어 아키텍처(예: Hopper)의 M128 매트릭스 곱셈 유닛을 효율적으로 활용하도록 최적화되었습니다. BF16 데이터 타입을 사용하여 연산 효율성을 높입니다.
Device-Scope Release/Acquire Handoffs: 중간 BF16 상태를 persistent CTA(CUDA Thread Array) 간에 효율적으로 전달하기 위해 device-scope release/acquire 메커니즘을 사용합니다. 이는 스레드 간의 동기화 및 데이터 공유를 최적화합니다.
Stream-Local Workspace Reuse: 최종 consumer가 각 handoff 카운터를 리셋하여, 동일한 stream-local 작업 공간을 후속 eager 호출에서 재사용할 수 있도록 합니다. 이는 메모리 할당 및 해제 오버헤드를 줄입니다.

2. 핵심 CUDA 코드 인용 및 설명

새로 추가된 CUDA 커널 파일의 일부를 살펴보며 구체적인 최적화 내용을 확인해 보겠습니다.

// ... (기존 헤더 및 정의) ...

// Piece-persistent M128 kernel definition (simplified)
__global__ void flashkda_bf16_persistent_m128_kernel(
    // Input tensors (Q, K, V, etc.)
    const bf16_t* q_ptr, const bf16_t* k_ptr, const bf16_t* v_ptr, /* ... other inputs ... */
    // Output tensor
    bf16_t* out_ptr,
    // Runtime parameters
    uint32_t seq_len, uint32_t num_heads, uint32_t head_dim, /* ... */
    // Workspace for intermediate states
    void* workspace_ptr,
    // Synchronization primitives
    mbarrier_t* mbar_producer, mbarrier_t* mbar_consumer
) {
    // ... kernel logic ...

    // Example: Loading a piece of K and V from global memory to shared memory
    // This part would be highly optimized using tmem_ld instructions
    // and careful shared memory banking.
    // ... load K, V to SMEM ...

    // Example: Performing MMA (Matrix Multiply-Accumulate) operations
    // utilizing the M128 units for BF16.
    // This involves using tcgen05.mma instructions.
    // ... mma operations ...

    // Example: Storing intermediate results back to workspace or global memory
    // ... store intermediate results ...

    // Synchronization between producer and consumer CTAs
    // mbarrier_arrive(mbar_producer);
    // mbarrier_wait(mbar_consumer, phase);

    // ... rest of the kernel logic ...
}

위 코드는 실제 커널의 일부를 간략화하여 보여줍니다. 실제 코드는 tcgen05.mma, tcgen05.ld, tcgen05.st와 같은 NVIDIA 아키텍처별 최적화된 PTX(Parallel Thread Execution) 명령어를 사용하여 매트릭스 곱셈, 데이터 로딩 및 저장을 수행합니다. 특히 mbarrier 관련 함수들은 CTA(CUDA Thread Array) 간의 동기화를 효율적으로 관리하여, 중간 상태를 안전하게 주고받을 수 있도록 합니다.

flashkda_bf16_persistent_m128.cu 파일의 다른 부분에서는 다음과 같은 저수준 최적화가 이루어집니다:

Shared Memory Management: SMEM_SMEM_QD_OFF, SMEM_SMEM_V_OFF 등 다양한 오프셋과 크기 정의는 공유 메모리를 효율적으로 분할하고 활용하기 위한 것입니다. 이를 통해 데이터 재사용성을 높이고 전역 메모리 접근을 최소화합니다.
Synchronization Primitives: mbarrier_init, mbarrier_try_wait, mbarrier_wait 등의 함수는 mbarrier 명령어를 사용하여 CTA 간의 동기화를 관리합니다. 이는 데이터 종속성을 해결하고 병렬 실행의 정확성을 보장합니다.
Tensor Core Operations: tcgen05_mma_f16, mma_ts_step과 같은 함수들은 Tensor Core를 활용한 BF16 매트릭스 곱셈을 수행합니다. 이는 LLM의 핵심 연산인 행렬 곱셈을 하드웨어 가속을 통해 매우 빠르게 처리할 수 있게 합니다.

3. Dispatcher 로직 변경 (추정)

PR 설명에 따르면, 디스패처는 실시간 SM(Streaming Multiprocessor) 개수와 물리적 점유율/roofline 모델을 사용하여 recurrent chain을 분할합니다. 이는 새로운 piece-persistent 경로를 언제 사용할지 결정하는 중요한 로직입니다. 기존의 CUDA Graph 경로는 변경되지 않고 이전 변형을 계속 사용합니다. 이는 다음과 같은 장점을 가집니다:

Dynamic Scheduling: 하드웨어 리소스와 워크로드 특성에 따라 최적의 경로를 동적으로 선택하여 성능을 극대화합니다.
Compatibility: 기존 CUDA Graph 사용자에게는 영향을 주지 않으면서 새로운 최적화 기능을 제공합니다.

왜 이게 좋은가: 성능 향상 및 일반적 교훈

이번 PR은 LLM 추론 성능을 크게 향상시키는 몇 가지 중요한 이유를 제시합니다.

1. 획기적인 성능 향상

PR의 성능 표는 이 최적화의 효과를 명확하게 보여줍니다. 주요 결과는 다음과 같습니다:

최대 2.725x 속도 향상: GB300 GPU에서 exported API 대비 frozen #4605 (이전 최적화) 대비 약 1.045배, 그리고 새로운 FlashKDA (이번 PR의 최적화) 대비 약 2.725배의 속도 향상을 기록했습니다. 이는 동일한 하드웨어에서 훨씬 더 빠른 추론 속도를 달성할 수 있음을 의미합니다.
다양한 GPU에서의 일관된 성능: B200, B300, GB200, GB300 등 다양한 최신 NVIDIA GPU에서 일관되게 높은 성능 향상을 보여줍니다. 이는 이 최적화가 특정 하드웨어에 국한되지 않고 광범위하게 적용될 수 있음을 시사합니다.
정확성 유지: 모든 성능 측정에서 정확성(Correctness)은 Cake 29/29 또는 Cake 28 comparable + 1 N/A로 표시되어, 성능 향상이 정확성 저하를 동반하지 않음을 보장합니다.

2. 일반적 교훈

이 PR은 LLM 추론 성능 최적화에 대한 몇 가지 중요한 교훈을 제공합니다:

저수준 CUDA 최적화의 중요성: LLM과 같이 계산 집약적인 워크로드에서는 CUDA 커널 수준의 최적화가 성능 향상의 핵심입니다. Tensor Core 활용, 공유 메모리 관리, 효율적인 동기화 메커니즘은 필수적입니다.
워크로드 특성에 맞는 동적 스케줄링: 모든 워크로드가 동일한 최적화 기법에 동일하게 반응하는 것은 아닙니다. 하드웨어 리소스(SM count)와 연산 패턴(roofline model)을 기반으로 최적의 커널 경로를 동적으로 선택하는 것은 성능을 극대화하는 데 중요합니다.
점진적 개선과 호환성: 기존 코드베이스에 새로운 최적화를 도입할 때는 호환성을 유지하는 것이 중요합니다. 이 PR은 CUDA Graph 경로를 그대로 유지하면서 새로운 eager 호출에 대한 최적화를 추가하여 점진적인 개선을 이루었습니다.
메모리 접근 패턴 최적화: piece-persistent 기법과 stream-local workspace reuse는 메모리 접근 패턴을 최적화하여 성능을 향상시키는 좋은 예시입니다. 불필요한 데이터 이동을 줄이는 것이 LLM 추론 속도에 큰 영향을 미칩니다.

리뷰 댓글 분석

제공된 리뷰 댓글은 주로 CI/CD 파이프라인 실행 및 테스트 결과 확인에 관한 내용입니다. @flashinfer-bot run 및 /bot run tests/kda와 같은 명령어는 자동화된 테스트 및 빌드 프로세스를 트리거하는 데 사용됩니다. [SUCCESS] Pipeline [#64559499]: 16/16 executed test jobs passed 결과는 이 PR의 변경 사항이 기존 테스트를 통과했으며, 코드 품질 및 기능적 무결성이 검증되었음을 나타냅니다. 이는 코드 변경이 안정적이고 예상대로 작동함을 보여주는 긍정적인 신호입니다.

결론

FlashInfer의 이번 PR은 recurrence-piece persistent M128이라는 새로운 CUDA 커널 최적화를 통해 LLM 추론의 prefill 단계를 획기적으로 개선했습니다. 최대 2.7배의 속도 향상은 LLM의 실질적인 응용 가능성을 더욱 높이는 중요한 성과입니다. 저수준 CUDA 최적화, 동적 스케줄링, 효율적인 메모리 관리 기법의 조합은 앞으로 LLM 추론 성능을 개선하는 데 있어 중요한 참고 자료가 될 것입니다. FlashInfer 팀의 지속적인 노력은 AI 모델의 접근성과 효율성을 높이는 데 크게 기여하고 있습니다.

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] SGLang: LongCat-Image DiT의 FFN 연산 최적화 - Tanh-GELU 퓨전 적용
현재글 : [flashinfer] FlashInfer, CUDA 커널 최적화를 통한 LLM 추론 속도 향상: Recurrence-Piece Persistent M128 도입
다음글 [vllm] vLLM의 Hopper(SM90) 아키텍처를 위한 GEMM 커널 최적화

댓글

관련 포스트

PR Analysis 의 다른글