본문으로 건너뛰기

[flashinfer] FlashInfer Kimi-K3 TP12 LatentMoE Tail 최적화: 8개 토큰까지 융합 및 FP64 게이팅 도입

PR 링크: flashinfer-ai/flashinfer#5668 상태: Merged | 변경: +5066 / -575

들어가며

최근 대규모 언어 모델(LLM)의 발전은 모델의 크기와 복잡성을 기하급수적으로 증가시키고 있습니다. 이러한 모델을 효율적으로 학습하고 서빙하기 위해서는 연산의 최적화가 필수적입니다. 특히, Mixture-of-Experts (MoE)와 같은 희소 활성화(sparse activation) 기법은 모델의 표현력을 높이는 동시에 연산량을 증가시키는 요인이 됩니다.

FlashInfer는 LLM 추론을 위한 고성능 커널 라이브러리로, 지속적으로 최신 하드웨어와 모델 아키텍처에 최적화된 커널을 개발하고 있습니다.

본 블로그 글에서는 FlashInfer 레포지토리의 flashinfer-ai/flashinfer에서 진행된 PR(#5666)을 분석하여, Kimi-K3 TP12 LatentMoE 연산의 성능을 크게 향상시킨 최적화 기법에 대해 자세히 알아보겠습니다. 이 PR은 특히 CAKE 팀에서 개발한 커널을 기반으로 하며, 기존의 State-of-the-Art (SOTA) 성능을 뛰어넘는 결과를 보여주었습니다.

주요 변경 사항은 다음과 같습니다:

  1. K23 연산의 토큰 처리 능력 확장: 최대 8개의 토큰까지 처리할 수 있도록 융합(fusion) 기능을 확장했습니다.
  2. 정확도 향상: FP64 게이팅을 도입하여 수치적 정확도를 개선했습니다.

이 PR이 어떻게 이러한 성능 향상을 달성했는지, 그리고 그 원리가 무엇인지 코드 diff와 함께 상세히 살펴보겠습니다.

코드 분석

이번 PR의 핵심 변경 사항은 flashinfer/experimental/kimi_k3_tp12_tail/cake_backend.py 파일과 관련 CUDA 커널 파일들(csrc/cake_kimi_k3_tp12_tail/)에 집중되어 있습니다. 주요 변경 내용을 파일별로 나누어 살펴보겠습니다.

1. flashinfer/experimental/kimi_k3_tp12_tail/cake_backend.py

이 파일은 Kimi-K3 TP12 LatentMoE 연산의 백엔드 로직을 정의합니다. 이번 PR에서는 K23 연산의 처리 능력을 확장하고, 이를 위한 새로운 용량 사다리(capacity ladder)를 도입했습니다.

주요 변경점:

  • K23_MAX_TOKENS 증가: 기존에는 최대 4개의 토큰까지만 융합된 K23 연산을 지원했지만, 이제는 최대 8개까지 지원합니다.

    -K23_MAX_TOKENS = 4
+K23_MAX_TOKENS = 8

  • K23_LADDER 도입: K23 연산이 처리할 수 있는 토큰 수에 따라 다른 용량의 커널을 사용하도록 정의합니다. K23_LADDER = (4, 8)는 4개 이하의 토큰은 4개 용량 커널을, 5개에서 8개 사이의 토큰은 8개 용량 커널을 사용하도록 합니다. 이는 k23_capacity_for 함수에서 활용됩니다.

    +#: K23 accumulator capacities compiled side by side (round 6, CAKE-740): ``M`` runs on the smallest capacity ``>= M``, so
+#: ``M <= 4`` keeps the round-5 four-token module bit for bit and ``5 <= M <= 8`` takes the eight-token module.  The
+#: eight-token rows round once (fp32 slice + shared sum -> BF16) where the round-5 chain rounded twice (cuBLAS BF16 slice,
+#: then the add): paired against an fp64 reference their mean / max absolute error is 0.78x / 0.64-0.93x of the round-5
+#: chain on every correctness shape (Cake design doc, round 6).
+K23_LADDER = (4, 8)

  • k23_capacity_for 함수 추가: 입력된 토큰 수(num_tokens)에 따라 적절한 K23 용량(c<capacity>)을 반환하는 함수입니다. 이는 커널 키 생성 시 사용됩니다.

    +def k23_capacity_for(num_tokens: int) -> int:
+    """The K23 accumulator capacity serving `num_tokens` (`<= K23_MAX_TOKENS`): the smallest ladder rung `>= M`."""
+    if not 1 <= num_tokens <= K23_MAX_TOKENS:
+        raise ValueError(
+            f"num_tokens must be in [1, {K23_MAX_TOKENS}] for the fused K23 tail, got {num_tokens}"
+        )
+    return next(c for c in K23_LADDER if c >= num_tokens)

  • 커널 키 생성 로직 변경: k3_kernel_key 함수에서 k23 폼을 사용할 때, 이제 새로운 k23_capacity_for 함수를 호출하여 c<capacity> 부분을 동적으로 생성합니다.

    -        return f

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글