Add Kimi-K3 for MI355X: vLLM, SGLang, and ATom - #182
Open
coketaste wants to merge 6 commits into
Open
Conversation
Add pyt_vllm_kimi-k3 benchmarking moonshotai/Kimi-K3, a 2.8T-parameter MoE (16 of 896 experts active, 1M context) whose weights ship natively MXFP4 from quantization-aware training, so unlike the K2.6 pair there is no separate amd/ FP4 repo and only one model entry is needed. K3 requires vLLM >= 0.27.0, which is not yet in a tagged vllm-openai-rocm release. Add docker/pyt_vllm_kimi_k3 pinned to the model-specific :kimi-k3 image rather than retargeting docker/pyt_vllm, which would move ~35 other vLLM entries onto a pre-release base. The Dockerfile body is otherwise identical and should be folded back once K3 lands in a versioned ROCm image. Serving flags and env mirror the single-node TP8 gfx950 profile from recipes.vllm.ai/moonshotai/Kimi-K3/hw/mi355x.json. TP8 only: the ~1680 GB minimum footprint fits an 8x MI355X node (2304 GB) but not TP4, and not a single 8x MI300X node, hence skip_gpu_arch gfx942. The model is deliberately left out of the vllm_default sweep so a tag run does not pull the ~1.56 TB checkpoint and hold 8 GPUs; it is invocable explicitly by name. Two deviations from the recipe are intentional. MAD adds --no-enable-prefix-caching for benchmark hygiene, as it does for every model here. The gsm8k accuracy stage is disabled because lm_eval drives /v1/completions, where K3's always-on reasoning is returned inline and exhausts the max_gen_toks=2048 budget hardcoded in run_vllm.py, making the score meaningless. Verified statically only, no hardware run: the emitted server command was rendered through run_vllm.expand_configs with MAD_SYSTEM_GPU_ARCHITECTURE=gfx950 and diffed against the recipe argv with no missing flags, and the vllm_default sweep and the expansion counts for K2.6, DeepSeek-R1 and GLM-5.2 are unchanged. Note that the 1800s server-start timeout in run_vllm.py may be tight for a 1.56 TB load and is left for a follow-up. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
coketaste
requested review from
Rohan138,
amathews-amd,
gargrahul and
ppalaniappan-amd
as code owners
July 28, 2026 20:39
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new MAD vLLM benchmark entry for moonshotai/Kimi-K3 targeting MI355X/MI350X-class (gfx950) nodes, using a dedicated ROCm vLLM image to avoid moving the rest of the vLLM suite onto a pre-release base.
Changes:
- Add a new
servingconfig stanza for Kimi K3 (TP=8 only) with recipe-aligned flags/env and accuracy disabled for gsm8k. - Register a new
models.jsonentry (pyt_vllm_kimi-k3) pointing to a dedicated Dockerfile and skipping gfx942. - Add a model-specific Dockerfile pinned to
vllm/vllm-openai-rocm:kimi-k3, and document the new model + exception in the vLLM benchmark README.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| scripts/vllm/configs/default.yaml | Adds Kimi K3 serving benchmark config (TP8, recipe-aligned flags/env, accuracy disabled). |
| models.json | Registers pyt_vllm_kimi-k3 using a dedicated Dockerfile and gfx942 skip. |
| docker/pyt_vllm_kimi_k3.ubuntu.amd.Dockerfile | Introduces a separate vLLM ROCm base image (:kimi-k3) for K3-only requirements. |
| benchmark/vllm/README.md | Documents the new model and explains why it uses a separate Docker image + how to run it. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Add pyt_sglang_kimi-k3 and pyt_sglang_kimi-k3_dspark benchmarking moonshotai/Kimi-K3 online serving, mirroring the vLLM K3 entry added in a40063d. Both track the AMD day-0 recipes in sgl-project/sglang#32548; K3 support itself is sgl-project/sglang#32541. Two entries rather than one because the two published recipes differ only in speculative decoding (DSPARK with the RadixArk/Kimi-K3-DSpark draft) and are reported as separate rows in the tracking issue, so both are worth tracking. They share a model repo, which is why the runner gains a --variant filter: the existing --model filter cannot select between them. Add scripts/sglang/run_sglang.py and configs/, a config-driven serving runner structured after scripts/vllm/run_vllm.py and emitting the same perf CSV schema, so multiple_results ingestion is shared. The existing scripts/sglang/run.sh is left untouched; it is bound to the offline getopts path and hard-exits on non-gfx94* GPUs. Add docker/pyt_sglang_kimi_k3 pinned to lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727, the day-0 image the issue's MI355X tables were measured against, rather than retargeting docker/pyt_sglang, which is still on lmsysorg/sglang:v0.4.5-rocm630 and would move the existing SGLang entry onto a pre-release base. A newer -20260728 tag exists but nothing published ties it to the recipe, so it is deliberately not adopted. Fold back once K3 lands in a versioned ROCm image. TP8 gfx950 only, hence skip_gpu_arch gfx942. Both are left out of any sweep tag so a tag run does not pull the checkpoint and hold 8 GPUs; they are invocable explicitly by name. Three deviations from the recipe are intentional. The runner emits --tp-size where the recipe writes --tp: there is no --tp server argument, tp_size declares the single alias --tensor-parallel-size, and --tp resolves only through argparse prefix matching, which any future --tp* option would silently break. The server-start timeout is raised to 5400s from run_vllm.py's 1800s, which that commit already flagged as tight for a checkpoint this size. There is no gsm8k accuracy stage, for the same reason as the vLLM entry: K3's always-on reasoning is returned inline over /v1/completions and exhausts the hardcoded 2048-token budget. The issue does not state its input and output lengths. They were recovered from its own tables, where (E2EL - TTFT) / TPOT + 1 lands on ~1024 output tokens on every row and concurrency x (inp + out) / E2EL reproduces the reported total throughput only at 8192 input tokens; keeping that shape is what makes these numbers comparable to the issue's. --random-range-ratio is pinned to 1.0 because it defaults to 0.0, which would randomize lengths down to a single token. Verified statically only, no hardware run. Config expansion yields 5 configs per variant and the variant filter separates them; an end-to-end dry run with subprocess stubbed produced the expected server and benchmark argv, both hf download calls, and correctly ignored a stale leading line in the appended JSONL output. All 14 server flags and 11 benchmark flags were validated against #32541's head branch rather than sglang main, which matters: the PR is still open and kimi_k3 is absent from main's reasoning DetectorMap, so the parser names and DSPARK are only valid inside the pinned image. Synthetic results were fed through madengine's handle_multiple_results, and filter_images_by_skip_gpu_arch skips both entries on gfx942 and runs them on gfx950 with no existing entry changed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
The Blueprints table in the root README is the entry point to each framework's landing page, so the SGLang row's Models column should name K3 now that both entries exist. The landing page itself, benchmark/sglang/README.md, was updated in ede3844. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…to coketaste/add_vllm_models
Introduces ATom (ROCm/ATOM) integration to MAD, modeled after MAD-Private client-perf. Includes Dockerfile, run scripts, and Kimi-K3 config based on ROCm/ATom PR #1718 recipe (gfx950 TP8, MXFP4 w/ FlyDSL SiTUv2, FP8 KV cache). Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds Kimi-K3 (moonshotai/Kimi-K3) benchmark support for MI355X (gfx950, TP8) across three serving frameworks:
pyt_vllm_kimi-k3vllm/vllm-openai-rocm:kimi-k3pyt_sglang_kimi-k3lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727pyt_sglang_kimi-k3_dsparkpyt_atom_kimi-k3rocm/atom-dev:latestAll entries use
skip_gpu_archto exclude non-gfx950 architectures (K3 MXFP4 weights require gfx95x).vLLM (
pyt_vllm_kimi-k3)docker/pyt_vllm_kimi_k3) pinningvllm/vllm-openai-rocm:kimi-k3(requires vLLM >= 0.27.0)scripts/vllm/configs/default.yamlwith AITER env varsSGLang (
pyt_sglang_kimi-k3,pyt_sglang_kimi-k3_dspark)docker/pyt_sglang_kimi_k3) pinning the day-0 imagescripts/sglang/run_kimi_k3.sh+run_sglang.pydriverscripts/sglang/configs/kimi_k3.yamlwith two YAML-anchored variants (nospec / dspark)ATom (
pyt_atom_kimi-k3)client-perfrecipes/Kimi-K3.md)ATOM_USE_TRITON_GEMM=1,AITER_FLYDSL_FORCE=1,ATOM_USE_TRITON_MOE=0, etc.--kv_cache_dtype fp8,--gpu-memory-utilization 0.93,--block-size 128,--no-enable_prefix_cachingrun_atom.pysupports serving benchmarks, accuracy tests, and kernel profilingTest plan