diff --git a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx index c9a88c033..6025fa67f 100644 --- a/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx +++ b/docs/cookbook/autoregressive/Moonshotai/Kimi-K3.mdx @@ -30,7 +30,7 @@ Then run the **Python** output of the command panel below in that environment. ```bash Command docker pull lmsysorg/sglang:latest # NVIDIA (CUDA) -docker pull lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910 # AMD MI350X / MI355X (ROCm) +docker pull lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916 # AMD MI350X / MI355X (ROCm) ``` For how to launch the image, see [Install → Method 3: Using Docker](../../../docs/get-started/install#method-3-using-docker). Substitute the inner `sglang serve ...` with what the command generator below produces. @@ -65,7 +65,7 @@ Pick your hardware, then the deployment shape and operating point. Node count fo **Strategy** — the operating point within that shape: - **Low-Latency** — no DCP, so the MLA KV stays TP-replicated. For chat. B200 splits its two nodes into PP2 × TP8; every other platform is flat TP. -- **Balanced** — the accuracy-preserving default: PP2 × DCPEP8 on B200 (the two pipeline stages and DCP8 split KV and KDA state), TP16/DCP16 on GB200, TP8/DCP8 on B300/GB300, TP8 ROCm/AITER on MI35x. +- **Balanced** — the accuracy-preserving default: PP2 × DCPEP8 on B200 (the two pipeline stages and DCP8 split KV and KDA state), TP16/DCP16 on GB200, TP8/DCP8 on B300/GB300, TP8/DCP8 ROCm/AITER on MI35x. - **High-Throughput** — the large-scale lane: pick a **Cluster Size** and **Large-Scale Preset** in the Playground ([details](#large-scale-presets)). The cell itself is Balanced, except on H100 (plus `extra_buffer_lazy`) and H200 (widens to 4×8 TP32/EP32 at `--mem-fraction-static 0.90`). `Long-Context` appears only under the `Prefill` PD mode; for long-context unified serving on B200, start from High-Throughput and raise `--context-length`. @@ -93,6 +93,22 @@ import { KimiK3MambaRatioCalculator } from "/src/snippets/_kimi_k3_mamba_ratio_c NVFP4 NOSPEC / NVFP4 DSPARK), which is why no point past concurrency 64 is published for Balanced. +### AMD AITER with DCP8 + +The MI350X/MI355X unified Balanced recipe uses TP8/DCP8 with AITER prefill and +decode attention. DCP shards the target MLA KV cache; RadixArk DSPARK's draft KV +remains replicated. The pinned `v0.5.19-rocm720-mi35x-20260916` image records +SGLang revision `e7f7447333`, which includes +[AITER DCP support (#34432)](https://github.com/sgl-project/sglang/pull/34432) and +the [DCP KV-free fix (#38941)](https://github.com/sgl-project/sglang/pull/38941). +No source overlay is required for DCP. + +Keep `SGLANG_K3_KDA_FUSED_BACKEND` unset with this image. The separate fused-KDA +opt-in requires the [deferred-gate fix (#39066)](https://github.com/sgl-project/sglang/pull/39066), +which is not included in this image. This updated recipe remains **Final +Verification In Progress**; the recorded speed numbers use their original +configurations and do not validate the new image or DCP8 recipe. + ### Mamba ratio calculator @@ -180,11 +196,11 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request | GB200 4×4 | TP16/DCP16 | MNNVL auto-detected | | H200 2×8 (4×8 on Unified High-Throughput) | TP16/EP16 + symm-mem, Marlin + FlashMLA; High-Throughput widens to TP32/EP32 over 4 nodes at mem-frac 0.90 with `extra_buffer_lazy` | same block on every node; export the cross-node NIC (`GLOO_SOCKET_IFNAME` / `NCCL_SOCKET_IFNAME`, `SGLANG_HOST_IP`); keep `NCCL_MNNVL_ENABLE=1 NCCL_CUMEM_ENABLE=1` | | H100 4×8 | TP32/EP32, Marlin + FlashMLA | SM90a build of the K3 image; pin NCCL/Gloo to the same NIC on all nodes; least post-weight headroom (80 GB) | -| MI350X/MI355X 1×8 | TP8 ROCm/AITER | AITER A8W4 FlyDSL MoE, Triton attention (`SGLANG_MLA_DECODE_TUNE=1` for gfx950 MLA decode geometry), graph bs up to 256, fp8 kvcache; DSPARK supported. Activation-quant and fused-KDA-decode knobs: [AMD ROCm/AITER environment](#amd-env) | +| MI350X/MI355X 1×8 | TP8/DCP8 ROCm/AITER (Unified Balanced) | AITER A8W4 FlyDSL MoE, AITER prefill/decode attention with sharded target MLA KV, graph bs up to 256, fp8 kvcache; DSPARK supported. Activation-quant and fused-KDA-decode knobs: [AMD ROCm/AITER environment](#amd-env) | | Ascend A3 Series 4×8 (32 cards / 64 dies) | TP64/DP4 + DeepEP | PD-mixed `Unified` only; DSPARK baked in; pin `GLOO`/`HCCL_SOCKET_IFNAME` on every node | | Ascend 950PR/DT Series 4×8 | TP32/dp1 + DeepEP | PD-mixed `Unified` only; DSPARK baked in; shared experts / dense MLP shard over attention-TP (`--shared-experts-tp-size 4`); radix cache off; pin `GLOO`/`HCCL_SOCKET_IFNAME` on every node | -**DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles: +**Blackwell DCP notes** — the DCP cells are Balanced and High-Throughput on every Blackwell platform, in both the `Unified` and `Decode` roles: - DCP is the only axis that shards the TP-replicated MLA KV; Low-Latency skips it. - Leave `--dcp-comm-backend` unset (fabric-resolved: `fi_a2a` on GB200/GB300, `a2a` on B200/B300). diff --git a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx index 87f803d7c..00daa4ce3 100644 --- a/docs/src/snippets/configs/moonshotai/kimi-k3.jsx +++ b/docs/src/snippets/configs/moonshotai/kimi-k3.jsx @@ -478,8 +478,8 @@ export const config = { gb200: "lmsysorg/sglang:kimi-k3", // 20260903 or newer: the AITER SiTU A4W4/A8W4 layout fix (sgl-project/sglang#33838, // merged Sep 3) and the fused gfx950 KDA decode boundary (#34198) first ship here. - mi350x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910", - mi355x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260910", + mi350x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916", + mi355x: "lmsysorg/sglang-rocm:v0.5.19-rocm720-mi35x-20260916", // NVFP4 needs a build with sgl-project/sglang#35077; the purpose-built dev // image is cut from that PR's head (CUDA 13). "b300|nvfp4": "lmsysorg/sglang:dev-dev-kimi-k3-nvfp4", @@ -1217,7 +1217,7 @@ export const config = { ], }, { - // MI350X and MI355X use the same single-node TP8 ROCm/AITER profile. + // MI350X and MI355X use the same single-node TP8/DCP8 ROCm/AITER profile. match: { hw: "mi350x", pdMode: "unified", strategy: "balanced" }, nnodes: 1, verified: false, @@ -1232,7 +1232,10 @@ export const config = { "--model-path {{MODEL_NAME}}", "--trust-remote-code", "--tp-size 8", - "--attention-backend triton", + "--dcp-size 8", + "--dcp-comm-backend a2a", + "--prefill-attention-backend aiter", + "--decode-attention-backend aiter", "--kv-cache-dtype fp8_e4m3", "--dtype bfloat16", "--mem-fraction-static 0.85", @@ -1259,7 +1262,10 @@ export const config = { "--model-path {{MODEL_NAME}}", "--trust-remote-code", "--tp-size 8", - "--attention-backend triton", + "--dcp-size 8", + "--dcp-comm-backend a2a", + "--prefill-attention-backend aiter", + "--decode-attention-backend aiter", "--kv-cache-dtype fp8_e4m3", "--dtype bfloat16", "--mem-fraction-static 0.85", diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 9d43cfdf5..383a6f478 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -2029,7 +2029,8 @@ class KimiK3DeltaAttention(nn.Module): qkv, g_proj_states, f_a, beta, _pad = torch.split( fused_states, self._qkvgbfa_sizes, dim=-1 ) - forget_gate = gemm(f_a, self._bfa_f_b_w) + # Fused KDA decode consumes f_a and applies f_b itself. + forget_gate = f_a if defer_f_b else gemm(f_a, self._bfa_f_b_w) return qkv, beta, forget_gate, g_proj_states if (