Skip to content

Add Kimi-K3 for MI355X: vLLM, SGLang, and ATom - #182

Open
coketaste wants to merge 6 commits into
ROCm:developfrom
coketaste:coketaste/add_vllm_models
Open

Add Kimi-K3 for MI355X: vLLM, SGLang, and ATom#182
coketaste wants to merge 6 commits into
ROCm:developfrom
coketaste:coketaste/add_vllm_models

Conversation

@coketaste

@coketaste coketaste commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds Kimi-K3 (moonshotai/Kimi-K3) benchmark support for MI355X (gfx950, TP8) across three serving frameworks:

Tag Framework Docker image Status
pyt_vllm_kimi-k3 vLLM vllm/vllm-openai-rocm:kimi-k3 Day-0 image available
pyt_sglang_kimi-k3 SGLang lmsysorg/sglang-rocm:rocm720-mi35x-k3-20260727 Day-0 image available
pyt_sglang_kimi-k3_dspark SGLang + DSpark same as above Day-0 image available (DSpark has known ROCm crash sglang#32569)
pyt_atom_kimi-k3 ATom rocm/atom-dev:latest Awaiting ROCm/ATOM#1718 merge

All entries use skip_gpu_arch to exclude non-gfx950 architectures (K3 MXFP4 weights require gfx95x).

vLLM (pyt_vllm_kimi-k3)

  • Separate Dockerfile (docker/pyt_vllm_kimi_k3) pinning vllm/vllm-openai-rocm:kimi-k3 (requires vLLM >= 0.27.0)
  • Config in scripts/vllm/configs/default.yaml with AITER env vars

SGLang (pyt_sglang_kimi-k3, pyt_sglang_kimi-k3_dspark)

  • Separate Dockerfile (docker/pyt_sglang_kimi_k3) pinning the day-0 image
  • Dedicated scripts/sglang/run_kimi_k3.sh + run_sglang.py driver
  • Config in scripts/sglang/configs/kimi_k3.yaml with two YAML-anchored variants (nospec / dspark)
  • Sweep: ISL 8192, OSL 1024, mc 2/4/8/16/32 (matches AMD day-0 tables in sglang#32548)

ATom (pyt_atom_kimi-k3)

  • New ATom integration for MAD (Dockerfile, run scripts, configs), ported from MAD-Private client-perf
  • Recipe from ROCm/ATom PR #1718 (recipes/Kimi-K3.md)
  • Env vars: ATOM_USE_TRITON_GEMM=1, AITER_FLYDSL_FORCE=1, ATOM_USE_TRITON_MOE=0, etc.
  • Server flags: --kv_cache_dtype fp8, --gpu-memory-utilization 0.93, --block-size 128, --no-enable_prefix_caching
  • Includes GSM8K 5-shot accuracy config (validated range: 0.9538-0.9591)
  • run_atom.py supports serving benchmarks, accuracy tests, and kernel profiling

Test plan

  • Validate JSON/YAML configs
  • Verify tags resolve: `madengine run --tags pyt_atom_kimi-k3 --dry-run`
  • Run on MI355X with K3-enabled images

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 coketaste self-assigned this Jul 28, 2026
Copilot AI review requested due to automatic review settings July 28, 2026 20:39

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 serving config stanza for Kimi K3 (TP=8 only) with recipe-aligned flags/env and accuracy disabled for gsm8k.
  • Register a new models.json entry (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.

coketaste and others added 5 commits July 28, 2026 17:22
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>
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>
Copilot AI review requested due to automatic review settings July 28, 2026 23:10

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review is ineligible. To be eligible to request a review, you need a paid Copilot license, or your organization must enable Copilot code review.

@itej89 itej89 changed the title Add vLLM Kimi K3 for MI355X [WIP] Add Kimi-K3 for MI355X: vLLM, SGLang, and ATom Jul 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants