Kimi K3, MI355X에서 성능 혁신을 이끌다

Kimi K3와 MI355X: 성능 혁신의 시작

최근 몇 달 동안 오픈 소스 모델의 능력이 급격히 향상되었습니다. DeepSeek V4-Pro와 GLM5.2는 Opus 수준에 가까운 지능을 보여주며, 오픈 소스는 이제 폐쇄형 모델에 대한 비용 효율적인 대안으로 부상했습니다. 그러나 Kimi K3와 같은 모델은 아직 보지 못했습니다. Fable/Sol 수준의 지능을 약속하는 Kimi K3는 오픈 소스의 새로운 시대를 여는 모델입니다. 그러나 더 똑똑한 모델은 더 큰 모델을 의미하며, 이러한 모델은 능력만큼이나 빠르게 크기가 커지고 있습니다. GLM5.2는 753B 파라미터, DeepSeek V4-Pro는 1.6T, 그리고 Kimi K3는 무려 2.8T 파라미터를 가지고 있습니다. 이는 1M 토큰의 컨텍스트를 위한 KV 캐시를 할당하기 전에 1.5TB 이상의 VRAM을 필요로 합니다. B200 노드(8개의 GPU)조차도 Kimi K3를 수용할 수 없습니다. 따라서 B300 노드(각 GPU당 288GB VRAM)를 사용하거나 두 개의 B200 노드(TP16)를 Kimi에 할당하는 제한된 옵션이 남습니다. 그러나 AMD의 MI355X도 288GB VRAM을 제공합니다. 이 칩을 좋아하지 않을 수 없습니다. 평균적으로 B300보다 약 2.4배, B200보다 약 1.7배 저렴한 MI355X는 하드웨어 사양이 비슷한 Blackwells에 대한 비용 효율적인 대안입니다.

AMD의 유일한 문제는 소프트웨어 지원입니다. 느린 커널과 추론 프레임워크에 대한 초기 지원 부족은 AMD에서 최첨단 모델을 제공하는 데 실질적인 엔지니어링 노력을 요구합니다. Wafer의 주장에 따르면 에이전트가 커널과 모델 최적화에서 개선되고 있으며, 이 격차를 줄이고 있습니다. 그러나 AMD가 Kimi K3에 대한 초기 지원을 제공하면서 대부분의 작업이 이미 완료되었습니다. 결과는 훌륭합니다. 1,024-토큰 입력 / 400-토큰 출력 벤치마크에서 MI355X는 노드당 952 tok/s, 단일 스트림에서는 118 tok/s를 기록했습니다. 이는 TP16 B200 배포의 단일 스트림 디코드(498 tok/s, 16-GPU, 2-노드 총계 — 노드당 약 249 tok/s)보다 3.8배 이상의 집계 처리량과 1.3배 이상의 단일 스트림 디코드를 보여줍니다. B300 노드는 여전히 MI355X보다 집계 처리량에서 약 1.65배 우세하지만, 가격이 2.4배인 상황에서 MI355X는 성능 대비 비용에서 B300을 압도합니다.

• 8× MI355X (TP8)

  • Decode tok/s per stream: 118 tok/s
  • Peak aggregate: 952 tok/s
  • Peak aggregate per GPU: 119 tok/s
  • Peak aggregate per $/GPU-hr: 48 tok/s/$

• 2×8 B200 (TP16)

  • Decode tok/s per stream: 90 tok/s
  • Peak aggregate: 498 tok/s
  • Peak aggregate per GPU: 31 tok/s
  • Peak aggregate per $/GPU-hr: 7 tok/s/$

• B300 (TP8+DCP8)

  • Decode tok/s per stream: 172 tok/s
  • Peak aggregate: 1,568 tok/s
  • Peak aggregate per GPU: 196 tok/s
  • Peak aggregate per $/GPU-hr: 33 tok/s/$

MI355X의 성능 대비 비용 비율은 GPU-hr당 $2.50, B300은 $6.00, B200은 $4.25입니다. B200의 경우, 디코드 중요 경로에서 크로스 노드 올-리듀스를 수행해야 하기 때문에 수치가 다소 낮게 나타납니다(RoCE v2에서 약 195 Gb/s). 이는 여기서 유일하게 두 노드를 사용하는 구성입니다. Kimi K3는 1M-토큰 KV 풀과 함께 단일 8×192GB 노드에 맞지 않기 때문입니다. 그러나 이것이 바로 포인트입니다. Kimi K3의 크기에서 MI355X의 HBM 용량 초점은 B200에 비해 실질적이고 측정 가능한 우위를 제공합니다.

성능 최적화 과정

Kimi K3는 기본적으로 제공되었지만, 현재의 처리량 수치에 도달하기 위해 추가 작업이 필요했습니다. 주요 레버는 추측 디코드였습니다. K3는 초안 텐서를 제공하지 않기 때문에 외부 블록 확산 초안인 RadixArk의 Kimi-K3-DSpark를 사용해야 했습니다. CUDA에서는 잘 작동하지만, ROCm에서는 첫 번째 실제 요청에서 스케줄러가 'NameError: name 'top_k_renorm_prob' is not defined'라는 오류로 중단되었습니다. 이 문제를 해결하기 위해 PyTorch 함수를 사용했습니다. Top-k renorm은 작은 연산으로, 모델의 확률 벡터에서 가장 높은 k개의 항목을 유지하고 나머지를 0으로 설정한 후 남은 항목을 1로 합산하도록 재조정합니다. 이 연산을 sglang의 ROCm 샘플링 브랜치에 추가하여 CUDA 빌드에서 sgl_kernel에서 가져오는 동일한 계산을 수행했습니다. 이 수정으로 단일 스트림 성능이 약 2.2배, 중간 부하에서 스트림당 약 1.7배, 피크 집계가 18% 향상되었습니다. 더 중요한 것은, 피크 집계 처리량이 더 높은 동시성(c64 vs c24 no-spec)에서 달성되었다는 점입니다.

프리필 최적화

모델 성능에 대한 논의는 주로 초당 디코드 토큰에 초점을 맞춥니다. 그러나 많은 경우 디코드 tok/s는 과대평가되며, 사용자가 가장 많이 느끼는 첫 번째 토큰까지의 시간이 간과됩니다. MI355X는 여기서 어려움을 겪습니다. 동일한 172k-토큰 콜드 프리필이 MI355X에서는 약 51초, B300에서는 약 23초가 걸렸습니다. 1M-컨텍스트 모델에서는 많은 작업이 대규모 프리필을 필요로 하며, GPU가 프리필에 몇 분 동안 소요되면 전체 노드 플릿이 무용지물이 될 수 있습니다. 이 격차는 거의 전적으로 하나의 커널 때문이었습니다. K3는 ROCm에서 느린 일반 Triton 주의를 사용하고 있었고, 빠른 AITER MLA 프리필 커널이 로드되지 않았습니다. 문제는 모양 불일치였으며, K3는 TP8에서 랭크당 12개의 주의 헤드를 제공하고, AITER의 MLA 경로는 4, 8 또는 16의 배수로 구축되었습니다. 해결책은 간단했습니다: 헤드 수를 12→16으로 제로 패딩하고, 빠른 커널을 실행한 후 실제 12개의 헤드를 출력에서 추출했습니다. 결과적으로 동일한 172k 콜드 프리필에서 AITER MLA 프리필 ASM은 Triton 대체의 ~4–7k보다 ~13k tok/s의 안정적인 상태로 실행되어 프리필 속도를 ~2–3배 향상시켰습니다. 이는 TTFT 레버이지이며, 집계 처리량에는 영향을 미치지 않습니다. 디코드는 변경되지 않으므로 위의 숫자에는 영향을 미치지 않지만, 사용자가 첫 번째 토큰이 나타나기까지 기다리는 시간을 줄입니다.