Skip to content

[Kimi-K3][LMCache] LMCache offload for Kimi-K3 #7598

[Kimi-K3][LMCache] LMCache offload for Kimi-K3

[Kimi-K3][LMCache] LMCache offload for Kimi-K3 #7598

name: ATOM SGLang Test
on:
pull_request:
branches: [main] # Triggers on PRs targeting `main`
types: [opened, synchronize, reopened, ready_for_review, labeled, closed]
paths-ignore:
- '**/*.md'
- 'docs/**'
- 'atom/plugin/vllm/**'
- 'atom/mesh/**'
- '.github/benchmark/models_atomesh.yaml'
- '.github/scripts/atomesh/**'
- '.github/workflows/atomesh-*.yaml'
- '.github/workflows/atom-vllm-*.yaml'
- '.github/benchmark/oot_models_accuracy.json'
- 'LICENSE'
- '.gitignore'
pull_request_review:
types: [submitted]
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: ${{ github.event_name == 'pull_request' || (github.event_name == 'pull_request_review' && github.event.review.state == 'approved') }}
env:
ATOM_BASE_NIGHTLY_IMAGE: rocm/atom-dev:latest
ATOM_PYTHON_TAG: "cp312"
GITHUB_REPO_URL: ${{ github.event.pull_request.head.repo.clone_url || 'https://github.com/ROCm/ATOM.git' }}
GITHUB_COMMIT_SHA: ${{ github.event.pull_request.head.sha || github.event.head_commit.id || github.sha }}
NIGHTLY_PLUGIN_IMAGE_TAG: rocm/atom-dev:sglang-latest
PLUGIN_CI_SLURM_ACCOUNT: amd-aifw-dev
PLUGIN_CI_SLURM_PARTITION: amd-spur
PLUGIN_CI_LOG_ROOT: /home/junyyang/ATOM_PLUGIN_RUNNER/LOG/
jobs:
ci-gate:
name: Check Heavy CI Gate
runs-on: ubuntu-latest
permissions:
contents: read
pull-requests: read
outputs:
should_run: ${{ steps.gate.outputs.should_run }}
reason: ${{ steps.gate.outputs.reason }}
steps:
- name: Checkout ATOM repo
uses: actions/checkout@v6
- name: Check approval and label gate
id: gate
env:
CI_GATE_LABELS: ci:full,ci:sglang
CI_GATE_PATHS_IGNORE: |
**/*.md
docs/**
atom/plugin/vllm/**
atom/mesh/**
.github/workflows/atomesh-*.yaml
.github/scripts/atomesh_*.sh
.github/dashboard/atomesh_*.html
.github/workflows/atom-vllm-*.yaml
.github/benchmark/oot_models_accuracy.json
LICENSE
.gitignore
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: bash ./.github/scripts/check_heavy_ci_gate.sh
check-signal:
if: ${{ needs.ci-gate.outputs.should_run == 'true' }}
needs: [ci-gate]
name: Check Pre Checkin Signal
runs-on: ubuntu-latest
permissions:
actions: read
contents: read
steps:
- name: Checkout ATOM repo
uses: actions/checkout@v6
- name: Wait for Pre Checkin workflow
run: bash ./.github/scripts/check_signal.sh
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_SHA: ${{ github.sha }}
CHECK_SIGNAL_EVENT_NAME: pull_request
CHECK_SIGNAL_HEAD_REF: ${{ github.event.pull_request.head.ref }}
CHECK_SIGNAL_SHA: ${{ github.event.pull_request.head.sha }}
download_aiter_wheel:
if: ${{ needs.ci-gate.outputs.should_run == 'true' && needs.check-signal.result == 'success' }}
needs: [ci-gate, check-signal]
name: Download aiter wheel
runs-on: ubuntu-latest
outputs:
aiter_artifact_id: ${{ steps.download.outputs.aiter_artifact_id }}
aiter_wheel_name: ${{ steps.download.outputs.aiter_wheel_name }}
steps:
- name: Checkout code
uses: actions/checkout@v6
- name: Prefer latest main aiter wheel manifest and fallback to artifact
id: download
timeout-minutes: 30
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: bash .github/scripts/download_aiter_wheel.sh
- name: Upload aiter wheel
uses: actions/upload-artifact@v7
with:
name: aiter-whl
path: aiter-whl/amd_aiter*.whl
retention-days: 7
atom-sglang-accuracy:
needs: [ci-gate, check-signal, download_aiter_wheel]
if: ${{ needs.ci-gate.outputs.should_run == 'true' && needs.check-signal.result == 'success' && needs.download_aiter_wheel.result == 'success' }}
name: Accuracy (${{ matrix.model_name }})
runs-on: ${{ matrix.runner }}
timeout-minutes: 240
permissions:
actions: read
contents: read
strategy:
fail-fast: false
max-parallel: 2
matrix:
include:
- model_name: "GLM-5.1-FP8 TP8"
model_path: "zai-org/GLM-5.1-FP8"
extra_args: "--trust-remote-code --tensor-parallel-size 8 --attention-backend aiter --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.85 --page-size 1 --disable-radix-cache"
env_vars: |
SGLANG_AITER_FP8_PREFILL_ATTN=0
SGLANG_USE_AITER=1
ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
lm_eval_num_fewshot: 5
accuracy_test_threshold: 0.93
runner: atom-mi355-8gpu-vllm-sgl-ci
- model_name: "DeepSeek-R1-FP4-V2 TP8 MTP3"
model_path: "amd/DeepSeek-R1-0528-MXFP4-v2"
extra_args: "--trust-remote-code --tensor-parallel-size 8 --attention-backend aiter --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.85 --page-size 1 --disable-radix-cache --speculative-draft-model-path SGLang/DeepSeek-R1-NextN --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --max-running-requests 256 --cuda-graph-bs 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 160 192 224 256"
env_vars: |
SGLANG_DEFAULT_SERVER_ARGS=
ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1
SGLANG_AITER_FP8_PREFILL_ATTN=0
SGLANG_USE_AITER=1
SGLANG_ENABLE_SPEC_V2=1
SGLANG_ENABLE_TORCH_COMPILE=1
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
TORCHINDUCTOR_COMPILE_THREADS=128
accuracy_test_threshold: 0.93
runner: atom-mi355-8gpu-vllm-sgl-ci
- model_name: "MiniMax-M3-MXFP4 TP4"
model_path: "amd/MiniMax-M3-MXFP4"
extra_args: "--trust-remote-code --tensor-parallel-size 4 --attention-backend aiter --page-size 128 --mem-fraction-static 0.8 --max-running-requests 128 --disable-radix-cache"
env_vars: |
SGLANG_DEFAULT_SERVER_ARGS=
AITER_QUICK_REDUCE_QUANTIZATION=INT4
SGLANG_USE_AITER=1
ATOM_FORCE_ATTN_TRITON=1
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE=atom.plugin.sglang.models
lm_eval_num_fewshot: 5
lm_eval_num_concurrent: 32
lm_eval_use_chat_completions: "1"
lm_eval_extra_model_args: "max_gen_toks=16384"
accuracy_test_threshold: 0.93
runner: atom-mi355-8gpu-vllm-sgl-ci
- model_name: "Qwen3.5-35B-A3B-FP8 TP2"
model_path: "Qwen/Qwen3.5-35B-A3B-FP8"
extra_args: "--tensor-parallel-size 2 --mem-fraction-static 0.9 --reasoning-parser qwen3 --disable-radix-cache"
env_vars: |
SGLANG_DEFAULT_SERVER_ARGS=
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
ATOM_ENABLE_QK_NORM_ROPE_CACHE_QUANT_FUSION=0
accuracy_test_threshold: 0.76
runner: atom-mi355-8gpu-vllm-sgl-ci
- model_name: "DeepSeek-V4-Flash TP8"
model_path: "deepseek-ai/DeepSeek-V4-Flash"
extra_args: "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.9 --swa-full-tokens-ratio 0.1 --max-running-requests 256 --page-size 256 --disable-radix-cache --disable-shared-experts-fusion --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4"
env_vars: |
SGLANG_DEFAULT_SERVER_ARGS=
AITER_BF16_FP8_MOE_BOUND=0
ATOM_MOE_GU_ITLV=1
SGLANG_DEFAULT_THINKING=1
SGLANG_DSV4_REASONING_EFFORT=max
SGLANG_USE_AITER=1
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
TORCHINDUCTOR_COMPILE_THREADS=128
lm_eval_num_fewshot: 5
lm_eval_num_concurrent: 8
accuracy_test_threshold: 0.94
runner: atom-mi355-8gpu-vllm-sgl-ci
- model_name: "DeepSeek-V4-Pro MTP3 TP8"
model_path: "deepseek-ai/DeepSeek-V4-Pro"
extra_args: "--trust-remote-code --tensor-parallel-size 8 --kv-cache-dtype fp8_e4m3 --mem-fraction-static 0.9 --swa-full-tokens-ratio 0.1 --max-running-requests 256 --page-size 256 --disable-radix-cache --disable-shared-experts-fusion --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4"
env_vars: |
SGLANG_DEFAULT_SERVER_ARGS=
AITER_BF16_FP8_MOE_BOUND=0
ATOM_MOE_GU_ITLV=1
SGLANG_DEFAULT_THINKING=1
SGLANG_DSV4_REASONING_EFFORT=max
SGLANG_USE_AITER=1
SGLANG_DSV4_FP4_EXPERTS=true
SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
TORCHINDUCTOR_COMPILE_THREADS=128
ATOM_USE_TRITON_MOE=0
ATOM_SGLANG_V4_DISABLE_PRO_DRAFT_TRITON_MOE=1
ATOM_SGLANG_V4_ENABLE_DRAFT_EXTEND_CG=1
lm_eval_num_fewshot: 5
lm_eval_num_concurrent: 8
accuracy_test_threshold: 0.94
runner: atom-mi355-8gpu-vllm-sgl-ci
env:
AITER_ARTIFACT_ID: ${{ needs.download_aiter_wheel.outputs.aiter_artifact_id }}
CONTAINER_NAME: plugin_ci_sglang_${{ github.run_id }}_${{ strategy.job-index }}
PLUGIN_CI_SLURM_NODES: crsuse2-m2m-v2-002,crsuse2-m2m-v2-027
RESULT_DIR: plugin-ci-results
steps:
- name: Checkout ATOM repo
uses: actions/checkout@v6
- name: Download aiter wheel
uses: actions/download-artifact@v8
with:
name: aiter-whl
path: aiter-whl
- name: Build SGLang Slurm cell
id: cell
env:
MATRIX_JSON: ${{ toJson(matrix) }}
PR_BASE_SHA: ${{ github.event.pull_request.base.sha }}
PR_HEAD_SHA: ${{ github.event.pull_request.head.sha }}
PLUGIN_CI_SLURM_SUBMIT_RUNNER: ${{ matrix.runner }}
SKIP_DOCKER_LOGIN: ${{ github.event.pull_request.head.repo.fork && '1' || '0' }}
run: |
set -euo pipefail
python3 .github/scripts/plugin_ci/build_cell_json.py \
--plugin sglang \
--matrix-json "${MATRIX_JSON}" \
--github-output
- name: Submit SGLang accuracy job to Slurm
env:
CELL_JSON: ${{ steps.cell.outputs.cell_json }}
DOCKER_USERNAME: ${{ secrets.DOCKER_USERNAME }}
DOCKER_PASSWORD: ${{ secrets.DOCKER_PASSWORD }}
HF_TOKEN: ${{ secrets.AMD_HF_TOKEN }}
run: |
set -euo pipefail
chmod +x .github/scripts/plugin_ci/*.sh
bash .github/scripts/plugin_ci/submit.sh \
--plugin sglang \
--cell-json "${CELL_JSON}" \
--result-dir "${RESULT_DIR}"
- name: Add SGLang result to job summary
if: always()
env:
CELL_JSON: ${{ steps.cell.outputs.cell_json }}
run: |
python3 - <<'PY'
import json
import os
from pathlib import Path
cell = json.loads(os.environ.get("CELL_JSON") or "{}")
result_root = Path(os.environ["RESULT_DIR"])
job_ids = list(result_root.glob("*.slurm-job-id"))
outputs = list(result_root.rglob("atom_sglang_accuracy_output.txt"))
with Path(os.environ["GITHUB_STEP_SUMMARY"]).open("a", encoding="utf-8") as summary:
summary.write(f"## {cell.get('id', 'SGLang plugin CI')}\n\n")
if job_ids:
summary.write(f"Slurm job: `{job_ids[0].read_text(encoding='utf-8').strip()}`\n\n")
if outputs:
text = outputs[0].read_text(encoding="utf-8", errors="replace")
lines = text.splitlines()
table = []
capture = False
for line in lines:
if "|Tasks|Version|" in line:
capture = True
if capture:
if not line.strip() and table:
break
table.append(line)
if table:
summary.write("\n".join(table) + "\n")
PY
- name: Upload SGLang Slurm artifacts
if: always()
uses: actions/upload-artifact@v7
with:
name: sglang-${{ matrix.model_name }}-artifacts
path: plugin-ci-results/
if-no-files-found: ignore