From afbaaeb72ce7cc2f9552b042cddd40454a79c70d Mon Sep 17 00:00:00 2001 From: Jasen2201 Date: Mon, 13 Jul 2026 06:38:55 +0000 Subject: [PATCH 1/3] Add Kimi-K2.5 MXFP4 PD disaggregation recipe for atomesh (1P+1D) Co-Authored-By: Claude --- recipes/mesh/kimi-k2.md | 169 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 recipes/mesh/kimi-k2.md diff --git a/recipes/mesh/kimi-k2.md b/recipes/mesh/kimi-k2.md new file mode 100644 index 0000000000..ae9ebf980b --- /dev/null +++ b/recipes/mesh/kimi-k2.md @@ -0,0 +1,169 @@ +# Kimi-K2.5 PD Disaggregation with atomesh + +PD-disaggregated serving for Kimi-K2.5 (MXFP4, text-only backbone) using the +ATOM native backend, Mooncake RDMA KV transfer, and atomesh routing. Single-node +1P+1D topology. + +## Prerequisites + +- AMD MI355X GPUs (8 GPUs per instance, TP=4) +- RDMA network connectivity (RoCE or InfiniBand) for KV cache transfer +- Model weights accessible at the same path on all nodes +- Checkpoint: [`amd/Kimi-K2.5-MXFP4`](https://huggingface.co/amd/Kimi-K2.5-MXFP4) + +## Quick Reference + +| Topology | Nodes | GPUs | Prefill flags | Decode flags | Typical CONC | +|----------|------:|-----:|---------------|--------------|-------------| +| 1P+1D (single-node) | 1 | 8 | TP=4 | TP=4 | 1–256 | + +## Environment Setup + +Start a container with the RDMA-aware docker script: + +```bash +DOCKER_IMAGE=rocm/atom-dev:latest bash atom/mesh/scripts/docker_start.sh +docker exec -it atom_mesh bash +``` + +All commands below run **inside the container**. + +### Common Env Vars + +```bash +export NODE_IP=$(ip route get 1.1.1.1 | awk '/src/ {print $7; exit}') +export PYTHONUNBUFFERED=1 +export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export ATOM_HOST_IP=${NODE_IP} +export LD_LIBRARY_PATH=$(python3 -c "import sysconfig; print(sysconfig.get_path('purelib'))")/mooncake:/opt/rocm/lib:${LD_LIBRARY_PATH:-} +rm -rf /root/.cache/atom/* 2>/dev/null || true +``` + +## Online Quant Config Reference + +Kimi-K2.5 MXFP4 uses a **layer-level** online quant config targeting only the +MLA attention linear weights (attention projections are quantized to PTPC FP8 at +load time; MoE expert weights are already in MXFP4 format): + +```json +{ + "global_quant_config": "", + "layer_quant_config": { + "model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", + "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", + "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", + "model.layers.*.self_attn.o_proj": "ptpc_fp8" + } +} +``` + +--- + +## 1P+1D — Single-Node (MXFP4) + +Prefill on GPU 0-3, decode on GPU 4-7, router on port 8000. + +### Prefill Server (GPU 0-3) + +```bash +export HIP_VISIBLE_DEVICES=0,1,2,3 + +HSA_NO_SCRATCH_RECLAIM=1 python3 -m atom.entrypoints.openai_server \ + --model amd/Kimi-K2.5-MXFP4 \ + --host 0.0.0.0 --server-port 8010 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv_cache_dtype fp8 \ + --block-size 16 \ + --gpu-memory-utilization 0.9 \ + --max-model-len 32768 \ + --max-num-seqs 256 \ + --max-num-batched-tokens 32768 \ + --online_quant_config '{"global_quant_config": "", "layer_quant_config": {"model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", "model.layers.*.self_attn.o_proj": "ptpc_fp8"}}' \ + --kv-transfer-config '{"kv_role":"kv_producer","kv_connector":"mooncake","handshake_port":6301}' \ + --no-enable_prefix_caching \ + 2>&1 | tee prefill.log +``` + +### Decode Server (GPU 4-7) + +```bash +export HIP_VISIBLE_DEVICES=4,5,6,7 + +HSA_NO_SCRATCH_RECLAIM=1 python3 -m atom.entrypoints.openai_server \ + --model amd/Kimi-K2.5-MXFP4 \ + --host 0.0.0.0 --server-port 8020 \ + --trust-remote-code \ + --tensor-parallel-size 4 \ + --kv_cache_dtype fp8 \ + --block-size 16 \ + --gpu-memory-utilization 0.9 \ + --max-model-len 32768 \ + --max-num-seqs 256 \ + --max-num-batched-tokens 32768 \ + --online_quant_config '{"global_quant_config": "", "layer_quant_config": {"model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", "model.layers.*.self_attn.o_proj": "ptpc_fp8"}}' \ + --kv-transfer-config '{"kv_role":"kv_consumer","kv_connector":"mooncake","handshake_port":6301}' \ + --cudagraph-capture-sizes "[8,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256]" \ + --no-enable_prefix_caching \ + 2>&1 | tee decode.log +``` + +### Router + +```bash +atomesh launch \ + --host 0.0.0.0 --port 8000 \ + --pd-disaggregation \ + --prefill "http://${NODE_IP}:8010" \ + --decode "http://${NODE_IP}:8020" \ + --policy random \ + --backend atom \ + --log-level info \ + --disable-health-check \ + --disable-circuit-breaker \ + --prometheus-port 29100 +``` + +--- + +## Verify + +After the router is up, send a test request: + +```bash +curl -sS http://127.0.0.1:8000/v1/completions \ + -H 'Content-Type: application/json' \ + -d '{"model":"amd/Kimi-K2.5-MXFP4","prompt":"Hello","max_tokens":32,"temperature":0}' +``` + +## GSM8K Accuracy (via Router) + +```bash +lm_eval --model local-completions \ + --model_args "model=amd/Kimi-K2.5-MXFP4,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,trust_remote_code=True" \ + --tasks gsm8k \ + --num_fewshot 3 +``` + +## Serving Benchmark (via Router) + +```bash +ISL=8192 +OSL=1024 +CONC=16 + +python -m atom.benchmarks.benchmark_serving \ + --model=amd/Kimi-K2.5-MXFP4 \ + --backend=vllm \ + --base-url=http://127.0.0.1:8000 \ + --dataset-name=random \ + --random-input-len="${ISL}" \ + --random-output-len="${OSL}" \ + --random-range-ratio=0.8 \ + --num-prompts=$(( CONC * 10 )) \ + --max-concurrency="${CONC}" \ + --request-rate=inf \ + --ignore-eos \ + --save-result \ + --percentile-metrics="ttft,tpot,itl,e2el" +``` From b1e6ad0582e2d96d56d64e3d50e2cfdb6ed02fff Mon Sep 17 00:00:00 2001 From: Jasen2201 Date: Mon, 13 Jul 2026 06:42:26 +0000 Subject: [PATCH 2/3] Update Kimi-K2.5 recipe: remove online_quant_config, fix cudagraph sizes Co-Authored-By: Claude --- recipes/mesh/kimi-k2.md | 24 +----------------------- 1 file changed, 1 insertion(+), 23 deletions(-) diff --git a/recipes/mesh/kimi-k2.md b/recipes/mesh/kimi-k2.md index ae9ebf980b..0a26c9c8ac 100644 --- a/recipes/mesh/kimi-k2.md +++ b/recipes/mesh/kimi-k2.md @@ -39,26 +39,6 @@ export LD_LIBRARY_PATH=$(python3 -c "import sysconfig; print(sysconfig.get_path( rm -rf /root/.cache/atom/* 2>/dev/null || true ``` -## Online Quant Config Reference - -Kimi-K2.5 MXFP4 uses a **layer-level** online quant config targeting only the -MLA attention linear weights (attention projections are quantized to PTPC FP8 at -load time; MoE expert weights are already in MXFP4 format): - -```json -{ - "global_quant_config": "", - "layer_quant_config": { - "model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", - "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", - "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", - "model.layers.*.self_attn.o_proj": "ptpc_fp8" - } -} -``` - ---- - ## 1P+1D — Single-Node (MXFP4) Prefill on GPU 0-3, decode on GPU 4-7, router on port 8000. @@ -79,7 +59,6 @@ HSA_NO_SCRATCH_RECLAIM=1 python3 -m atom.entrypoints.openai_server \ --max-model-len 32768 \ --max-num-seqs 256 \ --max-num-batched-tokens 32768 \ - --online_quant_config '{"global_quant_config": "", "layer_quant_config": {"model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", "model.layers.*.self_attn.o_proj": "ptpc_fp8"}}' \ --kv-transfer-config '{"kv_role":"kv_producer","kv_connector":"mooncake","handshake_port":6301}' \ --no-enable_prefix_caching \ 2>&1 | tee prefill.log @@ -101,9 +80,8 @@ HSA_NO_SCRATCH_RECLAIM=1 python3 -m atom.entrypoints.openai_server \ --max-model-len 32768 \ --max-num-seqs 256 \ --max-num-batched-tokens 32768 \ - --online_quant_config '{"global_quant_config": "", "layer_quant_config": {"model.layers.*.self_attn.fused_qkv_a_proj": "ptpc_fp8", "model.layers.*.self_attn.q_b_proj": "ptpc_fp8", "model.layers.*.self_attn.kv_b_proj": "ptpc_fp8", "model.layers.*.self_attn.o_proj": "ptpc_fp8"}}' \ --kv-transfer-config '{"kv_role":"kv_consumer","kv_connector":"mooncake","handshake_port":6301}' \ - --cudagraph-capture-sizes "[8,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256]" \ + --cudagraph-capture-sizes "[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256]" \ --no-enable_prefix_caching \ 2>&1 | tee decode.log ``` From 66f7519d72cb27ce8fe2b955adf0b1d6ce9a2a11 Mon Sep 17 00:00:00 2001 From: Jasen2201 Date: Mon, 13 Jul 2026 12:36:56 +0000 Subject: [PATCH 3/3] Add AITER_MXFP4_INTERMEDIATE=1 env var to Kimi-K2.5 recipe From PR #1580: enable MXFP4 intermediate output for INT4 quick reduce. Co-Authored-By: Claude --- recipes/mesh/kimi-k2.md | 1 + 1 file changed, 1 insertion(+) diff --git a/recipes/mesh/kimi-k2.md b/recipes/mesh/kimi-k2.md index 0a26c9c8ac..0811fb3fd9 100644 --- a/recipes/mesh/kimi-k2.md +++ b/recipes/mesh/kimi-k2.md @@ -34,6 +34,7 @@ All commands below run **inside the container**. export NODE_IP=$(ip route get 1.1.1.1 | awk '/src/ {print $7; exit}') export PYTHONUNBUFFERED=1 export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +export AITER_MXFP4_INTERMEDIATE=1 export ATOM_HOST_IP=${NODE_IP} export LD_LIBRARY_PATH=$(python3 -c "import sysconfig; print(sysconfig.get_path('purelib'))")/mooncake:/opt/rocm/lib:${LD_LIBRARY_PATH:-} rm -rf /root/.cache/atom/* 2>/dev/null || true