[flashinfer] FlashInfer, B200 GPU에서 Kimi-K3 MoE 라우터 성능 최적화
PR 링크: flashinfer-ai/flashinfer#8016 상태: Merged | 변경: +2224 / -6
들어가며
최근 대규모 언어 모델(LLM) 분야에서는 Mixture-of-Experts (MoE) 아키텍처가 큰 주목을 받고 있습니다. MoE는 모델의 파라미터 수를 크게 늘리면서도 추론 시에는 일부 전문가(expert)만 활성화하여 연산량을 효율적으로 관리할 수 있다는 장점이 있습니다. FlashInfer는 이러한 MoE 모델의 효율적인 추론을 지원하기 위한 라이브러리로, 지속적으로 성능 개선을 위한 노력을 기울이고 있습니다. 이번 글에서는 FlashInfer의 cake_kimi_k3_fused_router에서 발생한 성능 병목 현상을 해결하고, 특히 NVIDIA B200 GPU 환경에서 1-토큰 라우팅 성능을 향상시킨 PR(#1234)에 대해 자세히 분석해 보겠습니다.
이 PR은 flashinfer.fused_moe.kimi_k3_fused_router의 특정 시나리오, 즉 1개의 토큰만 처리하는 경우 NVIDIA B200 GPU (sm_100a)에서의 성능을 개선하는 것을 목표로 합니다. 기존에는 1-토큰 라우팅 시 약간의 오버헤드가 존재하여 최적의 성능을 내지 못했습니다. 이 PR은 커널 로직을 수정하여 이 오버헤드를 제거하고, 결과적으로 1-토큰 라우팅 속도를 약 1.04배 향상시켰습니다.
코드 분석
이번 PR의 핵심 변경 사항은 flashinfer/experimental/kimi_k3_fused_router/cake_jit.py 파일에서 관리되는 커널 설정에 있습니다. 특히, sm_100a 아키텍처에 대한 1-토큰 라우팅 커널의 설정을 변경하고, 새로운 커널을 추가하여 성능을 개선했습니다.
flashinfer/experimental/kimi_k3_fused_router/cake_jit.py 변경 사항
이 파일은 FlashInfer의 JIT 컴파일된 커널들을 관리하는 설정 파일입니다. 각 커널은 고유한 식별자(SHA256 해시 기반)를 가지며, 특정 GPU 아키텍처(arches) 및 설정(num_tokens, block_m 등)에 매핑됩니다.
1. sm_103a 아키텍처에서 sm_100a 아키텍처 분리 및 최적화:
기존에는 cake_kimi_k3_fused_router_27a24128e0b7468b7878 커널이 sm_100a와 sm_103a 모두를 지원했습니다. 하지만 이번 PR에서는 1-토큰 라우팅에 대한 최적화가 sm_100a에만 적용되었기 때문에, 해당 커널 설정을 sm_103a만 지원하도록 변경했습니다.
Before:
@@ -65,7 +65,7 @@
"cake_kimi_k3_fused_router_27a24128e0b7468b7878": {
"arm": "LC",
"num_tokens": 1,
- "arches": ["sm_100a", "sm_103a"],
+ "arches": ["sm_103a"],
"sources": [
"cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_27a24128e0b7468b7878_kernel.cu",
"cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_27a24128e0b7468b7878_binding.cu",
또한, sm_100a에 대한 closure_sha256 항목이 제거되었습니다. 이는 해당 커널의 sm_100a 버전이 더 이상 사용되지 않음을 의미합니다.
Before:
@@ -95,7 +95,6 @@
"dynamic_smem_bytes": 4352,
},
"closure_sha256": {
- "sm_100a": "9a6a45bd1325755c8aa0900295e52dcec6c758239f35cbca2e6367923f76209d",
"sm_103a": "e102766796fceebc25512769c7871764778f2d6622c4f34063a2f7bc920a3863",
},
"kernel_symbol": "kernel_cake_kimi_k3_fused_router_27a24128e0b7468b7878",
2. 새로운 1-토큰 최적화 커널 추가 (cake_kimi_k3_fused_router_5734ff3e8171d72c1539):
sm_100a 아키텍처에서 1-토큰 라우팅 성능을 개선하기 위해 새로운 커널 설정이 추가되었습니다. 이 커널은 sm_100a 아키텍처만을 대상으로 하며, num_tokens가 None으로 설정되어 있어 다양한 토큰 수를 처리할 수 있음을 시사합니다. (하지만 PR 설명에 따르면 실제 성능 개선은 1-토큰에서 발생합니다.)
After:
@@ -248,6 +247,43 @@
},
"kernel_symbol": "kernel_cake_kimi_k3_fused_router_55d217eef88b4aff9f1c",
},
+ "cake_kimi_k3_fused_router_5734ff3e8171d72c1539": {
+ "arm": "LC1_NC",
+ "num_tokens": None,
+ "arches": ["sm_100a"],
+ "sources": [
+ "cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_5734ff3e8171d72c1539_kernel.cu",
+ "cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_5734ff3e8171d72c1539_binding.cu",
+ ],
+ "compile_flags": [],
+ "ffi_entry": "run",
+ "arg_plan": [
+ ["buffer", "logits"],
+ ["buffer", "bias"],
+ ["buffer", "topk_weights"],
+ ["buffer", "topk_ids"],
+ ["buffer", "sorted_token_ids"],
+ ["buffer", "expert_ids"],
+ ["buffer", "num_tokens_post_padded"],
+ ["buffer", "expert_counts"],
+ ["buffer", "expert_offsets"],
+ ["buffer", "expert_scatter_offsets"],
+ ["parameter", "M"],
+ ["grid", "grid_x"],
+ ["grid", "grid_y"],
+ ["grid", "grid_z"],
+ ],
+ "launch": {
+ "block": [224, 1, 1],
+ "cluster": [1, 1, 1],
+ "cooperative": False,
+ "dynamic_smem_bytes": 4352,
+ },
+ "closure_sha256": {
+ "sm_100a": "9e59e5fb56dce765b95f8098a3e5ae176ed69eea619e25e125389e2bd783941d",
+ },
+ "kernel_symbol": "kernel_cake_kimi_k3_fused_router_5734ff3e8171d72c1539",
+ },
"cake_kimi_k3_fused_router_76a711ed994e53632512": {
"arm": "LPC",
"num_tokens": None,
3. 다른 커널 설정 변경:
마찬가지로 cake_kimi_k3_fused_router_0b365c12c38f335f62bc 커널도 sm_100a 지원을 제거하고 sm_103a만 지원하도록 변경되었습니다. 이는 1-토큰 라우팅 성능 개선이 sm_100a 특정 커널에서 이루어졌음을 다시 한번 보여줍니다.
Before:
@@ -784,7 +820,7 @@
"cake_kimi_k3_fused_router_0b365c12c38f335f62bc": {
"arm": "LC1_NC",
"num_tokens": None,
- "arches": ["sm_100a", "sm_103a"],
+ "arches": ["sm_103a"],
"sources": [
"cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_0b365c12c38f335f62bc_kernel.cu",
"cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_0b365c12c38f335f62bc_binding.cu",
그리고 sm_100a에 대한 closure_sha256 항목이 제거되었습니다.
Before:
@@ -814,11 +850,47 @@
"dynamic_smem_bytes": 4352,
},
"closure_sha256": {
- "sm_100a": "24a04e26eb4086bad1a3d8da3ad2339d83e0e574010f09bb4f0a164475e86c63",
"sm_103a": "c78989c7c995fe5e02cc5c136a86d385f8ab6f2d27d11f608b0ad5caee4eb3a5",
},
"kernel_symbol": "kernel_cake_kimi_k3_fused_router_0b365c12c38f335f62bc",
},
+ "cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf": {
+ "arm": "LC",
+ "num_tokens": 1,
+ "arches": ["sm_100a"],
+ "sources": [
+ "cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf_kernel.cu",
+ "cake_kimi_k3_fused_router/cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf_binding.cu",
+ ],
+ "compile_flags": [],
+ "ffi_entry": "run",
+ "arg_plan": [
+ ["buffer", "logits"],
+ ["buffer", "bias"],
+ ["buffer", "topk_weights"],
+ ["buffer", "topk_ids"],
+ ["buffer", "sorted_token_ids"],
+ ["buffer", "expert_ids"],
+ ["buffer", "num_tokens_post_padded"],
+ ["buffer", "expert_counts"],
+ ["buffer", "expert_offsets"],
+ ["buffer", "expert_scatter_offsets"],
+ ["parameter", "M"],
+ ["grid", "grid_x"],
+ ["grid", "grid_y"],
+ ["grid", "grid_z"],
+ ],
+ "launch": {
+ "block": [224, 1, 1],
+ "cluster": [1, 1, 1],
+ "cooperative": False,
+ "dynamic_smem_bytes": 4352,
+ },
+ "closure_sha256": {
+ "sm_100a": "ca6a0aaf169f08f2d400d52c2cf77e02cd7bdf8ea7d4ebdab624c5cb0d959953",
+ },
+ "kernel_symbol": "kernel_cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf",
+ },
"cake_kimi_k3_fused_router_006a7ad397982e4faa88": {
"arm": "LC",
"num_tokens": 16,
새로운 커널 cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf가 sm_100a 아키텍처에 대해 추가되었습니다. 이 커널은 num_tokens: 1로 명시되어 있어, 1-토큰 라우팅에 특화된 최적화가 적용되었음을 알 수 있습니다.
4. GPU 아키텍처 매핑 업데이트:
마지막으로, sm_100a 아키텍처에 대한 커널 매핑이 업데이트되었습니다. 기존에 sm_100a를 가리키던 커널들이 sm_103a 전용 커널로 변경되거나, 새로 추가된 최적화 커널(cake_kimi_k3_fused_router_5734ff3e8171d72c1539 및 cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf)로 대체되었습니다.
Before:
@@ -1094,9 +1166,9 @@
"sm_100a": {
"GW": "cake_kimi_k3_fused_router_a140805ed9d697427fbe",
"L": "cake_kimi_k3_fused_router_544948d65e43d405cb26",
- "LC1_NC": "cake_kimi_k3_fused_router_0b365c12c38f335f62bc",
+ "LC1_NC": "cake_kimi_k3_fused_router_5734ff3e8171d72c1539",
"LC2_CW": "cake_kimi_k3_fused_router_55d217eef88b4aff9f1c",
- "LC:1": "cake_kimi_k3_fused_router_27a24128e0b7468b7878",
+ "LC:1": "cake_kimi_k3_fused_router_0b2a25d8ce9c390209cf",
"LC:16": "cake_kimi_k3_fused_router_006a7ad397982e4faa88",
"LC:2": "cake_kimi_k3_fused_router_aea2ff5129b7b6374e5a",
"LC:4": "cake_kimi_k3_fused_router_9701331ccf9b00bc5af8",
이러한 변경은 sm_100a 아키텍처의 1-토큰 라우팅 시나리오에 대해 더 효율적인 커널을 사용하도록 시스템을 구성하는 역할을 합니다.
커널 내부 로직 변경 (추정)
cake_jit.py 파일 자체는 커널 설정을 관리할 뿐, 실제 CUDA 커널 코드를 포함하고 있지는 않습니다. 하지만 PR 설명과 diff에서 나타난 변경 사항을 통해 커널 내부 로직의 변화를 추론할 수 있습니다.
PR 설명에 따르면, 변경된 1-토큰 셀의 커널은 다음과 같은 최적화를 수행합니다:
- "rank their 16 winners from a certified window": 16개의 최고 전문가를 선정하는 과정을 최적화했습니다. 아마도 불필요한 비교나 정렬 과정을 줄였을 것입니다.
- "zero the expert counts and scatter offsets at kernel entry": 커널 시작 시점에
expert_counts와expert_scatter_offsets를 0으로 초기화합니다. 이는 이전 커널에서 사용되던 값들이 다음 커널 실행에 영향을 미치는 것을 방지하거나, 혹은 특정 조건에서 불필요한 계산을 줄이기 위한 조치일 수 있습니다.
이러한 로직 변경은 1-토큰 처리 시 발생하는 오버헤드를 줄여 직접적인 성능 향상을 가져온 것으로 보입니다.
왜 이게 좋은가
이번 PR은 다음과 같은 이유로 좋은 최적화라고 할 수 있습니다.
-
명확한 성능 향상: NVIDIA B200 (sm_100a) GPU에서 1-토큰 라우팅 시
block_m8 및 16 설정에서 각각 2.768 us -> 2.655 us (1.04x), 2.704 us -> 2.591 us (1.04x)의 속도 향상을 달성했습니다. 이는 특히 배치 크기가 작거나 시퀀스 길이가 짧은 경우 MoE 라우터의 성능을 직접적으로 개선하여 전체 추론 속도 향상에 기여합니다.tokens block_m FlashInfer main (us) this PR (us) speedup kernel 1 8 2.768 2.655 1.04x new 1 16 2.704 2.591 1.04x new -
아키텍처 특화 최적화: 모든 GPU 아키텍처에 대해 동일한 코드를 적용하는 대신, 특정 아키텍처(B200, sm_100a)의 특성을 고려하여 최적화된 커널을 별도로 제공했습니다. 이는 하드웨어의 성능을 최대한 활용하는 좋은 전략입니다.
sm_103a에서는 변경이 없다는 점은 이러한 특화 전략이 효과적이었음을 보여줍니다. -
코드의 명확성 및 유지보수성 향상:
cake_jit.py파일에서sm_100a와sm_103a를 분리하고, 1-토큰 최적화 커널을 명시적으로 추가함으로써 코드의 의도를 더 명확하게 만들었습니다. 이는 향후 코드 유지보수 및 다른 아키텍처에 대한 최적화를 용이하게 합니다. -
일반적인 교훈:
- 프로파일링의 중요성: 성능 병목은 특정 조건(여기서는 1-토큰 라우팅)에서만 나타날 수 있습니다. 따라서 다양한 시나리오에 대한 철저한 프로파일링이 필수적입니다.
- 하드웨어 특화 최적화: 최신 GPU 아키텍처는 고유한 특성을 가지므로, 이를 이해하고 활용하는 것이 성능 극대화의 핵심입니다. 때로는 특정 아키텍처에 대한 별도의 커널 구현이 필요할 수 있습니다.
- 작은 개선의 누적: 1.04배의 속도 향상은 작아 보일 수 있지만, LLM과 같이 수많은 연산이 반복되는 환경에서는 이러한 작은 개선들이 모여 상당한 전체 성능 향상을 이끌어낼 수 있습니다.
리뷰 피드백 반영
리뷰어 yyihuang님은 테스트 실행을 요청했습니다 (@flashinfer-bot run tests/experimental/test_cake_kimi_k3_fused_router.py). 이는 코드 변경이 기존 기능에 영향을 미치지 않고 의도한 대로 작동하는지 확인하는 중요한 절차입니다. PR 설명에 따르면 "Outputs are bitwise identical to FlashInfer main."라고 명시되어 있어, 기능적 변경 없이 성능만 개선되었음을 강조하고 있습니다. 이는 테스트 커버리지가 잘 갖춰져 있음을 시사하며, 리뷰어의 요청은 이러한 안정성을 재확인하는 과정으로 볼 수 있습니다.
결론
이번 PR은 NVIDIA B200 GPU 환경에서 FlashInfer의 Kimi-K3 MoE 라우터의 1-토큰 라우팅 성능을 성공적으로 개선했습니다. 커널 설정을 최적화하고 특정 하드웨어에 맞는 새로운 커널을 도입함으로써 약 4%의 속도 향상을 달성했으며, 이는 MoE 기반 LLM의 추론 효율성을 높이는 데 기여합니다. 이러한 최적화는 하드웨어 특성을 고려한 세심한 접근 방식과 철저한 성능 측정을 통해 이루어졌으며, 기술 블로그 독자들에게도 유용한 교훈을 제공합니다.
참고 자료
⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.
관련 포스트
- [flashinfer] FlashInfer, MoE 모델의 성능을 극적으로 향상시키는 융합 커널과 최적화된 스케줄러 도입
- [flashinfer] FlashInfer, CuTe DSL W4A16 MoE 최적화 및 비트 단위 정확도 테스트 추가
- [flashinfer] FlashInfer, NVFP4 SVDQuant GEMM 커널 최적화: 아키텍처별 디스패치 도입
- [flashinfer] FlashInfer, Muse-Spark-Open을 위한 융합 정규화 및 TP/EP MoE 최적화로 성능 극대화
- [flashinfer] FlashInfer, cuDNN Frost MoE 디코딩 최적화: 성능 향상과 코드 개선 분석
PR Analysis 의 다른글
- 이전글 [flashinfer] NVIDIA R200(SM107)을 위한 SwiGLU + MXFP8 커널 최적화: 소규모 워크로드 성능 향상 기법
- 현재글 : [flashinfer] FlashInfer, B200 GPU에서 Kimi-K3 MoE 라우터 성능 최적화
- 다음글 [flashinfer] FlashInfer의 B200 최적화: Cake 백엔드 라우팅 개선을 통한 성능 극대화
댓글