vLLM Kimi K3 지원

vLLM Kimi K3 지원

vLLM은 2.8조 개의 파라미터를 가진 오픈 웨이트 멀티모달 Mixture-of-Experts (MoE) 모델인 Kimi K3에 대한 효율적인 day-0 지원을 발표했습니다. 이번 통합을 통해 Kimi K3의 독특한 하이브리드 아키텍처에 대한 고성능 서빙이 가능해졌으며, 16개의 NVIDIA GB300 NVL72 GPU에서 DSpark speculative decoding을 사용할 경우 사용자당 최대 370 tokens per second를 달성합니다.

Kimi K3 아키텍처 개요

Kimi K3는 토큰당 896개의 전문가 중 16개를 활성화하는 2.8T 파라미터 모델입니다. 이 모델은 다음과 같은 주요 혁신을 통해 최대 100만 토큰의 컨텍스트 윈도우를 지원하며 긴 컨텍스트 효율성을 위해 설계되었습니다:

  • Kimi Delta Attention (KDA): 선형 어텐션 레이어(고정된 크기의 순환 상태를 유지)와 전역적 회상을 보존하기 위한 주기적인 풀 어텐션 레이어를 교차시키는 하이브리드 스택입니다.
  • Attention Residuals (AttnRes): 표준 잔차 누적을 depth-wise attention으로 대체하고, 학습된 pseudo-queries를 사용하여 이전 블록의 잔차 상태에 가중치를 부여하는 메커니즘입니다.
  • Stable LatentMoE: 라우팅된 전문가 계산을 위해 활성화 값을 더 좁은 잠재 차원으로 투영하는 설계로, Quantile Balancing을 사용하여 라우터 점수 분위수에 따라 전문가를 할당합니다.
  • Native MXFP4: 대역폭 요구 사항을 줄이기 위해 MoE 경로에서 네이티브 4-bit 가중치를 사용합니다.

Kimi K3를 위한 vLLM 서빙 최적화

Kimi K3가 표준 Transformer 아키텍처에서 벗어난 점을 처리하기 위해, vLLM은 몇 가지 특화된 서빙 메커니즘을 구현했습니다:

하이브리드 KV-Cache 관리

vLLM은 풀 어텐션 레이어를 위한 paged KV 블록과 KDA 레이어를 위한 컴팩트한 순환 상태 블록을 모두 처리하기 위해 단일 하이브리드 KV-cache 매니저를 사용합니다. 이 하이브리드 설계에서 프리픽스 캐싱을 활성화하기 위해, vLLM은 물리적 KDA 상태 블록을 세밀한 프리픽스 매칭에서 분리하여 공유 프롬프트가 KDA 상태와 paged KV를 모두 재사용할 수 있도록 합니다.

Fused Kernels 및 메타데이터 최적화

  • Attention Residuals: vLLM은 fused Triton 및 CUDA 커널을 사용하여 depth-wise attention logits, softmax 및 hidden-state aggregation을 하나의 연산으로 계산하여 메모리 트래픽을 줄입니다.
  • KDA Decode: 특화된 CUDA 커널이 causal convolution, recurrent updates 및 RMSNorm을 하나의 실행으로 융합하여 중간 텐서와 반복적인 상태 트래픽을 방지합니다.
  • KDA Prefill: vLLM은 prefill 단계를 가속화하기 위해 FlashKDA 및 후속 커뮤니티 최적화(Flash-Flash-KDA)를 통합했습니다.
  • Metadata Builder: 전용 Kimi K3 KDA 메타데이터 빌더가 eager PyTorch 연산을 fused Triton 커널로 대체하여, 배치 크기 1에서 메타데이터 준비 지연 시간을 96%(870°s에서 34°s로) 단축했습니다.

저지연 컴퓨팅

vLLM은 저배치 크기 설정에 skinnyGEMM을 채택하여 공유 메모리 데이터 스테이징을 우회함으로써 엔드 투 엔드 지연 시간을 약 10% 줄였습니다. 또한, LatentMoE tail-fusion 최적화는 선형 투영 단계에서의 중복 계산을 줄여 7°8%의 엔드 투 엔드 속도 향상을 가져옵니다.

프로덕션 기능 및 성능

DSpark를 이용한 Speculative Decoding

vLLM은 블록 확산(block-diffusion) speculative decoding 알고리즘인 DSpark를 지원합니다. vLLM과 TorchSpec으로 학습된 DSpark speculator를 사용하여, vLLM은 단일 사용자 요청에 대해 3.14배의 속도 향상을 달성하여 처리량을 118 tok/s에서 370 tok/s로 증가시켰습니다. 수용률(Acceptance rates)은 저엔트로피 작업(코딩)의 경우 스텝당 4.73 tokens에서 고엔트로피 작업(창작 글쓰기)의 경우 스텝당 2.61 tokens 사이입니다.

Prefill/Decode 분리 (Disaggregation)

고처리량 환경을 위해 vLLM은 prefill/decode (PD) 분리를 지원합니다. NIXL 커넥터는 별도의 prefill 및 decode 복제본 간에 토큰 레벨의 MLA 캐시와 요청 레벨의 KDA 상태 모두의 전송을 관리합니다.

에이전틱 캐시 유지 정책

긴 컨텍스트 에이전틱 워크로드에서 KDA 상태의 메모리 점유를 관리하기 위해 vLLM은 두 가지 유지 정책을 구현합니다:

  • 간격 기반 유지(Interval-based retention): 고정된 간격(예: 매 32K 토큰마다) 및 프롬프트 경계에서 KDA 상태를 체크포인트합니다.
  • Marconi 스타일의 선택적 유지(Marconi-style selective retention): 프리픽스의 두 번째 히트 시에만 KDA 상태를 캐싱하는 수요 기반 접근 방식으로, 자주 공유되는 프리픽스만 캐시 용량을 소비하도록 보장합니다.

벤치마크 및 하드웨어 요구 사항

성능 지표

GB300 NVL72 GPU에서 vLLM은 다음과 같은 단일 사용자 디코드 처리량을 달성합니다 (배치 크기 1):

구성 표준 디코드 DSpark Speculative Decode
TP8 111 tok/s 331 tok/s
TP16 118 tok/s 370 tok/s

정확도

vLLM 상의 Kimi K3는 최대 추론 노력(maximum reasoning effort) 시 GSM8K에서 0.976, GPQA-Diamond에서 0.939, OCRBench에서 0.889, MMMU Pro Vision에서 0.818을 기록했습니다.

하드웨어 요구 사항

Kimi K3를 서빙하려면 최소 8× B300 (또는 GB300 NVL72) GPU 또는 16× B200 GPU가 필요합니다. NVIDIA (Hopper 및 Blackwell)와 AMD (MI355X) 하드웨어 모두 지원됩니다.

배포 가이드 요약

Kimi K3를 실행하려면 vLLM은 다음 구성을 권장합니다:

  • Prefix Caching: --enable-prefix-caching을 통해 명시적으로 활성화해야 합니다.
  • MoE Backend: 분리형/전문가 병렬(DEP) 환경에는 deep_gemm_mega_moe를, TP > 1인 경우 flashinfer_trtllm을 사용하십시오.
  • All-to-all Backend: NVLink에는 flashinfer_nvlink_one_sided를, RDMA에는 deepep_v2를 사용하십시오.
  • Rust Frontend: VLLM_USE_RUST_FRONTEND=1을 통해 활성화하십시오.

Sources