[flashinfer] [FlashInfer] Device-side Gating을 통한 MoE GEMM2 경로 최적화: Kimi-K3 모델 성능 개선 사례

PR 링크: flashinfer-ai/flashinfer#6411 상태: Merged | 변경: +9762 / -330

들어가며

대규모 언어 모델(LLM)의 추론 성능을 높이기 위해 Mixture of Experts(MoE) 구조가 널리 사용되고 있습니다. 하지만 MoE는 각 토큰이 서로 다른 전문가(Expert)에게 할당되는 특성상, 연산 경로가 동적으로 결정되어야 한다는 기술적 난제가 있습니다. 특히 Kimi-K3와 같은 최신 모델에서 Long-prefill(T ≥ 16384) 상황을 처리할 때, 어떤 GEMM(General Matrix Multiply) 경로를 선택하느냐는 전체 지연 시간(Latency)에 결정적인 영향을 미칩니다.

기존 FlashInfer의 Cake 백엔드에서는 이러한 경로 선택(Dense vs Split)이 Host(CPU) 측에서 결정되었습니다. 이는 호스트와 디바이스 간의 동기화 오버헤드를 발생시키거나, CUDA Graph 캡처 시 유연성을 떨어뜨리는 원인이 되었습니다. 이번 PR(#4254 관련)은 이 의사결정 로직을 Device(GPU) 내부의 전용 게이트 커널로 옮겨 성능을 극대화한 사례입니다.

코드 분석: Device-side Gating의 도입

이번 변경의 핵심은 gemm2_gate라는 새로운 커널의 도입과 이를 지원하기 위한 매니페스트(Manifest) 및 플랜(Plan)의 업데이트입니다.

1. 새로운 커널 모듈 정의

cake_mxfp4_situ_moe_cute_manifest.json 파일을 보면, 기존에는 없던 gemm2_gate와 새로운 Dense Finalize 형태가 추가된 것을 확인할 수 있습니다.

[Before/After: Manifest의 커널 정의 추가]

// After (추가된 부분)
"gemm2_gate": {
  "backends": ["cuda_cpp", "cutedsl"],
  "gate": "e2e_registry",
  "kind": "gemm2_gate",
  "pdl": true,
  "plan_selected": true,
  "registry": "kimi_k3_mxfp4_situ_gemm2_gate",
  "template": "kimi_k3_mxfp4_situ_gemm2_gate",
  "token": "gemm2_gate",
  "use_pdl": true
},
"gemm2_dense_finalize_n192_c12_rg": {
  "kind": "gemm2_dense",
  "registry": "kimi_k3_mxfp4_situ_gemm2_dense_n192_c12_rg",
  "row_group": true,
  "use_pdl": true,
  // ... 생략 ...
}

여기서 gemm2_gate는 단 하나의 워프(One-warp)만 사용하는 매우 가벼운 커널입니다. 이 커널은 moe_sort 이후에 실행되어, 정렬된 테이블을 기반으로 Split 경로를 탈지, Dense 경로를 탈지 디바이스 상에서 즉시 결정합니다.

2. CUDA Graph를 위한 Fixed Launch Chain

기존에는 호스트가 경로를 결정했기 때문에, 경로가 바뀌면 실행해야 할 커널 리스트가 달라져 CUDA Graph를 다시 캡처하거나 여러 개의 그래프를 관리해야 했습니다.

이번 PR에서는 이를 해결하기 위해 12개의 고정된 런치 체인(Fixed 12-launch chain) 방식을 도입했습니다.

• gemm2_gate가 경로를 선택하면, 선택되지 않은 경로의 커널들은 실행되더라도 내부에서 처리할 아이템 수를 0으로 인지하고 즉시 종료(Exit with zero items)합니다.
• 결과적으로 CUDA Graph는 단 하나의 고정된 런치 리스트만 유지하면 되므로, 런타임 오버헤드가 획기적으로 줄어듭니다.

왜 이게 좋은 최적화인가?

1. 성능 수치상의 이점

B300 GPU에서 측정된 결과에 따르면, 32,768 토큰의 롱 프리필 상황에서 특정 행(Row)의 커널 실행 시간이 8.4%(cake)에서 9.5%(cake_cute)까지 감소했습니다.

• Speedup vs trtllm-gen: 기존 1.096배에서 1.214배로 향상되었습니다.
• Geomean 성능: 전체적인 기하 평균 성능 역시 1.2506에서 1.2536으로 개선되었습니다.

2. 호스트-디바이스 동기화 제거

데이터의 특성에 따라 실행 경로를 바꾸는 '조건부 실행'을 호스트가 담당하면, GPU의 연산 결과를 CPU가 읽어올 때까지 기다려야 하는 동기화 지점이 생깁니다. 이 로직을 GPU 내부 커널(gemm2_gate)로 옮김으로써 GPU는 멈춤 없이 연산을 이어갈 수 있습니다.

3. 수치적 안정성 유지

경로가 Dense에서 Split으로 변경되면 부동소수점 연산의 순서가 달라져 결과값이 미세하게 변할 수 있습니다. 하지만 이번 PR의 테스트 결과에 따르면, FP64 Oracle 대비 오차가 BF16 클래스 내에서 엄격하게 관리되고 있음을 확인했습니다(Max abs error 0.0142 이내).

시니어 엔지니어의 관점: 교훈

이 최적화 사례는 고성능 커널 설계 시 고려해야 할 중요한 원칙을 보여줍니다.

1. Decision on Device: 런타임 데이터에 의존하는 의사결정 로직은 가능한 한 디바이스 내부로 밀어 넣어야 합니다. CPU-GPU 통신은 언제나 병목의 주범입니다.
2. Static Graph, Dynamic Workload: CUDA Graph와 같이 정적인 실행 구조를 유지하면서도 동적인 워크로드를 처리하기 위해 'Zero-item exit' 전략을 사용하는 것은 매우 영리한 접근입니다.
3. Specialized Gating: 모든 스레드가 결정에 참여할 필요 없이, 단 하나의 워프(one-warp gate)만으로 전체 그리드의 운명을 결정하게 함으로써 게이팅 자체의 오버헤드도 최소화했습니다.

FlashInfer는 이번 개선을 통해 Kimi-K3와 같은 초거대 컨텍스트 모델의 추론 효율성을 한 단계 더 끌어올렸습니다. 롱 컨텍스트 처리가 중요해지는 트렌드에서 이러한 디바이스 중심의 세밀한 최적화는 필수적입니다.

참고 자료

• https://docs.nvidia.com/cuda/cuda-programming-guide/index.html
• https://github.com/NVIDIA/cutlass
• https://github.com/flashinfer-ai/flashinfer

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