[AMD] Update kimi-k3 amd cookbook 0903 (#37781)
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
@@ -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.17-rocm720-mi35x-20260817 # AMD MI350X / MI355X (ROCm)
|
||||
docker pull lmsysorg/sglang-rocm:v0.5.18-rocm720-mi35x-20260903 # 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.
|
||||
@@ -167,7 +167,7 @@ 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 |
|
||||
| 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) |
|
||||
| Atlas 800I A3 4×8 (32 cards / 64 dies) | TP64/DP4 + DeepEP | PD-mixed `Unified` only; DSPARK baked in; 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:
|
||||
@@ -181,6 +181,19 @@ Speculation: DSPARK holds block size + 1 (= 8) intermediate states per request
|
||||
|
||||
No cell has a serving round in this exact shape — treat them as starting points to verify.
|
||||
|
||||
<a id="amd-env" />
|
||||
|
||||
**AMD ROCm/AITER environment (MI350X / MI355X).** The MI35x cell emits `SGLANG_USE_AITER=1 SGLANG_AITER_K3_OPT=1 AITER_FLYDSL_FORCE=1 AITER_SITUV2_A8W4=1` — the first three turn on the AITER ROCm path, its K3-specific fused kernels, and the FlyDSL MoE kernels; the rest of this table is what you can change on top. Everything here is gfx950/ROCm-only and inert elsewhere. These knobs first ship in the `20260903` daily ROCm image pinned above — on an older image they are simply unread — and the two SiTU rows also need an AITER revision at or past [ROCm/aiter#4534](https://github.com/ROCm/aiter/pull/4534) (FlyDSL 0.3.0).
|
||||
|
||||
| Env var | Default | Effect |
|
||||
|---|---|---|
|
||||
| `AITER_SITUV2_A8W4=1` | unset | SiTU v2 MoE with A8W4 activation quantization, on AITER's GU-interleaved preshuffled weight layout. The performance default the cell ships. |
|
||||
| `AITER_SITUV2_A4W4=1` | unset | A4W4 instead, on the generic separated shuffle layout. Numerically correct but slower than A8W4 (530.8 vs 537.3 tok/s median output on 8×MI35x). Setting **both** gives A8W4 precedence — SGLang follows AITER and keeps the GU-interleaved layout. |
|
||||
| `SGLANG_K3_KDA_FUSED_BACKEND=aiter` | unset | Opt in to the fused ROCm KDA decode boundary: the `f_b` projection is deferred into the gfx950 FlyDSL kernel so decode fuses `f_b` + convolution + recurrent state update + gated RMSNorm. Any other value (or unset) keeps the unfused KDA path. |
|
||||
| `SGLANG_K3_FLYDSL_SOURCE` | `auto` | Which FlyDSL implementation backs the fused decode: `auto` prefers SGLang's vendored kernels and falls back to the AITER module, `sglang` / `aiter` pin one. Leave it alone unless you are bisecting the two. |
|
||||
|
||||
The fused KDA backend is fail-closed at two levels: it arms during model init only when the flag is exactly `aiter` **and** the gfx950 FlyDSL kernels are importable, and each decode step re-validates shapes, dtypes, strides, state indices, and output buffers before dispatch — anything unexpected falls back to the stock KDA implementation rather than erroring. Batch size 2 automatically picks a separately validated kernel schedule; every other batch keeps the original build options. Measured on a 69-layer graph the fused boundary is 9.20 → 8.38 µs/layer (−8.9%), with GSM8K 1319 at 0.950.
|
||||
|
||||
## 3. Advanced Usage
|
||||
|
||||
### 3.1 Reasoning
|
||||
|
||||
@@ -468,8 +468,10 @@ export const config = {
|
||||
gb300: "lmsysorg/sglang:kimi-k3",
|
||||
b200: "lmsysorg/sglang:kimi-k3",
|
||||
gb200: "lmsysorg/sglang:kimi-k3",
|
||||
mi350x: "lmsysorg/sglang-rocm:v0.5.17-rocm720-mi35x-20260817",
|
||||
mi355x: "lmsysorg/sglang-rocm:v0.5.17-rocm720-mi35x-20260817",
|
||||
// 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.18-rocm720-mi35x-20260903",
|
||||
mi355x: "lmsysorg/sglang-rocm:v0.5.18-rocm720-mi35x-20260903",
|
||||
// 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",
|
||||
@@ -843,6 +845,32 @@ export const config = {
|
||||
{ id: "on", label: "On", env: ["SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1"] },
|
||||
],
|
||||
},
|
||||
{
|
||||
// AITER SiTU v2 activation-quant mode (ROCm only). A8W4 rides AITER's
|
||||
// GU-interleaved preshuffled weights, A4W4 the generic separated layout;
|
||||
// SGLang mirrors AITER's precedence (A8W4 wins when both are set), so the
|
||||
// row emits exactly one. Needs AITER >= ROCm/aiter#4534.
|
||||
id: "situActMode", title: "SiTU MoE Activation Quant (AMD)",
|
||||
showWhen: (b) => ["mi350x", "mi355x"].includes(b.hw),
|
||||
stripEnv: ["AITER_SITUV2_A8W4", "AITER_SITUV2_A4W4"],
|
||||
options: [
|
||||
{ id: "a8w4", label: "A8W4 (default)", env: ["AITER_SITUV2_A8W4=1"] },
|
||||
{ id: "a4w4", label: "A4W4 — correct, ~1% slower", env: ["AITER_SITUV2_A4W4=1"] },
|
||||
],
|
||||
},
|
||||
{
|
||||
// Opt-in fused gfx950 KDA decode boundary (f_b + conv + recurrence +
|
||||
// gated RMSNorm). Fail-closed: init arms it only on gfx950 with the
|
||||
// FlyDSL kernels importable, and every decode step re-validates before
|
||||
// dispatch. Env var, not a flag, so it emits via env/stripEnv.
|
||||
id: "kdaFusedDecode", title: "Fused KDA Decode (AMD gfx950)",
|
||||
showWhen: (b) => ["mi350x", "mi355x"].includes(b.hw),
|
||||
stripEnv: ["SGLANG_K3_KDA_FUSED_BACKEND"],
|
||||
options: [
|
||||
{ id: "off", label: "Off" },
|
||||
{ id: "aiter", label: "On (AITER fused boundary)", env: ["SGLANG_K3_KDA_FUSED_BACKEND=aiter"] },
|
||||
],
|
||||
},
|
||||
{
|
||||
// Only meaningful with EP a2a on (MoE card or a large-scale preset).
|
||||
id: "eplb", title: "Expert Rebalancing (EPLB)",
|
||||
|
||||
Reference in New Issue
Block a user