vLLM Kimi K3.1 Pro 데이-0 지원: 희소 어텐션, DCP, 제로-미스매치 RL
TL;DR
vLLM은 Moonshot AI의 2.8T 파라미터(104B 활성) 오픈 가중치 멀티모달 MoE인 Kimi K3.1 Pro에 대해 NVIDIA GPU, AMD GPU, Google TPU, Alibaba T-HEAD Zhenwu PPU에서 데이-0 지원을 제공했습니다. 이번 릴리스가 중요한 이유는 K3.1 Pro가 기존 하이브리드 KDA + MLA 모델의 모든 24개 풀 어텐션 레이어에 희소 어텐션(DSA)을 추가했기 때문이며, vLLM은 이러한 하이브리드 희소 아키텍처를 효율적으로 서빙하기 위해 KV 캐시, 프리필 샤딩, 디코드 병렬 처리, 저지연 디코드 경로를 재설계해야 했습니다.
Kimi K3.1 Pro가 변경하는 것
Kimi K3.1 Pro는 Kimi K3의 아키텍처를 유지하면서 모든 풀 어텐션 레이어에 희소 어텐션 인덱서를 추가합니다. 모델은 93개 레이어 하이브리드(69개 KDA 레이어, 24개 풀 어텐션 레이어), 896개 MoE 전문가(16개 활성), MXFP4 MoE 가중치, 1M 토큰 컨텍스트 창을 유지합니다.
| Kimi K3 | Kimi K3.1 Pro | |
|---|---|---|
| 레이어 | 93: 69 KDA + 24 풀 어텐션 | 93: 69 KDA + 24 풀 어텐션 |
| 풀 어텐션 레이어 | dense MLA | MLA + DSA 인덱서 (top-2048), NoPE |
| 인덱서 | 없음 | 32 헤드 × 128 차원, FP8 키 |
| MoE | 896 전문가, 16 활성, MXFP4 | 변경 없음 |
| 컨텍스트 창 | 1M | 1M |
| 총 파라미터 | 2.8T | 2.8T |
핵심 기술 변경은 어텐션입니다. 각 희소 레이어는 MLA 어텐션 옆에 작은 인덱서를 유지합니다. 인덱서는 토큰당 하나의 FP8 키(스케일 포함 132바이트)를 저장하고, 32개의 경량 헤드를 사용하여 각 쿼리에 대해 모든 이전 토큰을 점수화합니다. 그런 다음 상위 2,048개 토큰을 유지하고, MLA 어텐션은 정확히 그 2,048개 토큰에 대해서만 실행됩니다. vLLM은 DSA가 DeepSeek-V3.2에서 도입되었으며 자세한 인덱서 설명은 이전 DeepSeek-V3.2 게시물을 참조한다고 밝혔습니다.
희소 어텐션이 긴 컨텍스트 디코드 대역폭을 줄임
DSA의 이점은 메모리 대역폭에 의해 제한되는 긴 컨텍스트 디코드에서 가장 큽니다. Dense MLA는 레이어당 컨텍스트 토큰당 576바이트의 FP8 잠재 표현을 읽습니다. DSA는 여전히 전체 컨텍스트를 스캔하지만, 스캔 중에는 132바이트 인덱스 키만 읽고, 그런 다음 2,048개 토큰에 대해서만 전체 잠재 표현을 읽습니다. 긴 컨텍스트에서 이는 레이어당 dense MLA가 읽는 바이트의 약 23%입니다.
이것이 K3.1 Pro의 핵심 서빙 주장입니다: 모델은 1M 토큰 컨텍스트를 유지하면서 긴 컨텍스트에서 디코드 단계의 레이어당 메모리 트래픽을 dense MLA 트래픽의 약 1/4로 줄입니다.
세 가지 상태를 위한 패킹된 KV 캐시
vLLM은 각 캐시 블록의 MLA 잠재 표현, 인덱서 키, KDA 상태를 함께 패킹하여 모든 레이어에 동일한 크기의 페이지를 제공하지 않습니다. 각 DSA 레이어는 서로 다른 너비의 토큰당 두 개의 캐시를 저장합니다: MLA 잠재 표현은 FP8에서 576바이트 또는 NVFP4에서 352바이트, 인덱스 키는 FP8에서 132바이트 또는 MXFP4에서 68바이트입니다. 두 너비 모두 서로의 배수가 아니며, 69개 KDA 레이어는 각각 요청당 재귀 상태를 추가하는데, 이는 한 레이어의 토큰 페이지보다 훨씬 큽니다. 기존의 레이어별 동일 크기 캐시는 작은 페이지를 가장 큰 페이지에 맞춰 패딩해야 합니다.
vLLM은 대신 캐시의 한 블록을 해당 블록의 MLA 페이지와 인덱스 키 페이지(모든 24개 DSA 레이어)와 블록을 공유하는 KDA 상태를 포함하는 단일 연속 행으로 만듭니다. 행 내부의 페이지는 서로 다른 크기를 가질 수 있습니다; 유일한 제약은 행이 KDA 상태를 보유한다는 것입니다. 블록 크기는 2의 거듭제곱이므로 서로 다른 TP 크기로 실행되는 프리필 및 디코드 인스턴스의 블록이 균등하게 나뉘고 KV 전송을 위해 짝지을 수 있습니다.
어텐션 커널은 여전히 예상하는 작은 고정 크기 페이지를 봅니다. vLLM은 각 단계 전에 블록 테이블과 캐시 슬롯을 커널 페이지 주소로 변환하므로 FlashInfer 및 FlashMLA의 희소 MLA 커널과 DeepGEMM의 인덱서 커널은 변경 없이 패킹된 캐시를 읽습니다.
이 레이아웃은 세 가지를 가능하게 합니다:
- 모든 희소 MLA 커널 및 KV 형식. Blackwell에서 기본값은 FP8 잠재 표현을 사용하는 FlashInfer의 희소 MLA입니다. FlashInfer의 PrimsTS 커널은 NVFP4 잠재 표현도 제공합니다.
- 분리형 서빙. 각 블록은 하나의 연속 메모리 범위이므로 프리필/디코드 전송은 수십 개의 작은 레이어별 읽기 대신 하나의 큰 RDMA 읽기로 이동합니다.
- 공유 인프라. 동일한 패킹 플래너는 혼합 페이지 크기 문제가 있는 캐시를 가진 GLM-5.3-Flash 및 DeepSeek-V4에도 제공됩니다.
토큰당 KV 캐시 풋프린트
텐서 병렬 처리에서 MLA 잠재 표현과 인덱스 키는 모든 GPU에 복제됩니다. KDA 상태는 K3에서와 같이 요청당입니다.
| GPU당 토큰당 KV 캐시 (24개 풀 어텐션 레이어) | 바이트 |
|---|---|
| Kimi K3, FP8 MLA 잠재 표현 | 13,824 |
| Kimi K3.1 Pro, FP8 잠재 표현 + FP8 인덱스 키 | 16,992 |
| Kimi K3.1 Pro, NVFP4 잠재 표현 + MXFP4 인덱스 키 | 10,080 |
잠재 표현과 인덱스 키는 독립적으로 양자화할 수 있습니다. NVFP4 잠재 표현만으로도 K3.1 Pro를 FP8의 K3보다 낮출 수 있습니다. 옵트인 MXFP4 인덱스 키는 추가로 11% 더 많은 토큰을 수용하고 긴 컨텍스트에서 인덱서를 10-20% 더 빠르게 만듭니다. 둘 다 사용하면 8개의 GB300 GPU가 FP8보다 1.7배 많은 토큰을 보유할 수 있으며, DCP8은 이를 7배로 곱합니다: 약 36개의 전체 1M 토큰 컨텍스트를 동시에, FP8의 TP8의 3개와 비교됩니다.
토큰 샤딩 프리필이 TTFT를 개선
큰 프리필 단계에서 vLLM은 헤드가 아닌 토큰으로 어텐션을 샤딩합니다. TP8에서 각 GPU는 K3.1 Pro의 96개 어텐션 헤드 중 12개를 소유하며, 이는 희소 MLA 프리필에 적합하지 않습니다: MLA의 잠재 KV는 모든 헤드가 공유하므로 각 GPU는 12개 헤드만 서빙하기 위해 전체 캐시를 읽고, 희소 커널은 이 12개 헤드를 최소 타일로 패딩합니다. 인덱서는 헤드별로 분할되지 않으므로 모든 GPU가 모든 토큰을 점수화하고 2차 비용이 8번 반복됩니다.
vLLM은 대신 all-to-all을 사용하여 각 GPU의 쿼리를 "모든 토큰, 12개 헤드"에서 "토큰의 1/8, 모든 96개 헤드"로 이동합니다. 각 GPU는 자체 토큰에 대해 인덱서와 희소 MLA를 실행하고, 두 번째 all-to-all이 출력을 헤드 레이아웃으로 반환합니다. KV 캐시는 이미 TP에서 복제되어 있으므로 다른 것은 변경되지 않습니다. MoE, 시퀀스 병렬 처리, 캐시 레이아웃은 그대로 유지됩니다. 인과 어텐션을 균형 있게 맞추기 위해 각 요청은 16개 조각으로 나뉘고 GPU r은 조각 r과 15 - r을 가져갑니다. 이는 프리필 컨텍스트 병렬 처리에서 사용되는 균형 분할과 동일합니다. 두 all-to-all 모두 피어의 대칭 메모리에 직접 씁니다.
8× GB300, TP8, 실제 가중치, 한 번에 하나의 요청에서:
| 프롬프트 토큰 | TTFT, 헤드 샤딩 | TTFT, 토큰 샤딩 | 속도 향상 |
|---|---|---|---|
| 27K | 0.82 s | 0.71 s | 1.16× |
| 108K | 3.72 s | 2.89 s | 1.29× |
| 340K | 15.89 s | 9.42 s | 1.69× |
에이전트 턴에서 흔한 긴 캐시된 프리픽스에 대한 짧은 추가도 이점을 얻습니다. 150K 토큰 캐시된 프리픽스에 2K 토큰 추가의 경우, 인덱서와 희소 MLA 커널은 8× GB300에서 FlashInfer 희소 MLA로 GPU당 약 15ms가 걸리며, 56ms 대신입니다. 비용은 교환 버퍼로 GPU당 약 0.8GiB의 KV 캐시입니다.
희소 인덱서를 사용한 디코드 컨텍스트 병렬 처리
텐서 병렬 처리는 MLA 잠재 표현을 모든 GPU에 복제하므로 GPU를 추가해도 KV 용량이 늘어나지 않습니다. 디코드 컨텍스트 병렬 처리(DCP)는 대신 토큰별로 캐시를 샤딩합니다: W 중 GPU r은 모든 W번째 토큰을 저장합니다. Dense MLA의 경우 각 GPU는 자체 샤드에 대해 어텐션을 수행하고 부분 출력은 로그-섬-지수(LSE)로 결합됩니다. DSA는 한 단계를 추가합니다: 각 GPU는 자체 토큰의 1/W만 순위를 매길 수 있으므로 전역 top-2048은 어텐션이 시작되기 전에 모든 GPU에서 병합되어야 합니다.
약 10배 적은 트래픽으로 정확한 top-k 병합
일반적인 병합은 모든 GPU의 전체 로컬 top-2048을 FP32 점수와 토큰 ID 쌍으로 다른 모든 GPU에 보냅니다. DCP8, 8× GB300, 각각 5개 쿼리 토큰을 가진 64개의 AgentX 유사 요청에서 어텐션 모듈 671µs 중 258µs가 소요됩니다.
vLLM은 훨씬 적게 보냅니다. 각 GPU는 BF16 점수에서 직접 최상의 1.5 × 2048 / W 후보만 선택하고 각각을 점수와 함께 32비트로 패킹하여 약 10배 적은 바이트를 보냅니다. 그런 다음 모든 GPU는 수집된 후보에서 동일한 전역 top-2048을 radix-select합니다. 병합은 정확하게 유지됩니다: GPU가 보낸 것보다 더 많은 top-2048을 보유하면 선택이 이를 감지하고 해당 GPU는 해당 행에 대한 전체 로컬 top-k를 다시 보냅니다.
정확한 단일 라운드 병합은 일반 병합보다 1.27-1.51배 빠릅니다. 배치 32 이상, 128K 컨텍스트에서 DCP8의 K3.1 Pro 희소 어텐션은 K3의 dense MLA보다 빠릅니다.
그런 다음 부분 어텐션 출력은 대칭 메모리에서 두 개의 CuTe DSL 커널로 LSE에 의해 교환 및 결합됩니다: 8× GB300에서 320개의 BF16 형태에 대해 기존 CUDA 구현보다 최대 3.7배 빠르며, 비트 단위로 동일한 출력을 생성합니다. vLLM은 이것이 DCP에서 모든 MLA 모델에 도움이 된다고 밝혔습니다.
병합 없는 DCP 프리필
일반 DCP 프리필은 모든 프롬프트 토큰에 대해 병합과 LSE 결합을 지불하고 모든 GPU에서 어텐션 작업을 반복합니다. DCP 그룹이 TP 그룹인 경우 vLLM은 대신 단계의 컨텍스트(MLA 잠재 표현과 인덱스 키 모두)를 모든 GPU에서 대칭 메모리 워크스페이스(FP8에서 GPU당 1M 토큰당 약 0.7GB)로 수집합니다. 그런 다음 위에서 설명한 토큰 샤딩 프리필을 실행합니다. 병합이나 결합이 남지 않으며 출력은 DCP 없는 토큰 샤딩 프리필과 비트 단위로 일치합니다.
| 하나의 DSA 레이어, DCP8, 8× GB300, AgentX 혼합 16K 토큰 단계 | DCP 프리필 | 수집된 컨텍스트 프리필 |
|---|---|---|
| FP8, FlashInfer trtllm-gen | 24.8 ms | 5.8 ms (4.3×) |
| FP8, FlashInfer PrimsTS | 24.8 ms | 5.7 ms (4.3×) |
| NVFP4, FlashInfer PrimsTS | 25.7 ms | 6.2 ms (4.1×) |
엔드 투 엔드로 단일 요청은 DCP8에서 TP8과 동일한 속도로 프리필됩니다: 27K, 108K, 340K 토큰에서 각각 1.10, 4.40, 14.61초, 8× GB300 FP8 KV에서 1.10, 4.39, 14.50초와 비교됩니다.
DCP 사용 시기
DCP는 단계당 지연 시간을 KV 용량과 교환합니다. DCP8은 GPU당 MLA 및 인덱스 키 캐시를 8배 줄이고 KDA 상태는 이미 헤드별로 분할되어 있습니다. 컨텍스트가 맞으면 일반 TP가 낮은 동시성에서 더 빠르게 디코드합니다. DCP는 KV가 제한일 때, 예를 들어 매우 긴 컨텍스트나 많은 장기 실행 에이전트 세션이 있을 때 유용합니다. 8× GB300에서 DCP8은 FP8 KV 21.8M 토큰을 보유하는 반면 TP8은 3.04M, NVFP4 잠재 표현으로는 32.1M 대 4.49M입니다. 동시성 1에서 TPOT는 32K 컨텍스트에서 TP8 9.35ms, DCP8 11.48ms, 128K에서 9.45 대 11.58ms입니다. 융합 디코드 경로는 DCP8을 5.96 및 6.04ms로 만듭니다.
TP8의 8개 GPU는 또한 전체 1M 토큰 요청을 서빙합니다: 8× GB300 FP8 KV에서 캐시는 3.04M 토큰, 즉 2.9개의 동시 1M 토큰 요청을 보유합니다.
DSA 레이어를 위한 융합 저지연 디코드
Kimi K3의 경우 vLLM은 작은 배치에 맞게 조정된 CuTe DSL 커널로 거의 모든 연산자를 대체하는 옵트인 디코드 경로를 제공했습니다. 통신을 컴퓨팅에 융합하고, MoE 스케줄을 재작업하고, 가중치를 L2로 프리페치합니다. K3.1 Pro의 경우 이 경로는 이제 DSA 레이어를 포함합니다. DSA 레이어는 다음과 같이 실행됩니다:
- 쿼리, KV 잠재 표현, 출력 게이트, 인덱서의 키 및 헤드 가중치를 위한 하나의 패킹된 입력 프로젝션;
- 융합 인덱서 프론트: 쿼리 정규화, 인덱서의 쿼리 프로젝션, 헤드 가중치가 포함된 FP8 또는 MXFP4 양자화, 새 토큰의 인덱스 키를 하나의 커널에서 처리;
- MLA 쿼리(업 프로젝션과 흡수된 키 프로젝션 포함, 새 토큰의 잠재 표현을 KV 캐시에 쓰는 것)는 인덱서와 겹쳐서 실행되며, 인덱서보다 먼저 실행되지 않음;
- 융합 인덱서: 키 균형 스코어러와 정확한 top-k 선택을 두 개의 커널로 실행하여 어텐션이 읽는 물리적 KV 행을 생성;
- 희소 어텐션: FlashInfer 희소 MLA(trtllm-gen 또는 PrimsTS, 배치 크기별 선택) 또는 FlashMLA 희소; DCP에서 배치 1-8일 때, 형태에 맞게 구축된 CuTe DSL 커널, 배치 1에서 레이어당 11.9µs 대 21.4µs;
- 값 프로젝션과 시그모이드 게이트를 하나의 커널로 처리한 다음, 융합된 all-reduce가 포함된 K3의 출력 프로젝션.
DCP에서 쿼리는 모든 GPU에 게시되고, 각 GPU의 top-k 후보는 교환 및 정확하게 병합되며, 부분 출력은 LSE로 결합됩니다. 이 모든 것은 커널 내부의 NVLink 메일박스를 통해 이루어지며 경로에 NCCL 집합이 없습니다.
융합 디코드 경로는 8× GB300, TP8, 1K 토큰 입력에서 배치 1에서 1.63배, 배치 64에서 여전히 1.18배 빠릅니다. DSpark 추측 디코딩과 7개의 드래프트 토큰을 사용하면 융합 경로는 동시성 1에서 토큰당 시간을 2.86ms에서 2.14ms로 낮추고 GSM8K 정확도는 변경되지 않습니다. DCP8에서 융합 경로는 동시성 1에서 TPOT를 32K 컨텍스트에서 11.48ms에서 5.96ms로, 128K에서 11.58ms에서 6.04ms로 줄입니다.
DSpark를 사용한 추측 디코딩
vLLM은 K3.1 Pro용 DSpark 스peculator를 훈련하고 미리보기 체크포인트인 Inferact/Kimi-K3.1-Pro-DSpark-preview를 출시했습니다. 온도 1.0, top-p 0.95의 현실적인 샘플링에서 미리보기는 7개의 드래프트 토큰으로 GSM8K에서 평균 수용 길이 4.597에 도달합니다. 8× GB300에서 TP8, 배치 크기 1에서 이는 3.0배 디코드 속도 향상입니다. 스peculator는 또한 K3.1 Pro의 최대 1M 토큰 전체 컨텍스트 창에서 정확도를 유지하며, 긴 컨텍스트 BEAM 벤치마크에서 측정되었습니다.
vLLM은 이 체크포인트를 미리보기 전용으로 설명하고 곧 출시할 더 정확하고 효율적인 스peculator를 훈련 중이라고 밝혔습니다.
NVIDIA GB300 및 Rubin에서의 에이전트 서빙
vLLM은 NVIDIA GB300 및 NVIDIA Rubin에서 SemiAnalysis AgentX 벤치마크로 데이-0 이미지를 벤치마킹했습니다.
NVIDIA GB300
프런티어의 지연 시간 민감 끝에서 Kimi K3.1 Pro는 사용자당 초당 375.9 토큰을 유지합니다. 프런티어를 따라 더 나아가면 GPU당 초당 19.2K 총 토큰을 서빙하면서 각 사용자를 초당 100 토큰 이상으로 유지합니다. 처리량 끝에서는 GPU당 초당 37.8K 토큰에 도달하며, 가장 인터랙티브한 지점의 약 26배이지만 각 사용자는 여전히 초당 약 54 토큰을 봅니다. GB300 스윕은 4개의 추측 토큰과 합성 수용 길이 3.36 토큰을 사용했습니다.
NVIDIA Rubin
NVIDIA Rubin에서 Kimi K3.1 Pro는 프런티어의 지연 시간 민감 끝에서 사용자당 초당 815.6 토큰을 유지합니다. 프런티어를 따라 더 나아가면 GPU당 초당 46.0K 총 토큰을 서빙하면서 각 사용자를 초당 100 토큰 이상으로 유지합니다. 처리량 끝에서는 GPU당 초당 68.7K 토큰에 도달하며, 가장 인터랙티브한 지점의 약 51배이지만 각 사용자는 여전히 초당 약 75 토큰을 봅니다. 동시성 64까지의 지점은 7개의 추측 토큰과 합성 수용 길이 3.84 토큰을 사용합니다. 더 높은 동시성은 3개의 추측 토큰과 수용 길이 3.0을 사용합니다.
데이 0에서의 제로-미스매치 RL
vLLM의 데이-0 지원은 처음으로 제로-미스매치 RL을 제공합니다: verl을 사용한 LoRA RL, 롤아웃 엔진으로 vLLM, 트레이너로 Megatron을 사용하며, 트레이너의 로그 확률이 롤아웃과 비트 단위로 일치합니다. 중요도 샘플링이나 라우팅 리플레이로 격차를 제한하는 대신 vLLM은 트레이너의 순방향 패스가 롤아웃을 비트 단위로 재현하도록 하여 커널 수준에서 격차를 제거합니다.
- 전체 모델에 대한 LoRA RL. Rank-32 LoRA는 모든 어텐션 프로젝션과 모든 896개 전문가를 포함하며, 기본 모델과 어댑터 모두 vLLM과 Megatron 간에 비트 단위로 동일합니다.
- 배치 불변 커널, 기본적으로.
VLLM_BATCH_INVARIANT=1을 사용하면 토큰의 로그 확률이 배치에 의존하지 않으며, vLLM의 고성능 커널 중 더 많은 것이 설계상 배치 불변입니다(예: DeepSelect top-k).
전체 모델에서 100단계 이상 동안 롤아웃/트레이너 KL은 모든 단계에서 정확히 0.0이며 AIME mean@8은 0.451에서 0.664로 상승합니다. vLLM은 배치 불변 경로의 성능 최적화가 진행 중이며 소스 코드는 Inferact/zero-mismatch-rl에 있다고 밝혔습니다.
AMD Instinct MI355X 지원
AMD Instinct MI355X에서 vLLM은 MXFP4 인덱싱을 희소 어텐션 및 디코드 컨텍스트 병렬 처리와 결합합니다. 최적화는 긴 컨텍스트의 메모리 요구 사항을 줄이고 프리필 및 디코드 효율성을 향상시킵니다.
네이티브 MXFP4 인덱서
인덱서는 어텐션을 위해 2,048개 토큰을 선택하기 전에 컨텍스트를 스캔합니다. MXFP4는 각 인덱스 키를 FP8의 132바이트에서 스케일 포함 68바이트로 줄여 인덱스 캐시 저장 공간을 48.5% 절감합니다. 점수화는 FP32 누적 및 헤드 간 가중 감소와 함께 네이티브 FP4 행렬 명령을 사용합니다. 키는 추측 검증 중에 쿼리 행 간에 재사용되어 반복 읽기를 줄입니다. AgentX 스윕에서 MXFP4 인덱서 구성은 동시성 48에서 총 워크로드 처리량을 5.3% 더 높였습니다.
컨텍스트 병렬 희소 어텐션 프리필
프리필 중에 GPU는 인덱스 키를 수집하고 쿼리 행을 서로 나눕니다. 각 GPU는 전체 컨텍스트에 대해 할당된 행을 점수화하고 전역 top-k를 직접 선택하여 순위별 후보 병합을 피합니다. 기본 어텐션은 통신을 줄이고 워크스페이스 예산에 맞을 때 KV를 별도로 수집한 다음 로컬 헤드에 대한 모든 쿼리 행을 처리할 수 있습니다.
| 입력 / 출력 토큰 | 총 처리량 (토큰/초) | 출력 처리량 (토큰/초) | 중앙값 TTFT (초) | 중앙값 TPOT (ms) |
|---|---|---|---|---|
| 8K / 1K | 5,357.8 → 6,015.4 (+12%) | 595.3 → 668.4 | 2.462 → 1.875 | 24.571 → 22.063 |
| 128K / 100 | 6,950 → 17,678 (+154%) | 5.30 → 13.48 | 268.7 → 105.5 | 331.5 → 129.6 |
설정: MI355X, TP8/DCP8, 동시성 16, 64개 프롬프트, 프리픽스 캐싱 비활성화.
FlyDSL 희소 MLA는 선택된 토큰을 FP8 KV 캐시에서 직접 읽습니다. 각 워크그룹은 로컬 어텐션 헤드에서 하나의 쿼리 행을 처리하여 프리필 단계의 많은 쿼리 행에 걸쳐 병렬 처리를 제공합니다. 더 작은 배치는 각 행 내에서 추가 병렬 처리를 노출하는 Triton 경로를 유지합니다.
| 구현 | 벤치마크 기간 (초) | 중앙값 TTFT (초) | 출력 처리량 (토큰/초) |
|---|---|---|---|
| Triton | 432.8 | 96.2 | 14.79 |
| FlyDSL | 404.4 | 89.8 | 15.83 |
| 개선 | 6.6% 낮음 | 6.6% 낮음 | 7.0% 높음 |
설정: MI355X, TP8/DCP8, 동시성 16, 64개 프롬프트, ISL/OSL 128K/100.
융합 저지연 디코드
작은 추측 검증 배치의 경우 DSA 인덱싱과 희소 어텐션이 별도로 실행된 다음 어텐션 출력 프로젝션, AttnRes, 전체 잠재 MoE 계산(텐서 병렬 all-reduce 포함)을 단일 실행으로 융합하는 영구 모노커널이 이어집니다. 이는 커널 실행 오버헤드를 줄이고 MoE 가중치 로딩을 이전 감소 및 잔여 처리와 겹칩니다. 전문화된 경로는 단계당 최대 8개의 검증 토큰을 지원하며, 7개의 드래프트 토큰이 있는 단일 요청을 포함합니다. 평균 TPOT는 6.02ms에서 4.40ms로 27.0% 개선되며, TP8/DCP1, FP8 인덱서, 8K/1K, 동시성 1, 7개 드래프트 토큰, 합성 수용 길이 3.84에서 측정되었습니다.
GPU 간 통신
디코드 중에 각 GPU는 토큰 ID를 로컬로 유지하면서 후보 점수를 교환합니다. 공유 임계값과 일관된 타이 처리로 어떤 로컬 후보가 선택되는지 결정합니다. 2,048개 후보의 경우 점수만 교환하면 후보 페이로드가 GPU당 쿼리당 16KiB에서 8KiB로 절반으로 줄어듭니다. AITER는 HIP IPC를 통해 피어 버퍼를 읽고 필요한 출력 레이아웃에 직접 쓰는 전용 GPU 커널로 DCP all-gather를 가속화하여 데이터 수집과 레이아웃 변환을 결합합니다. 8개의 MI355X GPU에서 독립형 마이크로벤치마크에서 AITER 경로는 테스트된 쿼리 및 후보 텐서 형태에 대해 RCCL보다 all-gather 지연 시간을 53-61% 줄였습니다.
Google TPU 지원
vLLM은 vllm-torchtpu를 통해 TPU에서 Kimi K3.1 Pro를 데이 0에 지원하며, TPU v7x의 16칩 슬라이스에서 검증되었습니다. 레이아웃은 MoE에 대해 TP32/EP32이며, 긴 컨텍스트를 위해 MLA 잠재 표현과 인덱스 키를 토큰별로 샤딩하는 DCP가 있습니다. DSA의 작업은 TPU의 두 종류의 컴퓨팅 유닛으로 분할됩니다: TensorCore가 FP8 인덱스 키를 점수화한 다음 SparseCore가 상위 2,048개 토큰을 선택하고 해당 잠재 표현을 TensorCore의 MXU가 어텐션을 실행할 수 있도록 밀집 버퍼로 수집합니다. GPU에서와 같이 DSA 및 KDA 상태는 하나의 블록 인덱스 KV 풀을 공유하여 하이브리드 캐시에서 프리픽스 캐싱이 계속 작동하도록 합니다.
Alibaba T-HEAD Zhenwu PPU 지원
Alibaba T-HEAD 팀은 vllm-sail 플러그인을 통해 Alibaba T-HEAD Zhenwu PPU(M890 및 810E 모두)에서 Kimi K3.1 Pro 실행 지원을 추가했습니다. 검증된 구성은 다음과 같습니다:
- 810E: 두 노드, TP32 / EP32, W4A16 MoE 백엔드 사용.
- M890: 네 노드, TP8 / DP4 / EP32, DeepEP 사용.
vLLM은 이것이 Zhenwu PPU가 주류 오픈소스 대형 모델에 대해 데이-0 지원을 제공하는 첫 번째 사례라고 설명하며, 곧 지속적인 지원이 제공될 것이라고 밝혔습니다.
빠른 시작
NVIDIA GPU에서 vLLM은 텐서 병렬 처리 8, fastsafetensors 로딩, Kimi K3 도구 호출 및 추론 파서, 5개의 추측 토큰이 있는 DSpark 추측 구성을 사용하는 vllm serve 명령을 제공합니다:
vllm serve moonshotai/Kimi-K3.1-Pro \
--tensor-parallel-size 8 \
--trust-remote-code \
--load-format fastsafetensors \
--enable-auto-tool-choice \
--tool-call-parser kimi_k3 \
--reasoning-parser kimi_k3 \
--speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'
AMD GPU에서 명령은 AITER 및 빠른 GPU safetensors 로딩을 활성화하고, 텐서 병렬 처리 8을 설정하고, 시퀀스를 16으로 제한하고, GPU 메모리 사용률을 0.85로 설정하고, 동일한 DSpark 추측 구성을 사용합니다:
VLLM_ROCM_USE_AITER=1 SAFETENSORS_FAST_GPU=1 \
vllm serve moonshotai/Kimi-K3.1-Pro \
--served-model-name kimi-k31-pro \
--tensor-parallel-size 8 \
--trust-remote-code \
--load-format fastsafetensors \
--max-num-seqs 16 \
--gpu-memory-utilization 0.85 \
--speculative-config '{"model":"Inferact/Kimi-K3.1-Pro-DSpark-preview","num_speculative_tokens":5,"method":"dspark","draft_sample_method":"probabilistic"}'
릴리스가 하이브리드 희소 모델 서빙에 대해 보여주는 것
Kimi K3.1 Pro는 토큰당 또는 요청당 세 가지 다른 종류의 상태(MLA 잠재 표현, DSA 인덱스 키, KDA 재귀 상태)를 가진 하이브리드 모델입니다. vLLM의 데이-0 지원은 이러한 모델을 효율적으로 서빙하려면 단일 어텐션 커널이 아니라 KV 캐시 레이아웃, 프리필 샤딩, 디코드 병렬 처리, 저지연 커널 융합에 걸친 조정된 변경이 필요함을 보여줍니다.
패킹된 블록 캐시는 가능하게 하는 추상화입니다: 어텐션 커널이 고정 크기 페이지를 유지하면서 서빙 스택이 혼합 너비 상태를 연속적으로 저장하여 RDMA 전송 및 프리픽스 캐싱을 가능하게 합니다. 토큰 샤딩 프리필은 텐서 병렬 처리에서 MLA의 헤드 복제 문제를 해결하고, 정확한 top-k 병합이 있는 DCP는 긴 컨텍스트에서 TP의 KV 용량 제한을 해결합니다. 융합 디코드 경로와 DSpark 추측 디코딩은 에이전트 워크로드가 의존하는 소형 배치 지연 시간을 목표로 합니다.
제로-미스매치 RL 결과는 추론 성능과는 별개이지만 훈련 인프라에 중요합니다: vLLM은 기본 모델과 모든 어텐션 프로젝션 및 모든 896개 전문가를 포함하는 rank-32 LoRA 어댑터 모두에 대해 vLLM 롤아웃과 Megatron 훈련 간에 비트 단위로 동일한 로그 확률을 보고하며, 100단계에 걸쳐 KL이 정확히 0.0입니다. 이는 중요도 샘플링이나 라우팅 리플레이로 제한하는 대신 RL 불안정성의 원인을 제거합니다.
플랫폼 범위
데이-0 지원은 NVIDIA GPU, AMD GPU, Google TPU, Alibaba T-HEAD Zhenwu PPU의 네 가지 하드웨어 제품군에 걸쳐 있습니다. NVIDIA 경로에는 SemiAnalysis AgentX에 대한 GB300 및 Rubin 벤치마크가 포함됩니다. AMD 경로에는 MXFP4 인덱싱, 컨텍스트 병렬 희소 프리필, 융합 저지연 디코드, AITER all-gather에 대한 MI355X 최적화가 포함됩니다. TPU 경로는 TP32/EP32 및 DCP를 사용하여 16칩 TPU v7x 슬라이스에서 vllm-torchtpu를 통해 실행됩니다. T-HEAD 경로는 810E 및 M890 구성에서 vllm-sail 플러그인을 통해 실행됩니다.
vLLM은 Kimi K3.1 Pro를 오픈소스로 공개한 Moonshot에 공을 돌리고, 데이-0 지원에 대해 Inferact, AMD, Alibaba T-HEAD에 감사하며, NVIDIA의 도움에 감사합니다.