[flashinfer] FlashInfer의 GDN 커널 런칭 오버헤드 80% 절감하기: 호스트 측 최적화 전략

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

PR 링크: flashinfer-ai/flashinfer#4374 상태: Merged | 변경: +585 / -300

들어가며

고성능 LLM 추론 엔진인 FlashInfer는 다양한 커널 최적화를 통해 낮은 지연 시간을 제공합니다. 특히 GDN(Gated Delta Network) 커널과 같은 복잡한 연산은 호스트(CPU)에서 디바이스(GPU)로 커널을 런칭하는 과정에서 상당한 오버헤드가 발생합니다. 본 PR은 SM90 및 SM120 아키텍처 환경에서 커널 런칭 오버헤드를 약 80% 감소시킨 최적화 사례를 다룹니다. 기존 434.7us였던 런칭 지연 시간을 98.2us까지 낮춘 핵심 전략을 살펴봅니다.

코드 분석

이번 최적화의 핵심은 '반복적인 객체 생성 방지'와 '캐싱 메커니즘의 고도화'입니다.

1. 커널 객체 캐싱 (Kernel Object Caching)

기존에는 매번 커널 객체를 새로 생성했으나, functools.cache를 사용하여 커널 인스턴스를 재사용하도록 변경되었습니다.

# Before
kernel = CPDeltaRuleTPrecomputeSm120(kernel_dtype)

# After
@functools.cache
def _get_t_precompute_kernel(kernel_dtype):
    return CPDeltaRuleTPrecomputeSm120(kernel_dtype)

2. 워크스페이스 버퍼 재사용

매번 torch.empty를 호출하여 메모리를 할당하는 대신, _get_cp_workspace 유틸리티를 통해 기존에 할당된 버퍼를 재사용함으로써 메모리 할당 오버헤드를 제거했습니다.

# Before
t = torch.empty((total_t_blocks, num_sab_heads, 64, 64), dtype=k.dtype, device=k.device)

# After
t = _get_cp_workspace("gdn_cp_sm120_t", (total_t_blocks, num_sab_heads, 64, 64), k.dtype, device)

3. 컴파일 옵션 및 커널 런칭 최적화

get_cached_compile을 도입하여 컴파일된 바이너리가 캐시에 존재할 경우 즉시 런칭하도록 로직을 개선했습니다.

# After
compiled = get_cached_compile(kernel, compile_options)
if compiled is None:
    # ... (compile logic)
    compiled = cached_compile(kernel, *kernel_args, compile_options=compile_options)
compiled(k_tma, ...)

왜 이게 좋은가

이 최적화는 '호스트 측 런칭 오버헤드'를 타겟팅합니다. GPU 연산 자체의 속도도 중요하지만, 짧은 연산이 반복되는 경우 호스트의 런칭 오버헤드가 전체 성능의 병목이 됩니다.

메모리 할당 비용 절감: torch.empty는 동기화가 필요한 메모리 할당을 유발할 수 있는데, 이를 캐싱된 버퍼로 대체하여 오버헤드를 제거했습니다.
객체 생성 오버헤드 제거: Python 레벨에서 CuTe 객체들을 매번 생성하는 것은 비용이 큽니다. functools.cache를 통해 이를 정적 객체로 관리하여 런타임 비용을 최소화했습니다.

결과적으로 SM120 환경에서 434.7us에서 98.2us로 약 80%의 성능 향상을 달성했습니다. 이는 실시간 추론 시스템에서 매우 유의미한 수치입니다.

교훈

Hot Path 최적화: 커널 런칭과 같이 빈번하게 호출되는 코드는 객체 생성과 메모리 할당을 최대한 피해야 합니다.
Caching 전략: 단순히 컴파일된 바이너리뿐만 아니라, 이를 런칭하기 위한 '커널 객체'와 '워크스페이스 버퍼'까지 캐싱하는 것이 전체 파이프라인의 속도를 결정짓습니다.