본문으로 건너뛰기

[flashinfer] [FlashInfer] Blackwell 아키텍처를 위한 Warp Level Split-K BF16 GEMM 최적화 분석

PR 링크: flashinfer-ai/flashinfer#4908 상태: Merged | 변경: +1099 / -74

들어가며

LLM(Large Language Model) 추론 서비스, 특히 vLLM과 같은 고성능 서빙 프레임워크에서 Decode 단계의 GEMM 성능은 전체 처리량(Throughput)과 지연 시간(Latency)에 결정적인 영향을 미칩니다. Decode 단계는 보통 배치 사이즈(M)가 매우 작고(1~32), 가중치 행렬(Weight)은 큰 전형적인 Memory-bound 워크로드입니다.

최근 NVIDIA의 차세대 아키텍처인 Blackwell(SM100/SM103)이 등장함에 따라, 기존의 Direct GEMM이나 Cluster Split-K 방식만으로는 특정 Shape(예: K=7168 등)에서 최적의 성능을 내지 못하는 경우가 발생했습니다. 이번 PR은 이를 해결하기 위해 Warp Level Split-K라는 새로운 커널 전략을 도입하고, 실제 추론 환경과 유사한 성능 측정을 위해 Autotuner의 Cold-L2 캐시 리플레이 기능을 개선한 내용을 담고 있습니다.

코드 분석

1. Warp Split-K 커널 도입 (dense_bf16_gemm_warp_splitk.py)

가장 핵심적인 변화는 K-차원을 Warp 레벨에서 분할하여 연산하는 커널의 추가입니다. 기존 방식이 CTA(Thread Block) 단위로 K를 나눴다면, 이 커널은 하나의 CTA 내에 있는 4개의 Compute Warp가 각각 K의 1/4씩 담당하여 연산한 뒤 Shared Memory 'Mailbox'를 통해 결과를 리덕션(Reduction)합니다.

또한, cp.async 명령어를 사용할 때 L2::evict_first 힌트를 활용하여, 한 번만 읽는 가중치 데이터가 L2 캐시를 오염시키지 않도록 하고, 여러 CTA가 반복해서 읽는 활성화 함수(Activation) 데이터가 L2에 상주할 수 있도록 최적화했습니다.

2. Runner 선택 로직 리팩토링 (flashinfer/gemm/gemm_base.py)

기존에는 M 값에 따라 하드코딩된 조건문으로 백엔드를 분기 처리했습니다. 이번 PR에서는 _CuteDSLBf16Runner라는 베이스 클래스를 도입하여 여러 알고리즘(Direct, Cluster Split-K, Warp Split-K)을 체계적으로 관리하도록 리팩토링되었습니다.

Before:

# 기존: M > 32일 때 즉시 cuBLASLt로 분기
if a.shape[0] > 32:
    return _cublaslt_mm_bf16_requirement(
        a, b, out, out_dtype, bias, False, "cublaslt"
    )

After:

# 변경: 하나의 Runner Set 내에서 M 범위에 따라 적절한 Runner를 선택
if a.shape[0] > _CUTE_DSL_BF16_MAX_M:
    # M > 32인 경우 cuBLASLt Fallback Runner가 처리하도록 통합
    return _cublaslt_mm_bf16_requirement(
        a, b, out, out_dtype, bias, False, "cublaslt"
    )

이러한 구조적 변화 덕분에 mm_bf16(backend="cute-dsl") 호출 시 M 값에 상관없이 일관된 인터페이스를 유지하면서도, 내부적으로는 최적의 커널을 동적으로 선택할 수 있게 되었습니다.

3. Autotuner: Cold-L2 Graph Replay (flashinfer/autotuner/autotuner.py)

GEMM 성능 측정 시 가장 흔한 실수는 'Warm 캐시' 상태에서 벤치마크를 수행하는 것입니다. 실제 추론 환경에서는 가중치가 캐시에 없는 경우가 많으므로, 이를 모사하기 위해 use_cold_l2_graph_replay 옵션이 추가되었습니다.

코드 변경 사항:

# flashinfer/autotuner/autotuner.py
if (
    tuning_config.use_cuda_graph
    and tuning_config.use_cold_l2_graph_replay
):
    # 1. 첫 번째 실행은 초기화 비용 흡수를 위해 비측정 실행
    graph.replay()
    if l2_eviction_buffer is not None:
        # 2. L2 캐시를 강제로 비움 (Cold-L2 상태 조성)
        l2_eviction_buffer.zero_()
    stream.synchronize()

이 최적화가 중요한 이유는, 캐시가 따뜻한 상태에서는 성능 차이가 미미해 보이던 알고리즘들이 실제 Cold-L2 환경에서는 큰 성능 차이를 보이기 때문입니다. 실제로 이 기능을 통해 M=9~16 구간에서 Cluster Split-K 대신 Warp Split-K가 더 우수하다는 것을 정확히 판별해낼 수 있었습니다.

왜 이게 좋은 최적화인가?

1. 압도적인 성능 향상

벤치마크 결과에 따르면, Kimi-K3 및 GLM-5.2 모델에서 주로 사용하는 Shape들에 대해 vLLM의 fused-A 커널 대비 최대 1.37x, 중앙값 1.08x의 성능 향상을 기록했습니다. 특히 K=7168과 같은 큰 K-차원 환경에서 Warp 레벨의 세밀한 병렬화가 빛을 발했습니다.

2. 하드웨어 특성(Blackwell)의 깊은 이해

ptxas 13.0의 버그(Encoding issue)를 피하기 위한 워크라운드를 적용하고, SM100/103의 메모리 계층 구조를 고려하여 L2 캐시 힌트를 정교하게 제어한 점은 시니어 엔지니어링의 정수를 보여줍니다.

3. Autotuning의 신뢰성 확보

단순히 빠른 커널을 만드는 것에 그치지 않고, 그 커널이 실제 환경에서 왜 빠른지를 증명할 수 있는 측정 도구(Cold-L2 Replay)를 개선했습니다. 이는 향후 다른 아키텍처나 Shape에 대한 최적화 시에도 강력한 기준점이 될 것입니다.

리뷰어 피드백 분석

리뷰 과정에서 vadiklyutiy는 Warp Split-K의 Tactic(설정 조합) 수가 수백 개에 달해 Autotuning 속도가 느려질 수 있음을 지적했습니다. 이는 성능과 튜닝 시간 사이의 Trade-off를 보여줍니다. 현재는 성능 극대화를 위해 전수 조사를 택했지만, 향후 Tactic Pruning(가지치기)을 통해 튜닝 시간을 단축하는 후속 작업이 논의될 수 있습니다.

또한, 현재 모든 검증이 B300(SM103)에서 이루어졌으므로, B200(SM100)에서의 추가 검증이 필요하다는 Maintainer의 의견이 반영되어 안정성을 높였습니다.

결론

이번 PR은 Blackwell 아키텍처의 잠재력을 끌어올리기 위한 정교한 커널 설계와 이를 뒷받침하는 인프라(Autotuner)의 개선을 동시에 달성했습니다. 특히 저지연이 생명인 LLM Decode 단계에서 이러한 최적화는 서비스 품질(QoS) 향상에 직접적인 기여를 할 것입니다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글