[flashinfer] FlashInfer, Kimi-K3 어텐션의 성능 최적화: SM100/SM103 아키텍처를 위한 중간 크기 라우팅 개선

2026년 10월 4일수정: 2026년 10월 4일

PR 링크: flashinfer-ai/flashinfer#6036 상태: Merged | 변경: +3406 / -4623

들어가며

최근 LLM(거대 언어 모델)의 발전은 어텐션 메커니즘의 효율적인 계산에 크게 의존하고 있습니다. 특히, FlashInfer는 Transformer 모델의 핵심인 어텐션 연산을 GPU에서 고속으로 처리하기 위한 라이브러리로 주목받고 있습니다. 이번 PR(#5994 이후)은 FlashInfer의 Kimi-K3 어텐션 커널에서 중간 크기(mid-M) 토큰 수에 대한 라우팅 전략을 최적화하여 성능을 크게 향상시키는 것을 목표로 합니다.

기존에는 특정 토큰 수 범위에서 사용되던 커널이 최적이 아닐 수 있었습니다. 이 PR은 SM100 및 SM103 아키텍처에서 다양한 토큰 수(M)와 시퀀스 길이(K) 조합에 대해 가장 효율적인 커널을 동적으로 선택하는 라우팅 로직을 개선합니다. 이를 통해 기존 vLLM의 네이티브 CUDA 커널 대비 상당한 성능 향상을 달성했습니다.

본 글에서는 이 PR의 코드 변경 사항을 상세히 분석하고, 왜 이러한 변경이 성능 향상으로 이어지는지, 그리고 기술적인 관점에서 어떤 의미를 가지는지 살펴보겠습니다.

코드 분석

1. 라우팅 로직 개선 (`cake_backend.py`)

이번 PR의 핵심은 flashinfer/experimental/cake_kimi_k3_attn_res/cake_backend.py 파일의 라우팅 로직 수정입니다. 특히, _SMALL_M_DIRECT_MAX_M 및 _SMALL_M_CLUSTER 딕셔너리의 값이 업데이트되어, 특정 M 및 K 값에 대해 더 효율적인 커널이 선택되도록 변경되었습니다.

변경 전 (_SMALL_M_DIRECT_MAX_M 예시):

# _SMALL_M_DIRECT_MAX_M = {
#     "sm_100a": {0: 16, 1: 16, 2: 16, 3: 16, 4: 16, 5: 16, 6: 16, 7: 16, 8: 16},
#     "sm_103a": {0: 16, 1: 16, 2: 16, 3: 16, 4: 16, 5: 16, 6: 16, 7: 16, 8: 16},
# }

변경 후 (_SMALL_M_DIRECT_MAX_M 예시):

# _SMALL_M_DIRECT_MAX_M = {
#     "sm_100a": {0: 256, 1: 512, 2: 256, 3: 256, 4: 128, 5: 128, 6: 128, 7: 128, 8: 64},
#     "sm_103a": {0: 512, 1: 512, 2: 256, 3: 256, 4: 128, 5: 128, 6: 128, 7: 128, 8: 64},
# }

_SMALL_M_DIRECT_MAX_M 딕셔너리는 small_m_direct 커널이 사용되는 최대 M 값을 아키텍처(sm_100a, sm_103a)와 K 값별로 정의합니다. 변경 후, 특히 K 값이 낮을 때(K=0, 1) M의 최대값이 16에서 256 또는 512로 크게 증가했습니다. 이는 더 넓은 범위의 중간 크기 토큰 수에 대해 small_m_direct 커널이 사용될 수 있음을 의미합니다.

_SMALL_M_CLUSTER 딕셔너리도 유사하게 업데이트되었습니다. 이 딕셔너리는 small_m_cluster 커널이 사용될 때 토큰당 CTA(Compute Thread Array) 수를 결정합니다. 예를 들어, SM100 아키텍처에서 K=6일 때, 이전에는 M=16까지만 클러스터링이 적용되었지만, 이제는 M=32까지 클러스터 2, M=64까지 클러스터 2가 적용됩니다. 이는 더 큰 M 값에서도 클러스터링을 통해 병렬성을 높일 수 있음을 시사합니다.

또한, _SMALL_M_CHUNK_BANDS 딕셔너리가 새로 추가되었습니다. 이는 small_m 커널의 온라인 소프트맥스 계산 시 사용되는 청크(chunk) 크기를 M 및 K 값에 따라 조절합니다. 이전에는 일부 K 값에 대해 고정된 청크 크기(예: nc4)를 사용했지만, 이제는 M 값에 따라 nc3 또는 nc4를 선택적으로 사용하여 메모리 접근 패턴을 최적화할 수 있습니다.

2. 라우팅 대체 로직 추가 (`cake_backend.py`)

RoutePlan 네임드 튜플에 fallback_from 필드가 추가되었습니다. 이는 특정 M, K 조합에 대해 최적화된 커널이 현재 체크아웃에 등록되어 있지 않을 경우, 가장 유사한 등록된 커널로 대체하고 그 정보를 기록하는 메커니즘입니다.

class RoutePlan(NamedTuple):
    # ... (기존 필드)
    fallback_from: Optional[str] = None

_persistent_fallback 함수는 이러한 대체 로직을 구현합니다. 만약 요청된 커널이 등록되어 있지 않으면, 동일한 프로그램 계열 내에서 가장 유사한 등록된 커널을 찾아 사용합니다. 이는 커널 커버리지를 넓히고 NotImplementedError 발생 가능성을 줄여줍니다.

3. README 업데이트 (`README.md`)

flashinfer/experimental/cake_kimi_k3_attn_res/README.md 파일도 업데이트되어, 변경된 라우팅 규칙과 커널 선택 기준을 반영했습니다. 특히 small_m_direct 및 small_m_cluster 커널에 대한 설명이 상세해졌습니다.

왜 이게 좋은가?

이 PR의 주요 개선점은 다음과 같습니다:

성능 향상: 중간 크기 토큰 수(M)에 대한 커널 선택 로직이 개선되면서, 해당 범위에서 더 효율적인 커널(예: small_m_direct, small_m_cluster)이 사용됩니다. 이는 GPU 활용률을 높여 전체적인 연산 속도를 향상시킵니다.
- 데이터 기반 성능: PR의 결과 표를 보면, Baseline ms 대비 Export ms (새로 최적화된 커널)의 비율이 대부분 1.0x에 가깝거나 더 좋습니다. 특히 Per-shape comparison 섹션의 Baseline / Export 비율은 새로운 커널이 기존 vLLM 네이티브 CUDA 커널보다 훨씬 빠르다는 것을 보여줍니다 (예: primary_m1_k0__sm_100a__pdl0에서 1.51x). Geomeans 섹션에서도 Source / Export 비율이 0.999x로 매우 우수하며, vLLM / export 비율은 최대 1.15x까지 성능 향상을 보여줍니다.
커널 커버리지 확장: fallback_from 메커니즘 도입으로 인해, 이전에 지원되지 않던 M, K 조합에서도 안정적으로 커널을 실행할 수 있게 되었습니다. 이는 라이브러리의 견고성을 높입니다.
세밀한 최적화: _SMALL_M_CHUNK_BANDS와 같이 온라인 소프트맥스 계산 시 청크 크기를 조절하는 기능은 메모리 접근 패턴을 최적화하여 캐시 효율성을 높이고 잠재적인 병목 현상을 줄입니다.
아키텍처 특화 튜닝: SM100과 SM103 아키텍처별로 다른 _SMALL_M_DIRECT_MAX_M 및 _SMALL_M_CLUSTER 값을 적용하여, 각 아키텍처의 특성에 맞는 최적의 성능을 이끌어냅니다.

일반적 교훈

이 PR은 다음과 같은 일반적인 최적화 교훈을 제공합니다:

프로파일링 기반 최적화: 특정 작업 부하(여기서는 중간 크기 토큰 수)에서 성능 병목을 식별하고 해당 부분을 집중적으로 최적화하는 것이 중요합니다. 단순히 모든 경우를 일반화하는 것보다 특정 시나리오에 대한 깊이 있는 분석이 큰 성능 향상을 가져올 수 있습니다.
동적 커널 선택 및 라우팅: 하드웨어 아키텍처, 입력 크기(M, K), 그리고 연산 특성(예: PDL 사용 여부)에 따라 최적의 커널을 동적으로 선택하는 라우팅 로직은 고성능 컴퓨팅 라이브러리에서 필수적입니다. 이는 다양한 사용 사례와 하드웨어 환경에 대한 성능을 극대화합니다.
점진적 개선과 커버리지: 모든 가능한 입력 조합에 대해 완벽한 커널을 미리 준비하는 것은 비효율적일 수 있습니다. 대신, 핵심 경로를 최적화하고, 지원되지 않는 경로는 가장 가까운 최적화된 경로로 대체하는 전략은 개발 효율성과 라이브러리 안정성을 동시에 확보하는 좋은 방법입니다.
코드와 문서의 일관성: 변경된 라우팅 로직과 커널 선택 기준은 README 파일에 명확하게 문서화되어야 합니다. 이는 라이브러리 사용자뿐만 아니라 개발자 간의 이해를 돕습니다.

결론

이번 PR은 FlashInfer의 Kimi-K3 어텐션 커널에서 중간 크기 토큰 수에 대한 라우팅 전략을 정교하게 개선함으로써 상당한 성능 향상을 이루었습니다. 이는 LLM 추론 성능 최적화에 있어 세밀한 커널 튜닝과 지능적인 라우팅 로직의 중요성을 다시 한번 강조합니다. 앞으로도 FlashInfer와 같은 라이브러리들이 지속적인 최적화를 통해 AI 모델의 효율성을 높여나갈 것으로 기대됩니다.

References

FlashInfer GitHub Repository
vLLM Project
torch.add documentation (Large token counts benchmark reference)
torch.clone documentation (Large token counts benchmark reference)

참고 자료

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

PR Analysis 의 다른글

이전글 [sglang] SGLang FLUX 3 Action 성능 최적화: CUDA Graphs 도입을 통한 Host Launch 오버헤드 제거
현재글 : [flashinfer] FlashInfer, Kimi-K3 어텐션의 성능 최적화: SM100/SM103 아키텍처를 위한 중간 크기 라우팅 개선
다음글 [onnxruntime] ONNX Runtime MLAS, AVX-512 NCHWc Depthwise Convolution 최적화: Sliding Window 커널 도입