[AMD] Add MI355X Kimi-K2.6 tuning artifacts by jhinpan · Pull Request #23381 · sgl-project/sglang

jhinpan · 2026-04-21T14:34:03Z

Summary

This PR unblocks Kimi-K2.6 W4A16 MoE serving on AMD Instinct MI355X without downstream forks.

Changes:

Add a tuned Triton fused-MoE config for E=384, N=256, device_name=AMD_Instinct_MI355X, dtype=int4_w4a16 under python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/.
Cover 13 batch-size buckets: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 8192].
Add SGLANG_DEEPSEEK_LOAD_MAX_WORKERS as an opt-in cap for the DeepSeek-style checkpoint loader thread pool. Default is None, so unset preserves Python ThreadPoolExecutor default behavior; large MoE deployments can set a smaller value such as 4 to reduce aggregate host I/O pressure across ranks.
Keep the fused-MoE tuner usable on ROCm by using local device id 0 inside Ray HIP workers.
Document the new env var in both docs trees and clean changed-file lint issues after merging current main.

Benchmark Evidence

Validated on 4x AMD Instinct MI355X with Kimi-K2.6, Triton 3.6.0, ROCm 7.2, --decode-attention-backend triton, --prefill-attention-backend aiter, and CUDA graph enabled.

sglang.bench_one_batch_server --batch-size 1 --output-len 1024:

input	output	TTFT (s)	Decode (tok/s)	Overall (tok/s)	Total latency (s)
1,024	1,024	0.44	40.24	80.47	25.45
2,048	1,024	0.44	40.01	120.04	25.59
4,096	1,024	0.66	39.07	195.33	26.21
8,192	1,024	0.70	38.05	342.44	26.91
16,384	1,024	1.17	35.85	609.49	28.56
32,768	1,024	2.42	32.03	1056.88	31.97

A/B against the same launch command with only the new MoE config file removed:

input	output	TTFT A (s)	TTFT C (s)	TTFT delta	Decode A (tok/s)	Decode C (tok/s)	Decode delta
1,024	1,024	0.857	0.437	-49.0%	37.41	40.24	+7.6%
2,048	1,024	0.723	0.443	-38.7%	37.39	40.01	+7.0%
8,192	1,024	1.846	0.702	-62.0%	34.60	38.05	+10.0%

The benchmark image used for the serving validation was built from this PR branch before the final current-main merge/docs/lint refresh; those final changes do not alter the tuned config or serving path.

Recommended Runtime Setting

For very large MoE checkpoints on NFS-backed storage, use an explicit cap such as:

export SGLANG_DEEPSEEK_LOAD_MAX_WORKERS=4

Leaving the variable unset preserves existing SGLang behavior.

Test Plan

Merged current sgl-project/sglang:main into the PR branch and resolved the conflict in python/sglang/srt/environ.py.
python -m ruff check benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py python/sglang/srt/environ.py python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
python -m ruff format --check benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py python/sglang/srt/environ.py python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py
python -m py_compile python/sglang/srt/environ.py python/sglang/srt/models/deepseek_common/deepseek_weight_loader.py benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py
python -m json.tool python/sglang/srt/layers/moe/moe_runner/triton_utils/configs/triton_3_6_0/E=384,N=256,device_name=AMD_Instinct_MI355X,dtype=int4_w4a16.json
Local PYTHONPATH=python smoke confirms envs.SGLANG_DEEPSEEK_LOAD_MAX_WORKERS defaults to None and parses 4 through EnvInt.override.
Served Kimi-K2.6 on 4x MI355X with the tuned config; /health returned 200 and one chat completion round trip worked.
sglang.bench_one_batch_server --batch-size 1 --input-len 1024 2048 4096 8192 16384 32768 --output-len 1024 produced the table above.

CI States

Latest PR Test (Base): ❌ Run #27216334917
Latest PR Test (Extra): 🚫 Run #27216335222

Add the tuned MI355X fused-MoE config used by Kimi-K2.6 and make the DeepSeek weight loader parallelism configurable for large MoE checkpoints. Also make the fused-MoE tuner pick a stable HIP device id so the config can be regenerated from upstream sources.

gemini-code-assist · 2026-04-21T14:34:09Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

- Register SGLANG_DEEPSEEK_LOAD_MAX_WORKERS via envs (EnvInt) instead of a raw os.environ.get; default is None so ThreadPoolExecutor keeps its prior pass-through behaviour (min(32, cpu_count() + 4)) on NV hosts that do not opt in. Large MoE checkpoints on MI355X can still set SGLANG_DEEPSEEK_LOAD_MAX_WORKERS=4. - Drop inline `import os as _os` in the weight loader and route through the already-imported `envs` module; fixes the black line-length CI failure as a side effect. - Use the module-level `_is_hip` cache in the MoE tuner instead of calling `is_hip()` per BenchmarkWorker init. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Extend the fused-MoE Triton config for (E=384, N=256, MI355X, int4_w4a16) from 4 to 13 batch-size buckets so nearest-neighbor fallback is tighter for mid-M workloads (e.g. moderate concurrency, chunked prefill). Final coverage: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 8192], which matches Arist12's K2.5 default bucket set plus our longer-prefill 8192 bucket. Each bucket was produced by the upstream benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py sweep over 400 candidate configs per bucket.

jhinpan · 2026-04-22T19:33:47Z

cc @sunxxuns @HaiShaw Can u help review and trigger CICD for this PR? Huge Thx!

Adds a hand-authored eval instance that targets the FlyDSL fused-MoE port from AMD's Kimi-K2.5 MI300X optimization blog to Kimi-K2.6 on 4x MI355X, on top of the baseline established by sgl-project/sglang#23381. Contents: - task.yaml : optimize-type task, phase1_baseline=true, frontier_model=true, TP=4, 4xMI355X, FlyDSL env OFF at start so Phase 1 reproduces PR #23381 numbers (DSL2_ROOT + MLIR_PATH + CK_TILE_FLOAT_TO_BFLOAT16 pre-set so import works as soon as agent flips gates during a trial) - Dockerfile : layers FlyDSL + AITER dev/kimi-K2.5 on top of jhinpan/sglang-k26-mi355x:v0.5.10rc0-rocm720-20260420 - bench_flydsl_k26.sh : single bench script emitting one canonical line: output_throughput_tok_s: <v> | concurrency=40 in=10240 out=512 decode_bs1_in8k=<guard> - task_description.md : full phased plan A-G, env-var reference table, BS=1 decode guard (>= 0.98x PR #23381 ~38.05 tok/s), front-loaded supervisor-visible gating block - test_harness.py : GSM8K accuracy gate (lm_eval limit=50 >= 0.90) - metadata.json : links to GH issue amdpilot-org/sglang#2 + PR #23381 - ISSUE.md : condensed body as filed on amdpilot-org/sglang#2 Tracks GitHub issue amdpilot-org/sglang#2.

* add primus-qwen3-30b-mfu eval instance for issue #1 8x MI355X Qwen3-30B-A3B MoE pretraining MFU optimization. Built on ghcr.io/amdpilot-org/primus-mi355x-ready:v1. Two task.yaml variants for A/B testing the Phase 1 baseline agent: - task.yaml: phase1_baseline: true - task_nophase1.yaml: phase1_baseline: false Both pin the executor to Kimi-K2.6 at 10.235.24.154:30000. * fix(primus): use primus-mi355x-flat:v1 as base + install uv The previous base (ghcr.io/amdpilot-org/primus-mi355x-ready:v1) was a slimmed copy that did NOT include /workspace/primus_train/Primus. Phase 1 spent its full max_turns budget trying to bootstrap from scratch. Switch to primus-mi355x-flat:v1 (locally available, has Primus + Primus-Turbo pre-installed and patched for triton 3.4.0). Also install uv at /root/.local/bin in our Dockerfile so the kimi-cli runtime's source $HOME/.local/bin/env succeeds — without this, executor trials exit 137 immediately after image switch. Tag the new image primus-qwen3-30b-mfu-base:v1 so amdpilot triggers a build the first time it runs. * fix(primus): add gfx950 env vars + runtime IFNAME detection Two cumulative fixes for n08-09 (8x MI355X): 1. Without PYTORCH_ROCM_ARCH=gfx950 (and AITER_ROCM_ARCH, HSA_NO_SCRATCH_RECLAIM, HIP_FORCE_DEV_KERNARG, etc) the torch HIP runtime can't dispatch kernels and benchmarks die immediately with hipErrorInvalidDeviceFunction. These were set in xiao/baizhou's working containers but missing from amdpilot's docker run line. Add to both Dockerfile ENV and task.yaml container.env so they apply via either path. 2. Replace /workspace/detect_interface.sh with a /proc-based detector (the original needed `ip` from iproute2, unavailable in the slim base). bench_mfu.sh now auto-detects GLOO/NCCL socket IFNAME at runtime if bench_config.env doesn't pin one — without this, Megatron's distributed init fails fast (3-5s) before training starts. * feat(primus): add phase1_publish config to task.yaml Enables tag + push of the Phase 1 baseline image to docker.io/jhinpan/primus-qwen3-30b-mfu-phase1 after a successful phase1 commit. Template: {date}-{metric} (e.g. 20260422-278p80) + :latest. Other nodes can then `docker pull` that tag and skip Phase 1 entirely. * feat(primus): enable phase1_publish to ghcr.io/amdpilot-org Switch repository from docker.io/jhinpan to ghcr.io/amdpilot-org so all nodes in the org can pull the verified phase1-baseline image directly. Bump push timeout_s to 9000 (2.5h) to absorb the one-time 42 GB base-layer seed; subsequent pushes only upload the phase1-commit delta (~500 MB - 1 GB) via GHCR cross-repo mount. * feat(flydsl): sglang-kimi-k26-flydsl-mi355x eval instance Adds a hand-authored eval instance that targets the FlyDSL fused-MoE port from AMD's Kimi-K2.5 MI300X optimization blog to Kimi-K2.6 on 4x MI355X, on top of the baseline established by sgl-project/sglang#23381. Contents: - task.yaml : optimize-type task, phase1_baseline=true, frontier_model=true, TP=4, 4xMI355X, FlyDSL env OFF at start so Phase 1 reproduces PR #23381 numbers (DSL2_ROOT + MLIR_PATH + CK_TILE_FLOAT_TO_BFLOAT16 pre-set so import works as soon as agent flips gates during a trial) - Dockerfile : layers FlyDSL + AITER dev/kimi-K2.5 on top of jhinpan/sglang-k26-mi355x:v0.5.10rc0-rocm720-20260420 - bench_flydsl_k26.sh : single bench script emitting one canonical line: output_throughput_tok_s: <v> | concurrency=40 in=10240 out=512 decode_bs1_in8k=<guard> - task_description.md : full phased plan A-G, env-var reference table, BS=1 decode guard (>= 0.98x PR #23381 ~38.05 tok/s), front-loaded supervisor-visible gating block - test_harness.py : GSM8K accuracy gate (lm_eval limit=50 >= 0.90) - metadata.json : links to GH issue amdpilot-org/sglang#2 + PR #23381 - ISSUE.md : condensed body as filed on amdpilot-org/sglang#2 Tracks GitHub issue amdpilot-org/sglang#2. * fix(evals): align flydsl eval metadata and runtime env Keep FlyDSL's runtime MLIR path consistent with the Dockerfile and make new eval metadata consumable by the registry tooling. Mirror the Primus no-phase1 runtime env so the control variant does not hit known HIP temp path failures.

# Conflicts: # python/sglang/srt/environ.py

jhinpan · 2026-06-09T15:18:33Z

/tag-run-ci-label

jhinpan · 2026-06-09T15:19:31Z

/rerun-failed-ci

jhinpan · 2026-06-09T15:21:18Z

/tag-and-rerun-ci extra

ch-wan · 2026-06-09T20:23:45Z

This is a simple PR. I have removed the run-ci-extra label to save CI resources.

jhinpan requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, Ying1123, ch-wan, fzyzcjy, ispobock and merrymercy as code owners April 21, 2026 14:34

github-actions Bot added the deepseek label Apr 21, 2026

jhinpan and others added 4 commits April 21, 2026 10:58

Merge branch 'main' into feat/kimi-k26-mi355x-moe-tuning

55dcad8

Merge branch 'main' into feat/kimi-k26-mi355x-moe-tuning

3015417

This was referenced Apr 22, 2026

[AMD][kernel][Tier 5] FlyDSL fused-MoE optimization for Kimi-K2.6 on MI355X (kernel-level) amdpilot-org/sglang#2

Open

feat(flydsl): sglang-kimi-k26-flydsl-mi355x eval instance amdpilot-org/amdpilot-evals#16

Merged

This was referenced Apr 29, 2026

[CI Monitor] Daily Report - 2026-04-29 bingxche/sglang-ci-bot#54

Open

[CI Monitor] Daily Report - 2026-04-30 bingxche/sglang-ci-bot#55

Open

jhinpan added 3 commits June 9, 2026 15:09

Merge remote-tracking branch 'origin/main' into review-pr-23381

2992406

# Conflicts: # python/sglang/srt/environ.py

docs: document DeepSeek load worker env

bdfaecf

lint: clean touched files

3d974df

jhinpan requested review from JustinTong0323, wisclmy0611 and zijiexia as code owners June 9, 2026 15:13

github-actions Bot added the documentation Improvements or additions to documentation label Jun 9, 2026

docs: keep env docs in docs_new only

6495580

github-actions Bot added the run-ci label Jun 9, 2026

github-actions Bot added the run-ci-extra label Jun 9, 2026

ch-wan removed the run-ci-extra label Jun 9, 2026

This was referenced Jun 10, 2026

[CI Monitor] Daily Report - 2026-06-10 bingxche/sglang-ci-bot#99

Open

[CI Monitor] Daily Report - 2026-06-13 bingxche/sglang-ci-bot#103

Open

[CI Monitor] Daily Report - 2026-06-17 bingxche/sglang-ci-bot#107

Open

amd-bot mentioned this pull request Jun 28, 2026

[CI Monitor] Daily Report - 2026-06-28 bingxche/sglang-ci-bot#118

Open

This was referenced Jul 5, 2026

[CI Monitor] Daily Report - 2026-07-05 bingxche/sglang-ci-bot#126

Open

[CI Monitor] Daily Report - 2026-07-06 bingxche/sglang-ci-bot#127

Open

michaelzhang-ai mentioned this pull request Jul 9, 2026

[DO NOT MERGE][AMD][DI][CI] Kimi-K2.6 disagg diagnostic recipes for non-MTP GSM8K drop #30433

Closed

3 tasks

Navigation Menu

Uh oh!

[AMD] Add MI355X Kimi-K2.6 tuning artifacts - #23381

[AMD] Add MI355X Kimi-K2.6 tuning artifacts#23381
jhinpan wants to merge 9 commits into
sgl-project:mainfrom
jhinpan:feat/kimi-k26-mi355x-moe-tuning

jhinpan commented Apr 21, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Apr 21, 2026

Uh oh!

jhinpan commented Apr 22, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

ch-wan commented Jun 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

jhinpan commented Apr 21, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Benchmark Evidence

Recommended Runtime Setting

Test Plan

CI States

Uh oh!

gemini-code-assist Bot commented Apr 21, 2026

Uh oh!

jhinpan commented Apr 22, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

jhinpan commented Jun 9, 2026

Uh oh!

ch-wan commented Jun 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jhinpan commented Apr 21, 2026 •

edited by github-actions Bot

Loading