diff --git a/scripts/vllm/kimik3_mi300x/README.md b/scripts/vllm/kimik3_mi300x/README.md new file mode 100644 index 00000000..5b2ffd36 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/README.md @@ -0,0 +1,72 @@ +# Kimi-K3 (MXFP4) serving on AMD Instinct MI300X (gfx942) + +[Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) is Moonshot AI's ~2.8T-parameter +Mixture-of-Experts model (natively MXFP4/QAT, hybrid MLA + Kimi-Delta-Attention, 896 +experts). These are **colocated** (single-instance) vLLM serving recipes for **MI300X / +gfx942**, complementing MAD's existing single-node gfx950/MI355X K3 recipe. + +Why MI300X needs its own recipe: MI300X has 192 GB/GPU, so the ~1.5 TB checkpoint does +**not** fit a single 8-GPU node under TP8. All recipes here shard across **2 nodes (16× +MI300X)** using PP2×TP8 (~102 GB/GPU per node). gfx942 also requires +`VLLM_ROCM_USE_AITER_MLA=0` (the AITER MLA kernel is gfx950-only). + +## Image & model + +- **Image:** `amdsiloai/vllm:kimi-k3-mi325x-release-v2` — a gfx942 K3 vLLM build. Public + and anonymously pullable (`docker pull amdsiloai/vllm:kimi-k3-mi325x-release-v2`); built + for MI325X, also runs MI300X. Override with `-e IMAGE=` if you have a newer tag. +- **Model:** `moonshotai/Kimi-K3` (MXFP4). Place on local NVMe for fast load; pass via `MODEL_DIR`. + +## The recipes + +| Folder | Parallelism | Expert all2all | MoE path | Use when | +|--------|-------------|----------------|----------|----------| +| [`pp2xtp8/`](pp2xtp8/) | PP2×TP8, no EP | — | a16w4 | Simplest baseline; lowest single-user latency. | +| [`wideep_int4_allgather/`](wideep_int4_allgather/) | PP2×TP8, EP8/node | `allgather_reducescatter` (generic) | a8w4 (`AITER_SITUV2_A8W4=1`) | Expert-parallel without MoRI kernels. | +| [`wideep_int4_moriep/`](wideep_int4_moriep/) | PP2×TP8, EP8/node | `mori_low_latency` (**MoRI-EP**) | a8w4 (`AITER_SITUV2_A8W4=1`) | MoRI-EP all2all expert dispatch (intra-node EP group). | +| [`wideep_disagg_2p2d/`](wideep_disagg_2p2d/) | 2P/2D disagg, TP2×DP8 per pool → EP16, no PP | `mori_low_latency` (**MoRI-EP**) + **MoRIIO** KV/state transfer | MXFP4 | Prefill/decode disaggregation across 4 nodes. Highest concurrent throughput; NIAH validated to 300K. | + +The first three are colocated (single-instance, no prefill/decode split); the +fourth splits prefill and decode across two pools. + +**`wideep_disagg_2p2d/` is validated** — single-needle NIAH passes deterministically +through **300K tokens** (all depths). Two root-cause fixes closed the decode-recall +bug (4-KV-group block routing + multi-chunk prefill transfer); see its +[`STATUS.md`](wideep_disagg_2p2d/STATUS.md) and +[`RESULTS.md`](wideep_disagg_2p2d/RESULTS.md). **Pick by workload:** the colocated +recipes give the lowest single-request latency (one request spans all 16 GPUs) — +best for interactive / low-QPS; disagg gives **5.7× throughput at concurrency 8** +(7.3× at 16) plus decode-latency isolation, at ~4× higher single-stream latency — +best for batch / high-QPS. + +> **EP scope:** the 896 experts split **8-way across each node's 8 GPUs** (112 experts/GPU → `[EP Rank x/8]`), and that EP8 group is replicated on each of the 2 pipeline stages. The expert all2all (incl. MoRI-EP) therefore runs **intra-node**; the only cross-node traffic is the PP activation hand-off, over NCCL. ("16" is the GPU count, not the EP width.) + +## Quick start + +Each recipe is self-contained (`run.sh` + `README.md` + `niah_probe.py`). Launch the +**worker (rank 1) first, then the head (rank 0)**: + +```bash +cd / +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=worker MASTER= bash run.sh # node 1 +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=head MASTER= bash run.sh # node 0 (API :8000) + +# verify +curl http://:8000/v1/models +python3 niah_probe.py --url http://:8000 --model kimi-k3 --ctx 8500 --depths 0.1,0.5,0.9 +``` + +See each recipe's `README.md` for its specific flags and env. + +## Common env (set in every run.sh) + +- `VLLM_ROCM_USE_AITER_MLA=0` — **required** on gfx942. +- `--trust-remote-code --reasoning-parser kimi_k3 --mm-encoder-tp-mode data + --safetensors-load-strategy prefetch`. +- ROCm 7.2.x: `HSA_ENABLE_IPC_MODE_LEGACY=0`, `PYTORCH_ALLOC_CONF` / + `PYTORCH_HIP_ALLOC_CONF=expandable_segments:False`. +- **RDMA fabric env is cluster-specific** (`NCCL_IB_HCA`, `NCCL_IB_GID_INDEX`, and the + `MORI_*` knobs in the moriep recipe) — adjust for your cluster's NICs/GIDs. + +Overridable per recipe: `IMAGE`, `MODEL_DIR`, `MASTER`, `PORT` (8000), `MAX_MODEL_LEN` +(10240), `MAX_NUM_SEQS` (8), `GPU_UTIL` (0.90). diff --git a/scripts/vllm/kimik3_mi300x/kimi_k3_moriep_report.html b/scripts/vllm/kimik3_mi300x/kimi_k3_moriep_report.html new file mode 100644 index 00000000..c825463d --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/kimi_k3_moriep_report.html @@ -0,0 +1,358 @@ + + + + + +Kimi-K3 MXFP4 — MoRI-EP Wide-EP Serving on Thor2 / MI325X + + + +
+ +
+
+

Kimi-K3 (MXFP4) — MoRI-EP Wide Expert-Parallel Serving

+

Colocated PP2 × TP8 = 16 GPUs (MoE runs EP8 within each node) · 2× nodes AMD Instinct MI325X (gfx942) · Broadcom Thor2 400G RoCE

+
+ +
+ +
+

Result in one line: Kimi-K3 (~2.8T-param MoE, 896 experts, MXFP4) served with MoRI-EP (mori_low_latency) expert all-to-all — 21/21 needle-in-a-haystack tests passed from 20K up to 800K tokens, with expert parallelism active (-ep- in fingerprint).

+
+ +
+

Topology note (read this): the GPU count is PP2 × TP8 = 16 GPUs (DP=1) — EP is not a separate multiplicative axis. Within each 8-GPU pipeline stage the same 8 GPUs are used two ways: dense/attention layers run TP8, and MoE layers run EP8 over those same 8 GPUs (896 experts ÷ 8 = 112/GPU). vLLM derives EP size = TP × DP = 8. So MoRI-EP all-to-all runs intra-node (across a node's 8 GPUs); the only cross-node traffic is the pipeline-stage hand-off, over NCCL, not MoRI-EP. This config does not exercise MoRI-EP cross-node scale-out (validated separately, see below). Both "EP16" and reading "PP2×TP8×EP8" as a product (which would imply 128 GPUs) are wrong — it's 16 GPUs, EP8 overlaying TP8 per node.

+
+ + +
+
Model
Kimi-K3
MXFP4 · ~2.8T MoE
+
Experts
896
112 / GPU · EP8 per node
+
Parallelism
PP2 × TP8
16 GPUs · EP8 overlays TP per node
+
All-to-all
MoRI-EP
mori_low_latency · intra-node
+
Max context tested
800K
786,713 tokens · PASS
+
NIAH accuracy
21/21
7 sizes × 3 depths
+
+ + +

Long-context retrieval — latency vs. context length

+
+
+
End-to-end latency per request single request
+
NIAH · needle retrieval · all depths PASS
+
+ + + +
Context in prompt tokens (log scale) · latency in seconds · every point retrieved the needle correctly at depths 0.1 / 0.5 / 0.9.
+
+ + +

Full NIAH results

+
+ + + + + + + + + +
TargetPrompt tokensDepth 0.1Depth 0.5Depth 0.9Latency
+
+ + +

How MoRI-EP was used for the 896-expert MoE (EP8, intra-node)

+
+

Kimi-K3 has 896 experts. The GPU count is PP2 × TP8 = 16 (DP=1). EP is not a third multiplicative factor — it reuses the tensor-parallel GPUs:

+
    +
  • PP2 (pipeline): the model's layers are split into 2 stages, one per node — for weight fit. Cross-node (this is the only inter-node axis).
    • +
  • TP8 (tensor): within a node, each dense/attention layer's math is split across the 8 GPUs. Intra-node.
    • +
  • EP8 (expert): for MoE layers, those same 8 GPUs instead hold 112 experts each (896 ÷ 8). vLLM sets EP size = TP × DP = 8. Intra-node — EP8 overlays TP8, it does not add GPUs.
    • +
+

The runtime confirms it: [EP Rank 0/8], Local/global experts 112/896 (896 ÷ 112 = 8), world_size=16. The 8-way expert split is replicated on each pipeline stage — node A holds experts 0–895 for stage-0 layers, node B holds experts 0–895 again for stage-1 layers. So a GPU is TP-rank-N for dense layers and EP-rank-N for MoE layers — same GPU, two roles.

+

So the MoRI-EP dispatch/combine all-to-all (--all2all-backend mori_low_latency) runs within each node's 8-GPU EP group. The only cross-node traffic is the PP activation hand-off (stage 0 → stage 1), which uses NCCL over the Thor2 rails — not MoRI-EP.

+
NODE A (PP stage 0) NODE B (PP stage 1) + ┌───────────────────────────────┐ ┌───────────────────────────────┐ + │ GPU0..GPU7 EP0..EP7 │ │ GPU0..GPU7 EP0..EP7 │ + │ 112 experts each (896 total) │ │ 112 experts each (896 total) │ + │ ▲ MoRI-EP all2all (intra-node)│ │ ▲ MoRI-EP all2all (intra-node)│ + └───────────────┬───────────────┘ └───────────────┬───────────────┘ + │ PP activations (NCCL over Thor2 RDMA) + └───────────────►───────────────────┘ + token ─► router (top-16 of 896) ─► MoRI-EP DISPATCH (within node) ─► experts + ─► MoRI-EP COMBINE (within node) ─► next layer
+

Fabric knobs (set so MoRI-EP can use the Thor2 rails; here the EP group is intra-node so they govern on-node dispatch): MORI_RDMA_DEVICES=rdma0..rdma7, MORI_IB_GID_INDEX=3 (RoCEv2), MORI_NUM_QP_PER_PE=8, MORI_IB_ENABLE_RELAXED_ORDERING=1, MORI_GPU_ARCHS=gfx942.

+
+ +
+

MoRI-EP cross-node scale-out — validated separately

+

Because this serve keeps EP intra-node, it does not put MoRI-EP's cross-node scale-out on the critical path. That path was validated independently on this cluster with the standalone MoRI internode test (test_dispatch_combine_internode.py): cross-node MoRI-EP dispatch/combine over Thor2 RDMA, 500/500 rounds, 0 errors, on both the v1 (VRAM-atomics) and async_ll kernels, across multiple node pairs. To make cross-node MoRI-EP the load-bearing path of a live serve, an EP-group-spans-nodes config is needed (e.g. DP-across-nodes wide-EP) rather than PP2-for-weight-fit with EP nested inside each node.

+
+ + +

Everything added over the ROCm/MAD PR recipe

+

The base recipe is ROCm/MAD PR #193 → wideep_int4_moriep (recipe 3 of 3). The PR ships Mellanox/mlx5 fabric defaults; this cluster is Broadcom Thor2. Below is the complete set of things that had to exist or change beyond the PR — split into fabric adaptation (in the run script) and the cluster groundwork the recipe silently assumes.

+ +
+
+

A · Fabric adaptation in run.sh

+

The only edits to the recipe itself — all cluster-specific env the PR README says to override:

+ + + + + + + + + + +
SettingPR defaultThis cluster
Boot / OOB NICeth0eno0
NCCL_IB_HCAmlx5_0,2..9rdma0..rdma7
MORI_RDMA_DEVICES—rdma0..rdma7
MORI_SOCKET_IFNAME—eno0
Model mountlocal NVMe-v /it-share-prj2-1 (WekaFS)
GID index33 (already matched)
+

Unchanged (kept exactly as PR): the whole vllm serve line — TP8 / PP2 / --enable-expert-parallel / --all2all-backend mori_low_latency / --trust-remote-code / --reasoning-parser kimi_k3 / --mm-encoder-tp-mode data / --safetensors-load-strategy prefetch — and every gfx942/MoRI/AITER knob (VLLM_ROCM_USE_AITER_MLA=0, AITER_SITUV2_A8W4=1, MORI_GPU_ARCHS=gfx942, MORI_NUM_QP_PER_PE=8).

+
+ +
+

B · Cluster groundwork prerequisites

+

The recipe silently assumes a healthy MoRI-capable cluster. Getting there took work the PR does not cover:

+
    +
  • NIC firmware/driver: brought all Thor2 NICs to fw 237.1.148.0 / bnxt_re 237.1.137.0 and confirmed MoRI-EP works natively on it (no 235 downgrade — the older PR#191 workaround does not apply to this newer Thor2).
    • +
  • MoRI-EP + MoRI-IO validation: both v1 (VRAM atomics) and async_ll EP kernels pass across the fabric; MoRI-IO ~48.5 GB/s; RDMA mesh 347–382 Gb/s on all pairs.
    • +
  • Docker daemon fix (worker node): a rogue snap dockerd was shadowing the system daemon with an AppArmor-confined data-root → caused the MoRI symmetric-memory illegal-access. Restored system docker (/var/lib/docker).
    • +
  • GPU recovery (g32): a GPU that fell off the PCI bus recovered via BMC hard reboot → back to 8/8.
    • +
  • Model staging: Kimi-K3-MXFP4 (1.5 TB, 96 shards) downloaded to WekaFS /it-share-prj2-1, readable ~1.5 GB/s/node, mounted into the serve container.
    • +
  • Context raise: relaunched at --max-model-len 900000 (recipe default 10240) to enable the long-context sweep.
    • +
+
+
+ + +

System & software stack

+
+
+

Hardware & fabric

+ + + + + + + + + + + +
Nodes2 (head + worker), 16 GPU
GPUAMD Instinct MI325X (gfx942)
GPU memory256 GB/GPU (~192 usable band)
NICBroadcom BCM57608 "Thor2" 400G
NIC firmware237.1.148.0
RDMA driverbnxt_re 237.1.137.0
TransportRoCE v2 · GID 3 · 8 rails
Shared FSWekaFS (/it-share-prj2-1)
+
+
+

Software

+ + + + + + + + + + + +
Imageamdsiloai/vllm:kimi-k3-mi325x-release-v2
vLLM0.1.dev19253+g5f76ae224
ROCm7.2.3
PyTorch2.11.0
QuantizationMXFP4 · a8w4 MoE (SiTU)
AttentionMLA + Kimi-Delta (DSA), AITER_MLA=0
KV cache8.4M tokens (9.33× @ 900K)
Model native max1,048,576 tokens
+
+
+ + +

Reproduce

+
+

Adapted recipe run.sh on the head node. Launch worker first, then head:

+
# worker (rank 1) on node B — FIRST +MODEL_DIR=/it-share-prj2-1/models/Kimi-K3-MXFP4 ROLE=worker \ + MASTER=<head eno0 IP> MAX_MODEL_LEN=900000 bash run_k3_thor2.sh + +# head (rank 0, serves API :8000) on node A +MODEL_DIR=/it-share-prj2-1/models/Kimi-K3-MXFP4 ROLE=head \ + MASTER=<head eno0 IP> MAX_MODEL_LEN=900000 bash run_k3_thor2.sh + +# verify — fingerprint contains -ep- when expert parallelism is live +curl http://<head IP>:8000/v1/models +python3 niah_sweep.py --url http://<head IP>:8000 --model kimi-k3 \ + --tokens 20000,100000,200000,500000,600000,750000,800000 --depths 0.1,0.5,0.9
+

Bring-up: ~18 min to load 1.5 TB (96 shards) + ~2 min AITER JIT compile (a8w4 MoE + MoRI-EP dispatch). AITER kernels cache after first run.

+
+ +
+ Generated on an MI325X + Thor2 cluster · Kimi-K3 MXFP4 wide-EP via MoRI-EP · base recipe ROCm/MAD PR #193 (wideep_int4_moriep). All figures measured on the live serve; NIAH needle HELIOTROPE-7492, greedy decode. +
+ +
+ + + + diff --git a/scripts/vllm/kimik3_mi300x/pp2xtp8/README.md b/scripts/vllm/kimik3_mi300x/pp2xtp8/README.md new file mode 100644 index 00000000..9ba570dd --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/pp2xtp8/README.md @@ -0,0 +1,45 @@ +# Kimi-K3 (MXFP4) on MI300X / gfx942 — PP2×TP8 baseline + +Serves Kimi-K3 (MXFP4) across **2 nodes** (16× MI300X): tensor-parallel 8 within each +node, pipeline-parallel 2 across nodes. Each node holds half the layers (~102 GB/GPU); +a single 8-GPU node cannot fit the model + KV. This is the simplest, lowest-latency K3 +serve on MI300X, with no expert parallelism. See `../wideep_int4_moriep` for wide-EP. + +- Image: `amdsiloai/vllm:kimi-k3-mi325x-release-v2` (gfx942 K3 vLLM build; public, anonymous pull) +- Model: `moonshotai/Kimi-K3` (MXFP4), on local NVMe or NFS +- Colocated (single instance; no prefill/decode disaggregation) + +## Run (worker FIRST, then head) + +```bash +# on the worker node (rank 1): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=worker MASTER= bash run.sh + +# on the head node (rank 0, serves the API on :8000): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=head MASTER= bash run.sh +``` + +- `MASTER` = the head node's **eth0** IP (NCCL/PP bootstrap). +- `MODEL_DIR` = path to the Kimi-K3-MXFP4 weights. Prefer **local NVMe** over NFS (much faster load). +- First start recompiles gfx942 AITER kernels (a few minutes), then cached. + +## Verify + +```bash +curl http://:8000/v1/models +python3 niah_probe.py --url http://:8000 --model kimi-k3 --ctx 8500 --depths 0.1,0.5,0.9 +``` + +## Key env (set in run.sh) + +| Var | Value | Why | +|-----|-------|-----| +| `VLLM_ROCM_USE_AITER_MLA` | `0` | Required on gfx942 — the AITER MLA kernel is gfx950-only and asserts at TP8. | +| `NCCL_SOCKET_IFNAME` / `GLOO_SOCKET_IFNAME` | `eth0` | Control-plane NIC. | +| `NCCL_IB_HCA` | `mlx5_0,2,3,4,5,7,8,9` | 8× RDMA NICs (**cluster-specific — override for your fabric**). | +| `NCCL_IB_GID_INDEX` | `3` | RoCE GID (**cluster-specific**). | +| `HSA_ENABLE_IPC_MODE_LEGACY` | `0` | ROCm 7.2.x IPC. | +| `PYTORCH_(HIP_)ALLOC_CONF` | `expandable_segments:False` | Required on ROCm 7.2.x. | + +Overridable env: `IMAGE`, `MODEL_DIR`, `MASTER`, `PORT` (8000), `MAX_MODEL_LEN` (10240), +`MAX_NUM_SEQS` (8), `GPU_UTIL` (0.90). diff --git a/scripts/vllm/kimik3_mi300x/pp2xtp8/niah_probe.py b/scripts/vllm/kimik3_mi300x/pp2xtp8/niah_probe.py new file mode 100644 index 00000000..0fed8766 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/pp2xtp8/niah_probe.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Lightweight NIAH (needle-in-a-haystack) + basic accuracy probe for a served +Kimi-K3 endpoint. No deps beyond stdlib (urllib). Run from any node that can +reach the head IP:port. + +Usage: + python3 niah_probe.py --url http://:8000 --model kimi-k3 \ + --ctx 8000 --depths 0.1,0.5,0.9 +""" +import argparse, json, time, urllib.request, urllib.error + +FILLER = ("The quick brown fox jumps over the lazy dog. " + "Grand vistas of rolling hills stretched to the horizon. ") + +def chat(url, model, prompt, max_tokens=256, timeout=180): + body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, "temperature": 0, + }).encode() + req = urllib.request.Request(url.rstrip("/") + "/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + dt = time.time() - t0 + ch = d["choices"][0]["message"] + txt = (ch.get("content") or "") + " " + (ch.get("reasoning") or "") + return txt.strip(), d.get("usage", {}), dt + +def make_haystack(approx_words, needle, depth): + # Build filler to approx word count; insert needle at fractional depth. + words = [] + while len(words) < approx_words: + words.extend(FILLER.split()) + n = len(words) + pos = int(n * depth) + needle_words = needle.split() + words[pos:pos] = needle_words + return " ".join(words) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", required=True) + ap.add_argument("--model", default="kimi-k3") + ap.add_argument("--ctx", type=int, default=8000, help="approx haystack words") + ap.add_argument("--depths", default="0.1,0.5,0.9") + args = ap.parse_args() + + print(f"== sanity ==") + txt, usage, dt = chat(args.url, args.model, "Reply with only: OK", max_tokens=8) + print(f" resp={txt!r} usage={usage} {dt:.1f}s") + + secret = "The secret passcode is HELIOTROPE-7492." + question = ("\n\nQuestion: What is the secret passcode? " + "Answer with only the passcode.") + passed = 0; total = 0 + for depth in [float(x) for x in args.depths.split(",")]: + hay = make_haystack(args.ctx, secret, depth) + prompt = hay + question + try: + txt, usage, dt = chat(args.url, args.model, prompt, max_tokens=256) + except Exception as e: + print(f" depth={depth}: ERROR {e}") + total += 1; continue + ok = "HELIOTROPE-7492" in txt.upper() + passed += ok; total += 1 + print(f" depth={depth:>4}: {'PASS' if ok else 'FAIL'} " + f"(prompt_tok={usage.get('prompt_tokens')}, {dt:.1f}s) resp={txt[:60]!r}") + print(f"== NIAH {passed}/{total} passed ==") + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/pp2xtp8/run.sh b/scripts/vllm/kimik3_mi300x/pp2xtp8/run.sh new file mode 100755 index 00000000..aa32dcfe --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/pp2xtp8/run.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# Kimi-K3 (MXFP4) on MI300X / gfx942 -- PP2xTP8 baseline (no expert parallelism). +# TP8 within each node, PP2 across 2 nodes -> each node holds half the layers +# (~102 GB/GPU); a single 8-GPU node cannot fit the model + KV. Simplest, lowest- +# latency K3 serve on MI300X. Colocated (single instance; no P/D disaggregation). +# +# Usage (worker FIRST, then head): +# ROLE=worker MASTER= bash run.sh # on node1 (rank1) +# ROLE=head MASTER= bash run.sh # on node0 (rank0, serves API) +set -euo pipefail + +IMAGE="${IMAGE:-amdsiloai/vllm:kimi-k3-mi325x-release-v2}" +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR=}" +ROLE="${ROLE:?set ROLE=head|worker}" +MASTER="${MASTER:?set MASTER=}" +PORT="${PORT:-8000}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-8}" +GPU_UTIL="${GPU_UTIL:-0.90}" +CONTAINER="k3_pp2tp8_${ROLE}" +LOGHOST="${LOGHOST:-$HOME/k3run/logs}"; mkdir -p "$LOGHOST" + +# Proven fabric env (from cluster_rdma_env_recommender.py) +BOOT_NIC="eth0" +IB_HCA="mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9" +GID_INDEX=3 + +if [ "$ROLE" = "head" ]; then + NODE_RANK=0; SERVE_EXTRA="--port ${PORT}"; HEADLESS="" +else + NODE_RANK=1; SERVE_EXTRA=""; HEADLESS="--headless" +fi + +echo "[pp2tp8] node=$(hostname -s) role=$ROLE rank=$NODE_RANK master=$MASTER" +[ -f "$MODEL_DIR/model.safetensors.index.json" ] || { echo "ERROR: model missing at $MODEL_DIR"; exit 1; } +docker rm -f "$CONTAINER" 2>/dev/null || true + +docker run -d --name "$CONTAINER" \ + --network host --ipc host \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband --group-add video \ + --cap-add SYS_PTRACE --security-opt seccomp=unconfined --shm-size 128g \ + --ulimit memlock=-1 --ulimit stack=67108864 \ + -e VLLM_ROCM_USE_AITER_MLA=0 \ + -e NCCL_SOCKET_IFNAME=$BOOT_NIC -e GLOO_SOCKET_IFNAME=$BOOT_NIC \ + -e NCCL_IB_DISABLE=0 -e NCCL_IB_HCA=$IB_HCA -e NCCL_IB_GID_INDEX=$GID_INDEX \ + -e NCCL_IGNORE_CPU_AFFINITY=1 -e NCCL_DEBUG=WARN \ + -e HSA_ENABLE_IPC_MODE_LEGACY=0 -e HSA_NO_SCRATCH_RECLAIM=1 \ + -e PYTORCH_ALLOC_CONF=expandable_segments:False \ + -e PYTORCH_HIP_ALLOC_CONF=expandable_segments:False \ + -v "$MODEL_DIR":/model:ro -v "$LOGHOST":/logs \ + --entrypoint bash \ + "$IMAGE" -c " + vllm serve /model --served-model-name kimi-k3 \ + --tensor-parallel-size 8 --pipeline-parallel-size 2 \ + --distributed-executor-backend mp \ + --nnodes 2 --node-rank ${NODE_RANK} --master-addr ${MASTER} --master-port 29500 ${HEADLESS} \ + --trust-remote-code --reasoning-parser kimi_k3 --mm-encoder-tp-mode data \ + --safetensors-load-strategy prefetch \ + --max-model-len ${MAX_MODEL_LEN} --max-num-seqs ${MAX_NUM_SEQS} \ + --gpu-memory-utilization ${GPU_UTIL} ${SERVE_EXTRA} 2>&1 | tee /logs/vllm_pp2tp8_${ROLE}.log + " +echo "[pp2tp8] $ROLE started. log: $LOGHOST/vllm_pp2tp8_${ROLE}.log" +[ "$ROLE" = head ] && echo "[pp2tp8] health: curl http://${MASTER}:${PORT}/v1/models" diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/Dockerfile.kimik3_disagg b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/Dockerfile.kimik3_disagg new file mode 100644 index 00000000..873c2a66 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/Dockerfile.kimik3_disagg @@ -0,0 +1,379 @@ +# CONTEXT {'gpu_vendor': 'AMD', 'guest_os': 'UBUNTU'} +############################################################################### +# +# MIT License +# +# Copyright (c) 2025 Advanced Micro Devices, Inc. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in all +# copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. +# +################################################################################# +# ============================================================================= +# vllm_disagg_inference.ubuntu.amd.Dockerfile +# ALL connectors in one image: moriio (TP + MoRI-EP wideEP) + rixl (NIXL TP + +# DeepEP wideEP). = the fullsource MoRI stack, plus a UCX/RIXL/rocSHMEM/DeepEP +# transport layer gated by --build-arg WITH_NIXL (default 1 = everything). +# +# docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile \ +# -t /vllm-disagg:local . +# export DOCKER_IMAGE_NAME=/vllm-disagg:local +# +# WITH_NIXL=1 (default) => builds UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from +# source, so all four connector combos (moriio TP/wideEP, rixl NIXL TP, DeepEP +# wideEP) are present (~+30-45 min build vs WITH_NIXL=0). +# WITH_NIXL=0 => MoRI-EP only (moriio TP/wideEP + deepep-from-base); lean, faster. +# +# STATUS: built + validated on ci_base (WITH_NIXL=1). The moriio paths (TP + MoRI-EP +# wideEP, incl. full DeepSeek-V3) are live-proven. rixl+TP builds and NIXL initializes, +# but has an image-level PyNCCL all-reduce issue on this stack (tracked separately). +# (BASE_IMAGE is a gated nightly; override --build-arg BASE_IMAGE=... as needed.) +# ============================================================================= +# Reconstructs the validated v1.2.1 (mori121) runtime stack by applying the recipe's +# component pins ON TOP of the open ROCm vLLM ci_base, cloning each source from +# public Git (no local build-contexts). Mirrors dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable: +# +# - BASE: rocm/vllm-dev:ci_base-0fcd9b99... (open ROCm 7.2 / cp312 CI base). +# - MoRI -> built from ROCm/MoRI @ v1.2.1 (BUILD_UMBP=OFF). +# - AITER -> 0.1.16.post3 prebuilt rocm7.2 wheel + flydsl 0.2.2; stale JIT wiped. +# - vLLM -> COMPILED from the Wide-EP multi-pod PD reference branch (the +# connector/router reference for the 2P2D DP=EP=16 topology). Full compile: it is +# a different commit than the base's, so a .py-only overlay would be ABI-mismatched. +# - RDMA fix (expandable_segments:False x2 + HSA_ENABLE_IPC_MODE_LEGACY=0) is NOT baked +# here — it lives in scripts/vllm_dissag/connectors/.env and the launcher +# forwards it via docker -e. ROCm 7.2.3 cannot dmabuf-export VMM memory, else MoRI +# RegisterRdmaMemoryRegion EFAULTs (errno 14) on the first disagg WRITE. +# - vllm-router (vllm-project/router PR#181 = DP-rank round-robin + 2P2D KV-notify +# dpfix) built in -> no external router binary needed. +# - validated recipe knobs baked as ENV. The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this vLLM (no runtime patcher). +# +# Build context = repo root: +# docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile -t / . +# +# BASE_IMAGE is the OPEN ROCm 7.2 / cp312 vLLM CI base. Everything Wide-EP (MoRI, +# rocSHMEM, DeepEP, UCX/RIXL) is built FROM SOURCE in the steps below, so this base +# only needs to supply ROCm 7.2.x + torch 2.11 + PyTorch. It is publicly pullable +# (`docker pull rocm/vllm-dev:ci_base-...`) so this recipe builds end-to-end with no +# gated images. It must share the ROCm/torch ABI of PROVEN_K3_IMAGE so the grafted +# AITER trees stay compatible. vLLM compile is long (~30-60 min). +# ============================================================================= + +# Image ARGs used by FROM must be declared BEFORE the first FROM (buildkit global +# scope). Declaring BASE_IMAGE after the first FROM scopes it to that stage only, +# so the second FROM resolves it blank -> "base name should not be blank". +ARG PROVEN_K3_IMAGE=amdsiloai/vllm:kimi-k3-mi325x-release-v2 +# Public open base — builds the full Wide-EP stack from source (steps 1-4b below). +ARG BASE_IMAGE=rocm/vllm-dev:ci_base-0fcd9b99cc9d63202da4c858d8ebc6582c9e2491 +# (internal-only alt base that pre-bakes MoRI 1.2.x, if you have registry access: +# rocmshared/pytorch-private:vllm-rocm_07_22_2026_shikpate_mori1.2.3 — then steps 1-4b +# are redundant/overridable. The public ci_base above is the supported default.) + +# K3-aware AITER donor stage (see step 5b): ships kimik3 tuned MXFP4 MoE configs, +# working MXFP4 codegen, flydsl 0.2.4 and aiter/ops/triton/conv. Same ROCm 7.2.3 + +# torch 2.11 base as our build base, so the grafted trees are ABI-compatible. +FROM ${PROVEN_K3_IMAGE} AS proven_k3_aiter + +FROM ${BASE_IMAGE} + +ENTRYPOINT [] +WORKDIR /app + +ARG GFX_COMPILATION_ARCH="gfx942" +ARG PYTORCH_ROCM_ARCH="gfx942" +ARG MAX_JOBS=32 +# NIXL/RIXL transport for the rixl connector. Default 1 => all connectors built +# (UCX/RIXL/rocSHMEM/DeepEP). Set --build-arg WITH_NIXL=0 for a lean MoRI-EP-only image. +ARG WITH_NIXL=1 +ARG NIC_COMPILATION_ARCH="cx7" + +# ----------------------------------------------------------------------------- +# 1. MoRI: replace the base's bundled MoRI with the validated ROCm/MoRI @ v1.2.1 +# (the version for the 06_29 mori121 image, dist-inf-cookbook +# Dockerfile.vllm.mori121_shareable). v1.2.1 carries the EP/RDMA correctness fixes +# plus the ROCm-7.2.3 dmabuf registration path used by the connector .env +# (expandable_segments:False). MoRI is JIT-built, so this swaps the JIT sources the +# kernels compile from at runtime. +# BUILD CONFIG: match the cookbook build — MORI_GPU_ARCHS=gfx942, BUILD_UMBP=OFF, +# DEFAULT NIC backends. Do NOT pass USE_IONIC=OFF / USE_BNXT=OFF: disabling NIC +# backends produced a MoRI that deadlocked at the cross-node EP all-to-all init. +# ----------------------------------------------------------------------------- +ARG MORI_REPO=https://github.com/ROCm/mori.git +ARG MORI_REF=v1.2.2 +ENV MORI_GPU_ARCHS=gfx942 +# Newer MoRI added the UMBP subsystem which requires gRPC (grpcpp/grpcpp.h) not +# present in this base; UMBP is unrelated to the EP dispatch/combine kernels, so +# disable it to avoid pulling in a gRPC build dependency. +ENV BUILD_UMBP=OFF BUILD_UMBP_SPDK=OFF +# Build/install matches dist-inf-cookbook Dockerfile.vllm.mori121_shareable for v1.2.1: +# `BUILD_UMBP=OFF pip install .` (default build isolation). apt/pip build tooling kept +# for bases that lack it; harmless where already present. +RUN sed -i 's|http://|https://|g' /etc/apt/sources.list 2>/dev/null || true && \ + sed -i 's|http://|https://|g' /etc/apt/sources.list.d/*.list 2>/dev/null || true && \ + apt-get update && apt-get install -y --no-install-recommends \ + git build-essential cmake ninja-build ccache libssl-dev pkg-config curl ca-certificates && \ + pip install meson==0.64.0 "pybind11[global]" tqdm prettytable && \ + pip uninstall -y amd_mori amd-mori amd-mori-nightly mori 2>/dev/null || true && \ + rm -rf /tmp/mori-src && \ + git clone --recursive "${MORI_REPO}" /tmp/mori-src && \ + cd /tmp/mori-src && git checkout "${MORI_REF}" && git submodule update --init --recursive && \ + BUILD_UMBP=OFF pip install . && \ + python3 -c "import mori, mori.io, mori.ops; print('MoRI OK at', mori.__path__[0])" && \ + mkdir -p /app && echo "MORI_REF=${MORI_REF}@$(git -C /tmp/mori-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/mori-src + +# ----------------------------------------------------------------------------- +# 2. AITER: install 0.1.16.post3 (prebuilt rocm7.2 wheel + flydsl 0.2.2), then +# invalidate the stale prewarmed AITER JIT cache compiled against the old .so. +# ----------------------------------------------------------------------------- +# K3 NOTE: use AITER 0.1.19 (NOT 0.1.16.post3). K3's vision tower imports +# aiter.ops.triton.conv.conv2d which first appears in 0.1.19; 0.1.16.post3 lacks +# it -> "No module named 'aiter.ops.triton.conv'". 0.1.19 (2026-07-27) also +# includes the #3658 top_k_top_p HSA-fault fix (merged 2026-06-16) needed for +# DP-EP disagg. The base image already ships 0.1.19.dev24 with conv+#3658, so +# this is effectively a no-op/keep — but pin the release wheel for reproducibility. +ARG AITER_VERSION=0.1.19 +ARG AITER_WHEEL_URL="https://github.com/ROCm/aiter/releases/download/v0.1.19/amd_aiter-0.1.19%2Brocm7.2.manylinux.2.28-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl" +RUN echo "Bumping AITER to ${AITER_VERSION} from ${AITER_WHEEL_URL}" && \ + _W="/tmp/$(basename "${AITER_WHEEL_URL}" | sed 's/%2B/+/g')" && \ + curl -fL --retry 3 --retry-delay 2 -o "${_W}" "${AITER_WHEEL_URL}" && \ + (pip uninstall -y amd_aiter amd-aiter aiter 2>/dev/null || true) && \ + pip install --no-deps "${_W}" && \ + pip install "flydsl==0.2.4" && \ + rm -f "${_W}" && \ + python3 - <<'PYEOF' +from importlib.metadata import version as v, PackageNotFoundError +vm = None +for n in ("amd-aiter", "amd_aiter", "aiter"): + try: vm = v(n); break + except PackageNotFoundError: pass +assert vm and vm.split("+", 1)[0].startswith("0.1.19"), f"AITER not 0.1.19: {vm!r}" +print("AITER OK:", vm) +PYEOF +RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER JIT cache" && \ + echo "AITER_VERSION=${AITER_VERSION}" >> /app/versions.txt + +# ----------------------------------------------------------------------------- +# 3. vLLM: compile from source at the 06_29 validated Wide-EP WRITE-mode branch +# (matches the published dist-inf-cookbook mori121 image). Full source compile +# (the base ships a different commit). The MoRIIO disagg fixes (#39276 notify, +# #41751 LL split, DP-rank hash-failsafe) are native in this branch, so no runtime +# patcher is needed. Override VLLM_REF to rebuild a different commit; build only +# committed commits (no working-tree edits). +# ----------------------------------------------------------------------------- +# VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the +# disagg image builds from): the K3 + MoRIIO connector branch with the folded +# connector fixes (mamba block-ids, remote_tp normalize, mamba N-1). Override to your +# own fork/branch. +ARG GH_TOKEN="" +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +# v3 branch = v2 (4-KV-group routing + multi-chunk compute-progress gate) PLUS +# the KDA gather sync-free fix that unblocks long context (>500K): removes a +# per-KDA-layer device->CPU sync that hung 750K/900K prefills. With this branch +# the image has ALL fixes BAKED IN; the runtime patchers in the recipe's +# patchers/ dir then detect "already applied" and no-op (idempotent). +ARG VLLM_REF=kimi-k3-wideep-disagg-fullsource-v3 +ENV VLLM_TARGET_DEVICE=rocm \ + PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ + MAX_JOBS=${MAX_JOBS} +RUN rm -rf /tmp/vllm-src && \ + _VURL="${VLLM_REPO}"; [ -n "${GH_TOKEN}" ] && _VURL=$(echo "${VLLM_REPO}" | sed "s|https://|https://x-access-token:${GH_TOKEN}@|"); \ + git clone "${_VURL}" /tmp/vllm-src && \ + cd /tmp/vllm-src && git checkout "${VLLM_REF}" && \ + echo "VLLM_REF=${VLLM_REF}@$(git rev-parse HEAD)" >> /app/versions.txt && \ + pip uninstall -y vllm 2>/dev/null || true && \ + MAX_JOBS="${MAX_JOBS:-32}" NVCC_THREADS="${NVCC_THREADS:-8}" \ + pip install --no-deps --no-build-isolation -v . && \ + python3 -c "import vllm; print('vLLM', vllm.__version__, 'from', vllm.__file__)" && \ + rm -rf /tmp/vllm-src + +# Cross-check MoRI + AITER survived the vLLM install (no silent downgrade). +RUN python3 - <<'PYEOF' +from importlib.metadata import version as v, PackageNotFoundError +def get(names): + for n in names: + try: return v(n) + except PackageNotFoundError: pass + return None +av = get(("amd-aiter", "amd_aiter", "aiter")) +assert av and av.split("+", 1)[0].startswith("0.1.19"), f"AITER not 0.1.19: {av!r}" +import mori, mori.io, mori.ops +print("Post-vLLM check OK: AITER", av, "+ MoRI importable") +PYEOF + +# ----------------------------------------------------------------------------- +# 4. vllm-router (DP-rank round-robin + MoRIIO connector) — built in, so NO +# external vllm-router binary is needed (leave ROUTER_BINARY unset). +# Source = vllm-project/router PR #181 branch, which now carries BOTH the +# round-robin DP-rank fix (11841c0d) AND the 2P2D KV-notify fix (6409ac1: +# remote_dp_rank_override + remote_dp_size). The KV-notify fix is REQUIRED: +# without it the 2P2D EP=16 run reproducibly wedges with "remote blocks never +# arrived" deferred-write expiries (decode notify targets the wrong DP rank). +# This is the exact source of the validated vllm-router-2p2d-dpfix binary. +# Pinned Rust toolchain (>=1.88: router deps time/home require rustc 1.88). +# ----------------------------------------------------------------------------- +ARG ROUTER_REPO=https://github.com/raviguptaamd/router.git +ARG ROUTER_REF=ravgupta/discovery-dp-rank-roundrobin +ARG RUST_TOOLCHAIN=1.88.0 +RUN if ! command -v cargo >/dev/null 2>&1; then \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain "${RUST_TOOLCHAIN}"; \ + fi && \ + export PATH="/root/.cargo/bin:${PATH}" && \ + rm -rf /tmp/vllm-router-src && \ + _RURL="${ROUTER_REPO}"; [ -n "${GH_TOKEN}" ] && _RURL=$(echo "${ROUTER_REPO}" | sed "s|https://|https://x-access-token:${GH_TOKEN}@|"); \ + git clone --filter=blob:none "${_RURL}" /tmp/vllm-router-src && \ + cd /tmp/vllm-router-src && git checkout "${ROUTER_REF}" && \ + cargo build --release && \ + install -m 755 target/release/vllm-router /usr/local/bin/vllm-router && \ + vllm-router --help 2>&1 | grep -q moriio && \ + echo "VLLM_ROUTER_REF=${ROUTER_REPO}@${ROUTER_REF}@$(git -C /tmp/vllm-router-src rev-parse HEAD)" >> /app/versions.txt && \ + rm -rf /tmp/vllm-router-src + +# ----------------------------------------------------------------------------- +# 4b. WITH_NIXL=1 (default): UCX + RIXL(+nixlbench) + rocSHMEM + DeepEP from source, +# so the rixl connector (NIXL TP + DeepEP wideEP) is present. Single guarded RUN so +# WITH_NIXL=0 skips it entirely (no layers, no cost). Build-verified on ci_base. +# ----------------------------------------------------------------------------- +ENV _ROCM_DIR=/opt/rocm \ + _UCX_SOURCE=https://github.com/ROCm/ucx.git \ + _UCX_BRANCH=da3fac2a \ + _UCX_INSTALL_DIR=/usr/local/ucx/ \ + _RIXL_SOURCE=https://github.com/ROCm/RIXL.git \ + _RIXL_BRANCH=f33a5599 \ + _RIXL_INSTALL_DIR=/usr/local/RIXL/install \ + _NIXLBENCH_INSTALL_DIR=/usr/local/RIXL +RUN if [ "${WITH_NIXL}" != "1" ]; then \ + echo "WITH_NIXL=${WITH_NIXL}: skipping UCX/RIXL/rocSHMEM/DeepEP (MoRI-EP + base DeepEP only)"; \ + else set -e && \ + echo "WITH_NIXL=1: building UCX + RIXL + rocSHMEM + DeepEP" && \ + apt-get update && apt-get install -y \ + autoconf automake libtool autogen pkg-config m4 gcc make \ + librdmacm-dev rdmacm-utils infiniband-diags ibverbs-utils perftest ethtool \ + libibverbs-dev rdma-core strace libgflags-dev \ + libaio-dev liburing-dev libcpprest-dev libgrpc-dev libgrpc++-dev \ + libprotobuf-dev protobuf-compiler-grpc wget && \ + pip install meson==0.64.0 "pybind11[global]" pyyaml && \ + # UCX + cd /tmp && git clone "${_UCX_SOURCE}" && cd ucx && git checkout "${_UCX_BRANCH}" && \ + ./autogen.sh && mkdir -p build && cd build && \ + ../configure --prefix="${_UCX_INSTALL_DIR}" --with-rocm="${_ROCM_DIR}" \ + --disable-go --disable-java --disable-assertions --enable-mt && \ + make -j && make install && \ + # googletest (RIXL dep) + cd /tmp && wget -q https://github.com/google/googletest/archive/refs/tags/v1.14.0.tar.gz && \ + tar -xzf v1.14.0.tar.gz && cd googletest-1.14.0 && mkdir -p build && cd build && \ + cmake -DBUILD_SHARED_LIBS=on .. && make -j && make install && \ + # RIXL + python bindings + cd /tmp && git clone "${_RIXL_SOURCE}" && cd RIXL && git checkout "${_RIXL_BRANCH}" && \ + meson setup build/ --prefix="${_RIXL_INSTALL_DIR}" -Ducx_path="${_UCX_INSTALL_DIR}" \ + -Ddisable_gds_backend=true -Dcudapath_inc="${_ROCM_DIR}/include" -Dcudapath_lib="${_ROCM_DIR}/lib" && \ + cd build && ninja && ninja install && cd /tmp/RIXL && \ + pip install --config-settings=setup-args="-Dcudapath_inc=${_ROCM_DIR}/include" \ + --config-settings=setup-args="-Dcudapath_lib=${_ROCM_DIR}/lib" \ + --config-settings=setup-args="-Ducx_path=${_UCX_INSTALL_DIR}" \ + --config-settings=setup-args="-Ddisable_gds_backend=true" . && \ + # rocSHMEM (DeepEP dep) + cd /tmp && git clone --no-checkout --filter=blob:none https://github.com/ROCm/rocm-systems.git && \ + cd rocm-systems && git sparse-checkout set --cone projects/rocshmem && git checkout develop && \ + mkdir -p /tmp/rocshmem-build && cd /tmp/rocshmem-build && \ + /tmp/rocm-systems/projects/rocshmem/scripts/build_configs/all_backends \ + -DUSE_EXTERNAL_MPI=OFF -DGPU_TARGETS="${GFX_COMPILATION_ARCH}" && \ + # DeepEP (build develop against the installed vLLM/torch) + cd /tmp && git clone https://github.com/ROCm/DeepEP.git && cd DeepEP && \ + PYTORCH_ROCM_ARCH="${GFX_COMPILATION_ARCH}" CFLAGS="-O3 -fPIC" \ + CXXFLAGS="-O3 -fPIC --offload-arch=${GFX_COMPILATION_ARCH}" HIP_CXX_FLAGS="-O3 -fPIC" \ + python3 setup.py --variant rocm --nic "${NIC_COMPILATION_ARCH}" build develop && \ + echo "WITH_NIXL build complete" >> /app/versions.txt && \ + rm -rf /tmp/ucx /tmp/googletest-1.14.0 /tmp/v1.14.0.tar.gz /tmp/rocm-systems /tmp/rocshmem-build; \ + fi +ENV LD_LIBRARY_PATH="/usr/local/ucx/lib:/usr/local/lib:/usr/local/RIXL/install/lib:${LD_LIBRARY_PATH}" \ + PATH="/usr/local/ucx/bin:${PATH}" + +# ----------------------------------------------------------------------------- +# 5. Cache locations (structural: WHERE the JIT/compile caches live in the image). +# These are the mount target for the launcher's persistent host JIT cache. +# ----------------------------------------------------------------------------- +# The image ships NO runtime recipe / tuning / platform ENV. By design, everything +# run-tunable is applied at launch, so this image stays a clean binary/library artifact +# and the same image serves any model/cluster without a rebuild: +# - model-serving recipe (KV_BLOCK_SIZE, KV_CACHE_DTYPE, *_CUDAGRAPH_MODE, *_MORI_BACKEND, +# GPU_MEMORY_UTILIZATION, KV_CACHE_MEMORY_BYTES, VLLM_ROCM_USE_AITER_MLA, ...) +# -> scripts/vllm_dissag/models.yaml (per-model env:, so dense vs MoE differ) +# - ROCm-7.2.3 GPU-RDMA platform env (expandable_segments:False x2, MORI_GPU_ARCHS, +# HSA_ENABLE_IPC_MODE_LEGACY=0, HSA_NO_SCRATCH_RECLAIM) and the MoRI/RDMA fabric +# tuning (MORI_RDMA_TC/SL, MORI_IB_GID_INDEX, MORI_NUM_QP_PER_PE, VLLM_MORIIO_*, ...) +# -> scripts/vllm_dissag/connectors/.env (cluster-editable, no rebuild) +# The slurm launcher forwards both via `docker -e` (platform env must reach PID 1 - +# PyTorch reads alloc-conf at import). Running this image WITHOUT the launcher: set the +# vars you need yourself (see connectors/moriio.env + models.yaml for the values). +ENV AITER_JIT_DIR=/opt/vllm_cache/aiter_jit \ + VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + COMGR_CACHE_DIR=/opt/vllm_cache/comgr + +# ----------------------------------------------------------------------------- +# 5b. K3-AWARE AITER GRAFT (the crux for K3 MXFP4 MoE on gfx942). +# The 0.1.19 release wheel installed in step 2 has NO Kimi-K3 MoE tuning. At the +# K3 MoE profiling shape (gfx942, M=EPxmax_tokens=131072, N=3584, K=3072, Situv2, +# mxfp4) it finds no tuned FlyDSL config and falls back to a heuristic kernel +# (flydsl_moe1_abf16_wfp4_bf16_t64x128x256_w4_bnt0) whose buffer.load.lds intrinsic +# aborts LLVM: "Do not know how to expand this operator's operand!" -> the worker +# dies natively in determine_available_memory (no Python traceback), engine init fails. +# The proven colocated image (PROVEN_K3_IMAGE, same ROCm 7.2.3 + torch 2.11 base) +# ships a K3-AWARE AITER: configs/model_configs/kimik3_{a8w4,fp4}_tuned_fmoe.csv + a +# working MXFP4->CK/int4 routing + flydsl 0.2.4 + aiter/ops/triton/conv (the K3 +# vision-tower dep, our original reason for 0.1.19). Grafting its aiter + aiter_meta +# (prebuilt hsaco) + flydsl trees over the 0.1.19 install makes K3 MXFP4 MoE compile. +# The `proven_k3_aiter` stage is declared at the top of the file (next to the main +# FROM); this COPY stays in the main stage. Placed AFTER the vLLM/router/NIXL layers +# so a later edit here does not invalidate the ~40-min vLLM compile cache. +# ----------------------------------------------------------------------------- +RUN rm -rf /usr/local/lib/python3.12/dist-packages/aiter \ + /usr/local/lib/python3.12/dist-packages/aiter_meta \ + /usr/local/lib/python3.12/dist-packages/flydsl \ + /usr/local/lib/python3.12/dist-packages/aiter*.dist-info \ + /usr/local/lib/python3.12/dist-packages/amd_aiter*.dist-info 2>/dev/null || true +COPY --from=proven_k3_aiter /usr/local/lib/python3.12/dist-packages/aiter /usr/local/lib/python3.12/dist-packages/aiter +COPY --from=proven_k3_aiter /usr/local/lib/python3.12/dist-packages/aiter_meta /usr/local/lib/python3.12/dist-packages/aiter_meta +COPY --from=proven_k3_aiter /usr/local/lib/python3.12/dist-packages/flydsl /usr/local/lib/python3.12/dist-packages/flydsl +# The proven_k3 donor's flydsl tree is 0.2.2, which the COPY above drops OVER the +# pip flydsl==0.2.4 from step 2. But K3's int4 SiTUv2 MoE path +# (_setup_kernel_k3_situ_gfx942 -> compile_moe_gemm1) hard-requires flydsl>=0.2.4 +# (ImportError otherwise -> WorkerProc init fails -> pool never starts). Re-pin 0.2.4 +# AFTER the graft so it wins, and verify. The K3-tuned MoE configs live in aiter/aiter_meta +# (still grafted); flydsl is the codegen frontend and 0.2.4 is ABI-compatible with them. +RUN pip install --no-cache-dir --force-reinstall "flydsl==0.2.4" && \ + python3 -c "import importlib.metadata as m; v=m.version('flydsl'); assert v=='0.2.4', f'flydsl {v}!=0.2.4'; print('flydsl OK', v)" && \ + echo "FLYDSL_REPIN=0.2.4 (after proven_k3 graft)" >> /app/versions.txt +RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && \ + echo "AITER_GRAFT=proven_k3 (kimik3 tuned fmoe configs)" >> /app/versions.txt + +# ----------------------------------------------------------------------------- +# 6. CRITICAL: scrub build-time MoRI JIT state. The `import mori` verification +# steps above compile/lock MoRI EP kernels under /root/.mori/jit on THIS build +# host, leaving stale .hsaco.lock files (ep_internode_v1, ep_internode_v1ll, ...). +# At runtime on the cluster, MoriAll2AllManager finds those locks, waits on a +# build-in-progress whose owner PID is long gone, and DEADLOCKS at ep:0 init. +# A clean image ships /root/.mori empty -> runtime compiles fresh. +# Clearing these makes the from-source image boot clean on 2P2D/4P4D. +# ----------------------------------------------------------------------------- +RUN rm -rf /root/.mori /tmp/mori_jit_* && mkdir -p /root/.mori && \ + echo "JIT_SCRUBBED: /root/.mori + /tmp/mori_jit_* cleared at build end" >> /app/versions.txt + +RUN cat /app/versions.txt 2>/dev/null | tail -20 || true diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/OPTIMIZATION.md b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/OPTIMIZATION.md new file mode 100644 index 00000000..bb26d2df --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/OPTIMIZATION.md @@ -0,0 +1,72 @@ +# Optimization notes — EP16 2P/2D disagg + +The recipe is tuned for correctness + long-context accuracy first. These are the +serving-capacity levers, with the one validated change called out. + +## ✅ Raise the KV cache (validated — big win, no downside seen) +`KV_CACHE_MEMORY_BYTES` is **pinned** (default `8e9`) to skip a boot `profile_run` +that hangs under TP2×DP8 + MoRI all2all — it is NOT a memory limit. At `8e9` the +GPU KV cache is only **542,372 tokens** (→ 1.69× concurrency at 320K ctx), while +there is **~72 GB/GPU free**. + +Raising it to `40e9` measured: +- **GPU KV cache: 2,836,158 tokens** (5.2×) +- **Maximum concurrency at 1,000,000-token ctx: 2.84×** (was 1.69× at 320K) +- Also required for any **single request > ~600K tokens** (a 900K request needs + 900K tokens of KV — the 542K cache can't hold it). + +```bash +KV_CACHE_MEMORY_BYTES=40000000000 ... bash run_2p2d_launch.sh +``` +Leave ~8–10 GB/GPU slack for activations + the MoRI heap. Recommended default for +throughput-oriented deployments. + +## `max_num_seqs` (raise with the KV cache) +Decode caps concurrent requests at `max_num_seqs` (default 8 = DP8 replicas × 1 +seq). With the larger KV cache above, raising to 16–32 lets each replica batch +multiple decodes → higher aggregate tok/s under load. Bounded by KV capacity; too +high with a small cache causes preemption. + +## `max_num_batched_tokens` — a dead end (keep 2048) +Raising 2048 → 8192 was measured **worse**: 200K single-stream 88s → 113s, and +20K conc=8 throughput 0.353 → 0.245 req/s. The prefill is compute-bound, not +chunk-overhead-bound, and larger prefill batches contend more across 8 replicas. +Keep `MAX_NUM_BATCHED_TOKENS=2048`. + +## Single-stream latency — architectural, little to gain +A single request runs on one DP replica (TP2 = 2 GPUs); a colocated PP2×TP8 serve +spreads one request across all 16 GPUs, so it wins single-stream by ~4×. That is +the design trade for concurrent throughput (5.7× at conc 8, 7.3× at 16 — see +[RESULTS.md](RESULTS.md)). Use disagg for high-QPS/batch, colocated for +low-latency interactive. No config closes the single-stream gap. + +## Residual write race → multi-needle 10/10 (correctness/quality, code change) +Multi-needle NIAH dips to ~9/10 at ≥ 20K (single-needle unaffected). Root cause: +decode can read a block before its RDMA write is globally visible in decode HBM +(`wait_for_layer_load()` is a no-op; `write_done` races the RDMA write). Proper +fix = a decode-side per-request KV-ready barrier before the model forward +(vLLM MoRIIO change, not a knob). Sender-side knobs (delay fence, `post_batch_size` +split) were measured useless and add latency. + +## >500K prefill hang — FIXED (KDA gather sync-free) +Contexts above ~500K used to hang (GPUs 100%, no progress). Root cause found with +py-spy: `gather_initial_states` ran a `bool((indices>=n).any())` device→CPU sync +per KDA layer per prefill chunk (~25k full stream drains at 750K) purely to log a +warning; the index clamp above it already made the address safe. Gating that +diagnostic behind `K3_KDA_GATHER_LOG=1` (default OFF) removed the drains: +**750K and 900K now pass (542s / 717s), 500K unchanged (301s), sub-quadratic +scaling.** Baked into vLLM branch `-v3`; also shipped as +`patchers/apply_kimik3_kda_gather_nosync.py`. This was both the hang and a real +perf drag (fewer syncs = faster). It also removed the biggest remaining +high-context perf overhead, so no further prefill-latency lever is outstanding +beyond the architectural O(n²) of attention itself. + +## Build / durability +- **Base image pin by digest for reproducibility.** The Dockerfile pins + `BASE_IMAGE` to a dated ROCm CI/`-complete` tag (good — reproducible). For + long-term durability, pin `@sha256:…` and mirror to a registry you control, so a + rebuild survives upstream tag GC. +- **First-launch JIT compile** (AITER + MoRI-EP kernels) is the ~4–8 min warmup; + it caches to node-local `/tmp/$USER/vllm_jit_cache` and is fast thereafter. A + fresh image pays this once (and can transiently trip a WorkerProc init race — + relaunch the affected pool once and it comes up clean). diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/README.md b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/README.md new file mode 100644 index 00000000..3ae3ea09 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/README.md @@ -0,0 +1,230 @@ +# wideep_disagg_2p2d — Kimi-K3 MI300X 2P/2D EP16 MoRIIO disagg + +> **✅ VALIDATED.** Single-needle NIAH passes **deterministically through 300K +> tokens** (all depths) on this 2-prefill + 2-decode EP16 disagg serve. The +> previously-open decode-recall bug is fixed (two root causes: 4-KV-group block +> routing + multi-chunk prefill transfer). See [`STATUS.md`](STATUS.md) for the +> root-cause write-up and [`RESULTS.md`](RESULTS.md) for the full NIAH + +> latency/throughput tables. +> +> **When to use this vs. colocated:** disagg (TP2×DP8) is for **concurrent +> throughput** and decode-latency isolation — 5.7× throughput at concurrency 8, +> 7.3× at 16. A *single* request runs on one DP replica (2 GPUs), so single-stream +> latency is ~4× a colocated PP2×TP8 serve (which spreads one request across all +> 16 GPUs). For low-latency interactive / low-QPS, use the colocated recipes +> ([`../pp2xtp8`](../pp2xtp8), [`../wideep_int4_moriep`](../wideep_int4_moriep)); +> for high-QPS / batch, use this. + +## What this is + +Prefill/decode **disaggregated** Kimi-K3 across **4 MI300X nodes**: + +``` + Prefill pool Decode pool + ┌───────────────┐ MoRIIO (RDMA) ┌───────────────┐ + │ PM + PW │ KV + KDA state │ DM + DW │ + │ 2 nodes = 16 │ ────────────────► │ 2 nodes = 16 │ + │ GPU, EP16 │ write + notify │ GPU, EP16 │ + └───────────────┘ └───────────────┘ + ▲ router (:30000) on PM fans requests P → D +``` + +- **Per pool**: TP2 × DP8 → **EP16** expert-parallel via **MoRI-EP** all2all. +- **No pipeline parallelism.** Disaggregation is the only cross-stage split; + everything else is DP/EP. +- **Connector**: MoRIIO in **WRITE** mode — prefill RDMA-writes both the MLA + attention KV *and* the Kimi-Delta-Attention (KDA/mamba) recurrent+conv state + into the decode engine's blocks, then notifies over TCP/ZMQ. + +K3 is hybrid: 24 MLA full-attention layers (paged fp8 KV) + 69 KDA layers +(recurrent + conv state). The two live in **separate** vLLM KV-cache groups, and +the connector transfers both — routing the KDA state by the *mamba* group's block +ids (see fix 1 below). + +## Quick start + +Runs from a control host that can `ssh` to all four nodes. Edit the node +IPs/hostnames (or export `PM_NODE/PM_IP/…`) and point `MODEL_DIR` at your +Kimi-K3-MXFP4 weights (local NVMe on every node recommended). + +```bash +cd wideep_disagg_2p2d/ + +# 1. build or obtain the disagg image (see "Image" below), tag it kimik3-wideep-disagg:latest +# 2. edit the 4 node IPs at the top of run_2p2d_launch.sh (or export PM_NODE=… etc.) + +MODEL_DIR=/path/to/Kimi-K3-MXFP4 \ +PM_NODE=… PM_IP=… PW_NODE=… PW_IP=… DM_NODE=… DM_IP=… DW_NODE=… DW_IP=… \ +AUTO_ROUTER=1 \ +K3_GROUP_ROUTING=1 K3_EXTRA_FIXES=1 LOAD_STRATEGY=lazy \ +MAX_MODEL_LEN=320000 MAX_NUM_BATCHED_TOKENS=2048 GPU_UTIL=0.85 \ +bash run_2p2d_launch.sh + +# watch both masters for "Application startup complete", then (AUTO_ROUTER does this +# for you) a single vllm-router comes up on PM:30000. Verify + probe: +curl http://:30000/v1/models +python3 niah_probe.py --url http://:30000 --model kimi-k3 \ + --ctx 50000 --depths 0.1,0.5,0.9 # PASS (deterministic to 300K) +``` + +### The winning config (flags that matter) + +| Flag | Value | Why | +|------|-------|-----| +| `K3_GROUP_ROUTING` | `1` | **Fix #1** — 4-KV-group block routing (always on). | +| `K3_EXTRA_FIXES` | `1` | **Fix #2** — multi-chunk compute-progress gate + all-group accumulation. Required for recall past `max_num_batched_tokens`. | +| `LOAD_STRATEGY` | `lazy` | `prefetch` double-loads RAM when the model is on tmpfs → decode OOM. | +| `MAX_NUM_BATCHED_TOKENS` | `2048` | Best measured throughput; raising to 8192 did **not** cut latency and hurt throughput (compute-bound prefill). | +| `MAX_MODEL_LEN` | `320000` | Needed for > 131K-token NIAH (default 131072 caps ~120K). Raise to `1000000` for the full native ctx. | +| `KV_CACHE_MEMORY_BYTES` | `8e9` | KV cache budget (pinned to skip a profile_run hang, NOT a mem limit). `8e9` = 542K tokens. **Raise to `40e9` (→ 2.84M tokens, ~72 GB/GPU free) for high throughput or single requests > ~600K.** See [OPTIMIZATION.md](OPTIMIZATION.md). | +| `GPU_UTIL` | `0.85` | 0.88 razor-misses KV headroom on some nodes. | +| `KV_CACHE_DTYPE` | `fp8` (default) | Transfer geometry assumes 1-byte elements; bf16 corrupts. | +| `PREFILL_BACKEND` | `mori_low_latency` | V1 high_throughput dispatch warmup crashes on this stack. | + +**Model on tmpfs (recommended):** loading the 1.5 TB checkpoint from a tmpfs RAM +cache (`/mnt/rammodel/Kimi-K3-MXFP4`) with `LOAD_STRATEGY=lazy` is ~2 min vs +~20 min from NFS (whose page cache gets evicted between runs). + +`run_2p2d_launch.sh` deploys the scripts + patchers + image to all four nodes, +starts **workers first, then masters**, then (with `AUTO_ROUTER=1`) waits for both +masters' `/v1/models` before launching exactly one router. `run_2p2d.sh` is the +per-node entrypoint (dispatches on `ROLE=prefill_master|prefill_worker| +decode_master|decode_worker`). + +## Load-bearing env (set by the launcher) + +| Var | Default | Meaning | +|-----|---------|---------| +| `TP_SIZE` / `DP_SIZE` / `DP_LOCAL` | 2 / 8 / 4 | TP2×DP8 → EP16 per pool; 4 DP ranks per node. | +| `PREFILL_BACKEND` | `mori_low_latency` | MoRI-EP all2all backend. | +| `KV_CACHE_MEMORY_BYTES` | 8e9 | Per-engine KV cache budget. | +| `PMASTER`/`DMASTER`/`PROXY_IP` | node IPs | Pool masters + router/proxy host. | +| `PREFILL_POD_HOSTS`/`DECODE_POD_HOSTS` | node IP lists | Pool membership. | +| `MODEL_DIR` | *(required)* | Kimi-K3-MXFP4 weights path (must exist on every node). | + +### RDMA fabric (overridable; validated defaults ON) + +`run_2p2d.sh` reads all fabric from env, so **nothing is hardcoded**, but the +defaults are the validated Broadcom **Thor2 (bnxt RoCE)** values so it works +out-of-the-box on the reference cluster: + +| Var | Default (Thor2) | Override for e.g. Mellanox | +|-----|-----------------|---------------------------| +| `SOCKET_IFNAME` | `eno0` | `eth0` | +| `NCCL_IB_HCA` / `RDMA_DEVICES` | `rdma0..rdma7` | `mlx5_0,mlx5_2,…,mlx5_9` | +| `IB_GID_INDEX` | `3` | your GID | +| `THOR2_BNXT_FIX` | `1` | `0` (non-Thor2) | + +`THOR2_BNXT_FIX=1` mounts the host **v34** `libibverbs`/`libbnxt_re` onto the +image's resolved soname (the image ships v59, but the Thor2 `bnxt_re` kernel +driver only accepts v34 → otherwise 0 RDMA devices). On non-bnxt fabric set +`THOR2_BNXT_FIX=0`. Example override: +```bash +SOCKET_IFNAME=eth0 NCCL_IB_HCA=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9 \ +RDMA_DEVICES=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9 \ +IB_GID_INDEX=3 THOR2_BNXT_FIX=0 bash run_2p2d_launch.sh + +## The connector fixes (folded into the vLLM source branch) + +The two root-cause fixes are **folded into vLLM source** on branch +`kimi-k3-wideep-disagg-fullsource-v2` of `raviguptaamd/vllm` (the `-v2` = the +base `fullsource` branch + these fixes baked in). The Dockerfile builds that +branch (`VLLM_REF=kimi-k3-wideep-disagg-fullsource-v2`), so the **image has the +fixes baked in**. The same fixes also ship as runtime patchers under +[`patchers/`](patchers/); on a v2 image they detect "already applied" and no-op +(idempotent) — so the recipe also works on an older/unfolded image. To rebuild +the image from scratch see the [Image](#image) section. + +**The two root-cause fixes that make NIAH pass** (full write-up in +[`STATUS.md`](STATUS.md)): + +1. **4-KV-group block routing** (`apply_kimik3_moriio_group_routing.py`, always + on, `K3_GROUP_ROUTING=1`) — K3 has **4** KV-cache groups (3 KDA/mamba + 1 MLA); + the shipped connector hardcoded 2-group indices and sent MLA KV to mamba + block-ids. Fix carries all groups' block-ids end-to-end and routes each layer + by its own group. Fixes short (≤ 1 block) recall. +2. **Multi-chunk prefill transfer** (`apply_kimik3_chunk_gate_fix.py` + + `apply_kimik3_chunked_allgrp.py`, `K3_EXTRA_FIXES=1`) — the connector detected + the final prefill chunk by *block count*, which fires after chunk 1 when a + prompt fits in ≤ 1 padded block → only `max_num_batched_tokens` of KV crossed. + Fix gates on **compute progress** from `scheduler_output` (build map → entry + defer → accumulation final-detect → post-loop sweep). Removes the razor cliff + at `max_num_batched_tokens`; recall now scales to 300K. + +Pre-existing connector fixes (still required, always on): + +3. **mamba block-id routing** — transfer KDA/mamba state by the *mamba* KV-cache + group's block ids (superseded in the general case by fix 1; kept as fallback). +4. **remote_tp_size normalize** — degenerate `remote_tp_size ≤ 1` → `world_size`, + so KV fans out to **all** decode TP ranks (not just rank 0). +5. **mamba N−1 boundary** — producer computes through token N−1, decoder recomputes + token N (matches vLLM's nixl/mooncake hybrid-PD handling). + +[`patchers/diagnostics/`](patchers/diagnostics/) holds element-wise transport +probes (all env-gated OFF). See the Debugging section. + +## Image + +[`Dockerfile.kimik3_disagg`](Dockerfile.kimik3_disagg) is **fully self-contained and +builds from a public base** — no gated images required: + +- **Base** (`BASE_IMAGE`): `rocm/vllm-dev:ci_base-0fcd9b99...` — the open ROCm 7.2 / + cp312 vLLM CI base (publicly pullable). It only supplies ROCm + torch; the entire + Wide-EP stack is built **from source** on top: + - **MoRI** `ROCm/mori @ v1.2.2` (gfx942, `BUILD_UMBP=OFF`) — the EP all2all kernels; + - **AITER** `0.1.19` wheel + `flydsl 0.2.4`; + - **vLLM** compiled from `VLLM_REPO`/`VLLM_REF` = `raviguptaamd/vllm` branch + `kimi-k3-wideep-disagg-fullsource-v2` (the folded connector fixes above); + - **vllm-router** (DP-rank round-robin + MoRIIO KV-notify) built in; + - with `WITH_NIXL=1` (default): UCX + RIXL + rocSHMEM + DeepEP from source too. +- **K3 AITER graft** (`PROVEN_K3_IMAGE`): `amdsiloai/vllm:kimi-k3-mi325x-release-v2` + (public) — only the Kimi-K3 tuned MXFP4 MoE configs / conv kernels are copied from it. + +```bash +docker build -f Dockerfile.kimik3_disagg \ + --build-arg GH_TOKEN=$(gh auth token) \ + -t kimik3-wideep-disagg:latest . +``` + +Build args (`BASE_IMAGE`, `PROVEN_K3_IMAGE`, `MORI_REF`, `VLLM_REPO`, `VLLM_REF`, +`WITH_NIXL`, `GH_TOKEN`) let you override any component. Push the result to a registry +you control and set `HUB_IMAGE` for [`load_image.sh`](load_image.sh) to pull it onto +each node. **Note:** the image builds `VLLM_REF` from GitHub, so push the vLLM fork +branch before building for a reproducible image. + +## Tests + +| Script | What | +|--------|------| +| [`niah_probe.py`](niah_probe.py) | Single-needle NIAH via the router (`--ctx --depths`). The deliverable metric — PASS to 300K. | +| [`benchmark_niah.py`](benchmark_niah.py) | Stricter 10-needle multi-needle stress (`NIAH_WORDS=…`). | +| [`concurrency_bench.py`](concurrency_bench.py) | Concurrent throughput (req/s, tok/s, latency percentiles). | + +See [`RESULTS.md`](RESULTS.md) for the full NIAH sweep + latency/throughput tables. + +## Debugging (opt-in, all default OFF) + +Turn diagnostics up/down via env flags; none change default behavior: + +| Flag | Effect | +|------|--------| +| `K3_CHUNK_GATE_DEBUG=1` | Log the chunk-gate decision per request (`entry`/`accum`/`sweep`, computed/scheduled/npt). | +| `K3_XFER_PROBE=1` | Producer offsets + src checksum, decode read-block + dst checksum. | +| `K3_DECODE_RECV_PROBE=1` | Decode reads its own KV slot norm on write-completion (~0 = bytes didn't land). | +| `K3_KDA_STATE_PROBE=1` | KDA recurrent/conv state norm at the decode read slot. | +| `K3_WRITE_BC=1` / `K3_HS_BC=1` / `K3_INPUTS_PROBE=1` | Write-delivery / handshake-dial / decode-inputs breadcrumbs. | +| `K3_WRITE_FENCE=delay K3_WRITE_FENCE_MS=…` | Sender write→notify delay (investigative; did not fix the residual race). | + +**Verify a patch landed in-container** (`/patchers` is bind-mounted; a stale file +silently re-applies old behavior): +```bash +docker exec k3disagg_prefill_master bash -lc \ + 'B=/usr/local/lib/python3.12/dist-packages/vllm/distributed/kv_transfer/kv_connector/v1/moriio; \ + grep -c k3-group-routing $B/moriio_connector.py; grep -c _k3_prog $B/moriio_connector.py' +``` + +## Status + +**VALIDATED** — single-needle NIAH passes deterministically to 300K. See +[`STATUS.md`](STATUS.md) for root cause, the fix, and the known residual +(multi-needle write race at ≥ 20K, single-needle unaffected). diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/RESULTS.md b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/RESULTS.md new file mode 100644 index 00000000..4be89153 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/RESULTS.md @@ -0,0 +1,99 @@ +# Results — Kimi-K3 MI300X 2P/2D EP16 MoRIIO disagg + +Serve: 2 prefill + 2 decode nodes; per-pool TP2×DP8 → EP16; MoRIIO WRITE; +router on the prefill master :30000. Config: `K3_GROUP_ROUTING=1 +K3_EXTRA_FIXES=1 LOAD_STRATEGY=lazy KV_CACHE_DTYPE=fp8 +MAX_NUM_BATCHED_TOKENS=2048 GPU_UTIL=0.85 PREFILL_BACKEND=mori_low_latency`, +model in tmpfs. `MAX_MODEL_LEN=320000` for the ≥ 150K rows (131072 otherwise). + +## Single-needle NIAH — the deliverable metric +`niah_probe.py`, needle = `HELIOTROPE-7492`, greedy (temp=0), depths +0.1 / 0.5 / 0.9. **All PASS, deterministic.** + +| context (tokens) | result | eval time / request | +|------------------|------------|---------------------| +| 10K | 3/3 PASS | 5.3s | +| 50K | 3/3 PASS | 19.5s | +| 100K | 3/3 PASS | ~47s | +| 120K | 3/3 PASS | ~54s | +| 150K | 3/3 PASS | ~84s | +| 200K | 3/3 PASS | ~88s | +| **300K** | **3/3 PASS** | **~150s** | +| 500K | 3/3 PASS | ~301s | +| 750K | PASS | ~542s | +| 900K | **PASS** | **~717s** | + +**Full native-context range (10K–900K) passes** with the KDA gather sync-free fix +(vLLM branch `-v3`; see below). Eval time is sub-quadratic (500K 301s → 750K 542s +→ 900K 717s). For ctx > ~120K raise `MAX_MODEL_LEN` (default 131072 caps ~120K, +use `1000000` for full native); for a single request > ~600K raise the KV cache +`KV_CACHE_MEMORY_BYTES=40e9` (2.84M-token cache; default `8e9` = 542K tokens is +too small for one 600K+ request). See [OPTIMIZATION.md](OPTIMIZATION.md). + +**The >500K unblock (KDA gather sync-free).** Contexts above ~500K previously +hung: `gather_initial_states` ran a diagnostic `bool((indices>=n).any())` per KDA +layer per prefill chunk, each forcing a device→CPU sync (full stream drain) — +~25k drains at 750K, stalling so hard it looked like a deadlock (py-spy: the DP +rank with real work stuck in that sync while the other DP ranks waited at the +batch-coordination all_reduce). The fix gates that diagnostic behind +`K3_KDA_GATHER_LOG=1` (default OFF); the index clamp still applies, so correctness +is unchanged. Folded into vLLM branch `kimi-k3-wideep-disagg-fullsource-v3` (the +Dockerfile builds it) and shipped as `patchers/apply_kimik3_kda_gather_nosync.py`. + +**One-time warmup:** a *fresh* serve pays a single aiter MLA-kernel JIT compile +(`fmha_fwd_hd192x128`, ~15 min) on the first ≥ ~200K-token request; cached +thereafter, so the per-request times above are the warm times. + +Note: `MAX_NUM_BATCHED_TOKENS=8192` was retested and still **corrupts generation** +on this stack (garbage output) — keep the default `2048`. + +Reproduce: +```bash +python3 niah_probe.py --url http://:30000 --model kimi-k3 \ + --ctx-list 10000,50000,100000,150000,200000,300000 --depths 0.1,0.5,0.9 \ + --timeout 400 +``` + +## Multi-needle stress — 10-animal (stricter) +`benchmark_niah.py`, 10 animals hidden across a word haystack, scored found/10. + +| context (words ≈ 1.3× tokens) | found/10 | +|-------------------------------|----------| +| 2000–5000 | 10/10 | +| 10000 | 10/10 | +| 15000 | 10/10 | +| 20000 | 9/10 (typical; 10/10 seen) | +| 50000 | 6–10/10 (nondeterministic) | +| 80000 | 9/10 | + +The occasional ≤ 1-needle miss at ≥ 20K is the residual RDMA write race +(see `STATUS.md` § Known residual). Single-needle retrieval is unaffected. + +## Latency vs. throughput +`concurrency_bench.py`, 20K-token prompts, 64 output tokens, batched=2048. + +| concurrency | throughput (req/s) | vs single-stream | p50 lat | p99 lat | +|-------------|--------------------|------------------|---------|---------| +| 1 | 0.062 | — | 16.2s | 16.2s | +| 8 | **0.353** | **5.7×** | 22.7s | 22.7s | +| 16 | **0.455** | **7.3×** | 34.5s | 35.1s | + +8 requests finish in 1.4× the wall time of 1 → 5.7× throughput. That is the DP8 +payoff; scaling flattens past 8 (= 8 replicas). **Single-stream latency is ~4× +a colocated PP2×TP8 EP8 serve** (1 request on 2 GPUs vs all 16) — architectural, +not tunable. + +### `max_num_batched_tokens` — a dead end (kept at 2048) +| metric | batched=2048 | batched=8192 | +|--------|--------------|--------------| +| 200K single-stream | ~88s | 112.8s (worse) | +| 20K conc=8 throughput | 0.353 req/s | 0.245 req/s (worse) | + +Raising it did not cut latency and hurt throughput (compute-bound prefill; larger +batches contend more across 8 replicas). **Keep `MAX_NUM_BATCHED_TOKENS=2048`.** + +## Use which serve +- **Colocated PP2×TP8 EP8** (`../wideep_int4_moriep`) — lowest single-request + latency; interactive / low-QPS. +- **Disagg TP2×DP8 EP16** (this) — highest concurrent throughput + decode-latency + isolation; batch / high-QPS. NIAH to 300K. diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/STATUS.md b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/STATUS.md new file mode 100644 index 00000000..e7109e1e --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/STATUS.md @@ -0,0 +1,110 @@ +# Status — Kimi-K3 MI300X 2P/2D EP16 MoRIIO disagg — **VALIDATED** + +**This recipe is validated.** Single-needle NIAH passes **deterministically +across the full native context range 10K–900K** (all depths) on the 2 prefill + +2 decode EP16 disagg serve. The decode-recall bug that previously blocked this is +**fixed**, and the >500K prefill hang is **fixed** (KDA gather sync-free — see +below + RESULTS.md). + +**Validated from a scratch build.** The two fixes are folded into vLLM source on +`raviguptaamd/vllm` branch `kimi-k3-wideep-disagg-fullsource-v2`; the Dockerfile +builds that branch, so the **image ships the fixes baked in**. A from-scratch +image build was deployed and re-validated (NIAH 12/12 to 300K) with the runtime +patchers all reporting "already applied" — i.e. the image is self-sufficient and +the patchers are idempotent no-ops on it. + +## Result + +Single-needle NIAH (needle = `HELIOTROPE-7492`, greedy `temperature=0`, depths +0.1 / 0.5 / 0.9) — **PASS at every size 10K → 900K**, deterministic: + +| ctx (tokens) | result | eval time / req | +|--------------|--------|-----------------| +| 10K | 3/3 PASS | 5.3s | +| 50K | 3/3 PASS | 19.5s | +| 100K | 3/3 PASS | ~47s | +| 200K | 3/3 PASS | ~88s | +| **300K** | **3/3 PASS** | **~150s** | +| 500K | 3/3 PASS | ~301s | +| 750K | PASS | ~542s | +| **900K** | **PASS** | **~717s** | + +Throughput (the point of DP disaggregation), 20K ctx, batched=2048: +conc=1 → 0.062 req/s; conc=8 → **0.353 req/s (5.7×)**; conc=16 → **0.455 req/s +(7.3×)**. See `RESULTS.md`. + +## Root cause (two independent bugs, both fixed) + +Kimi-K3's attention is hybrid — 24 MLA full-attention layers + 69 KDA +(Kimi-Delta-Attention, recurrent) layers → vLLM allocates **4 KV-cache groups** +(idx 0/1/2 = MambaSpec/KDA, 23 layers each; idx 3 = MLAAttentionSpec, 24 layers). +Both bugs stem from this multi-group hybrid. + +### Fix #1 — 4-KV-cache-group block routing (`patchers/apply_kimik3_moriio_group_routing.py`) +The shipped connector assumed 2 groups and hardcoded group indices `[0]`/`[1]` +when computing RDMA transfer offsets, sending MLA KV (group 3) to *group-0 +(mamba)* block-ids. Decode read the group-3 blocks, which were never written → +fluent but context-free output. Fix carries **all 4 groups' block-id lists +end-to-end** and routes each layer's transfer by its own group index. Always on +(`K3_GROUP_ROUTING=1`). This fixed short recall (≤ 1 block). + +### Fix #2 — multi-chunk prefill transfer (`apply_kimik3_chunk_gate_fix.py` + `apply_kimik3_chunked_allgrp.py`) +A razor-sharp cliff: recall died at exactly `max_num_batched_tokens`. The +connector detected the "final prefill chunk" by **block count** +(`num_prompt_tokens > len(block_ids) * self.block_size`), but the padded scheduler +block (~5760) holds a whole prompt in one block, so the KV transfer fired after +**chunk 1** (only `max_num_batched_tokens` computed); tokens past chunk 1 (e.g. a +needle at the end) were never transferred. When a prompt fits in ≤ 1 block, +block-count can never detect chunk completion. + +Fix: gate on **compute progress** from fresh `scheduler_output` +(`num_computed_tokens` + `num_scheduled_tokens`), not block count — applied at 4 +points: (A) build a per-step progress map, (B) entry defer, (C) accumulation +final-detect, (D) post-loop sweep for the final-chunk-adds-no-new-block case +(else the request deadlocks: `unmap MISS table_size=0`). `chunked_allgrp` +accumulates every group's block-ids across chunks. `SLACK=2` +(`K3_CHUNK_GATE_SLACK`) absorbs the mamba N-1 truncation. Enabled with +`K3_EXTRA_FIXES=1`. + +Verified on a 2611-token prompt: +``` +[k3-chunk-gate-entry] nblk=1 npt=2611 done=False prog=(0,2048) # chunk 1 -> defer +[k3-chunk-gate-sweep] nblk=1 npt=2611 done=True prog=(2048,563) # chunk 2 -> emit full blocks +``` + +### Fix #3 — KDA gather sync-free (`patchers/apply_kimik3_kda_gather_nosync.py`) +Contexts above ~500K hung. Root-caused with py-spy: the DP rank holding the real +request was stuck in `gather_initial_states` (the KDA recurrent-state gather) +while all other DP ranks waited at the `coordinate_batch_across_dp` all_reduce. +The native stack showed the block was a `bool((indices>=n).any())` device→CPU +sync (`_local_scalar_dense` → `memcpy_and_sync`, a full stream drain) run purely +to emit a diagnostic warning — **per KDA layer, per prefill chunk** (~25k drains +at 750K). The index clamp above it already guarantees a valid GPU address, so the +sync is pure overhead. Fix gates it behind `K3_KDA_GATHER_LOG=1` (default OFF); +correctness is unchanged (indices still clamped). With this, 750K (542s) and 900K +(717s) now pass where they previously hung indefinitely; 500K unchanged (301s), +scaling sub-quadratically. Folded into vLLM branch `-v3`. + +## Known residual (does not block single-needle NIAH to 300K) + +An RDMA write-visibility race: `write_done` travels ZMQ/TCP, a different path than +the RDMA write, and `wait_for_layer_load()` is a no-op — so at high context (many +blocks) decode can occasionally read a block before its RDMA write is globally +visible in decode HBM. Effect: the **stricter 10-needle** stress +(`benchmark_niah.py`) dips to ~9/10 at ≥ 20K; **single-needle** NIAH is +unaffected (deterministic to 300K). Sender-side mitigations (delay fence, +`post_batch_size` split) don't help and add latency. Proper fix = a decode-side +per-request KV-ready barrier before the model forward; tracked as future work. + +## Prior connector fixes (still required, pre-existing) +mamba/KDA state routed by the mamba KV-cache group's block ids; degenerate +`remote_tp_size` normalized so writes fan out to all decode TP ranks; mamba N-1 +prefill/decode boundary (producer computes h(N-1), decoder recomputes token N). +These are load-bearing and remain in `patchers/`. + +## Diagnostic knobs (opt-in, default OFF) +Gated knobs from the investigation remain available for debugging (turn up/down +as needed): `K3_XFER_PROBE`, `K3_DECODE_RECV_PROBE`, `K3_KDA_STATE_PROBE`, +`K3_WRITE_BC`, `K3_HS_BC`, `K3_INPUTS_PROBE`, `K3_CHUNK_GATE_DEBUG`, +`K3_WRITE_FENCE`, `K3_ENABLE_CLAMP`. None change default behavior. See each +patcher's docstring and the README "Debugging" section. diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/benchmark_niah.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/benchmark_niah.py new file mode 100644 index 00000000..0cdd027e --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/benchmark_niah.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +# Needle-in-a-haystack long-context retrieval test. +# Adapted from vllm-project/vllm issue #47042 (GLM-5.2 sparse-MLA decode collapse), +# generalized to run against any OpenAI-compatible endpoint / model. +# +# Env: +# NIAH_URL endpoint (default http://127.0.0.1:30000/v1/chat/completions) +# NIAH_MODEL model name/tag the server serves (required — the served path) +# NIAH_WORDS comma list of context sizes in words (default 2000,8000,20000,35000) +# NIAH_MAXTOK max_tokens for the answer (default 2048) +# NIAH_TIMEOUT per-request timeout seconds (default 1800) +import os, sys, json, random, urllib.request + +URL = os.environ.get("NIAH_URL", "http://127.0.0.1:30000/v1/chat/completions") +MODEL = os.environ.get("NIAH_MODEL", "") +WORDS = [int(x) for x in os.environ.get("NIAH_WORDS", "2000,8000,20000,35000").split(",") if x.strip()] +MAXTOK = int(os.environ.get("NIAH_MAXTOK", "2048")) +TIMEOUT = float(os.environ.get("NIAH_TIMEOUT", "1800")) + +FILLER = ( + "table chair window bottle pencil garden river mountain coffee planet " + "engine guitar pillow ticket basket candle market silver button orange " + "rocket napkin ladder pepper carpet helmet jacket mirror anchor pocket " + "branch copper saddle tunnel violin wallet zipper meadow cactus pebble" +).split() +ANIMALS = ["elephant", "giraffe", "kangaroo", "penguin", "dolphin", + "tiger", "rhinoceros", "octopus", "crocodile", "panda"] + +SYSTEM = ( + "You read a word list and pick out the animals. Reply with a single " + "comma-separated list of lowercase animal names. Output nothing else." +) + + +def make_haystack(n_words, seed=0): + rng = random.Random(seed) + words = [rng.choice(FILLER) for _ in range(n_words)] + step = max(n_words // (len(ANIMALS) + 1), 1) + for i, animal in enumerate(ANIMALS): + words[min((i + 1) * step, len(words) - 1)] = animal + return " ".join(words) + + +def run(n_words): + body = { + "model": MODEL, + "messages": [ + {"role": "system", "content": SYSTEM}, + {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words)}, + ], + "temperature": 0.0, + "max_tokens": MAXTOK, + } + data = json.dumps(body).encode() + req = urllib.request.Request(URL, data=data, headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=TIMEOUT) as r: + msg = json.loads(r.read())["choices"][0]["message"] + except Exception as e: + print("words=%6d ERROR %s" % (n_words, e), flush=True) + return None + text = ((msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")).lower() + found = sorted(a for a in ANIMALS if a in text) + print("words=%6d found=%2d/10 %s" % (n_words, len(found), found), flush=True) + return len(found) + + +def main(): + if not MODEL: + print("NIAH_MODEL must be set (the served model path/name)", file=sys.stderr) + sys.exit(2) + print("=== NIAH retrieval test ===", flush=True) + print("url=%s model=%s sizes=%s" % (URL, MODEL, WORDS), flush=True) + results = {} + for n in WORDS: + results[n] = run(n) + print("=== NIAH summary ===", flush=True) + for n in WORDS: + v = results[n] + print(" words=%6d found=%s/10" % (n, "ERR" if v is None else v), flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/concurrency_bench.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/concurrency_bench.py new file mode 100644 index 00000000..7494f123 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/concurrency_bench.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +"""Concurrency / throughput bench for the EP16 2P/2D disagg serve. + +Fires N identical requests concurrently at the router and reports wall time, +throughput (req/s and output tok/s), and per-request latency percentiles. +This is the metric disagg is FOR (DP8 = 8 concurrent replicas), unlike +single-stream NIAH. + +Usage: concurrency_bench.py [max_out] + e.g. concurrency_bench.py 20000 8 64 +""" +import os, sys, json, time, urllib.request, concurrent.futures as cf + +# Router endpoint. Override with ROUTER_URL, e.g. +# ROUTER_URL=http://:30000/v1/completions +ROUTER = os.environ.get("ROUTER_URL", "http://127.0.0.1:30000/v1/completions") +FILLER = ("The quick brown fox jumps over the lazy dog near the riverbank while " + "the morning sun rises over the distant mountains and birds sing. ") + +def make_prompt(ctx_tokens): + body = (FILLER * (ctx_tokens * 4 // len(FILLER) + 1))[:ctx_tokens * 4] + return body + "\n\nSummarize the above in one word:" + +def one(prompt, max_out): + data = json.dumps({"model": "kimi-k3", "prompt": prompt, + "max_tokens": max_out, "temperature": 0}).encode() + req = urllib.request.Request(ROUTER, data=data, + headers={"Content-Type": "application/json"}) + t0 = time.time() + try: + with urllib.request.urlopen(req, timeout=1800) as r: + out = json.load(r) + dt = time.time() - t0 + n = out.get("usage", {}).get("completion_tokens", max_out) + return dt, n, True + except Exception as e: + return time.time() - t0, 0, False + +def pct(xs, p): + if not xs: return 0.0 + xs = sorted(xs); i = min(len(xs) - 1, int(p / 100 * len(xs))) + return xs[i] + +def main(): + ctx = int(sys.argv[1]) if len(sys.argv) > 1 else 20000 + conc = int(sys.argv[2]) if len(sys.argv) > 2 else 8 + max_out = int(sys.argv[3]) if len(sys.argv) > 3 else 64 + prompt = make_prompt(ctx) + print(f"ctx={ctx}tok concurrency={conc} max_out={max_out}", flush=True) + t0 = time.time() + with cf.ThreadPoolExecutor(max_workers=conc) as ex: + res = list(ex.map(lambda _: one(prompt, max_out), range(conc))) + wall = time.time() - t0 + lats = [d for d, n, ok in res if ok] + okn = sum(1 for _, _, ok in res if ok) + outtok = sum(n for _, n, ok in res if ok) + print(f" ok={okn}/{conc} wall={wall:.1f}s", flush=True) + print(f" throughput: {okn/wall:.3f} req/s | {outtok/wall:.1f} out-tok/s", flush=True) + print(f" latency: mean={sum(lats)/len(lats):.1f}s p50={pct(lats,50):.1f}s " + f"p99={pct(lats,99):.1f}s min={min(lats):.1f}s max={max(lats):.1f}s", flush=True) + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/load_image.sh b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/load_image.sh new file mode 100755 index 00000000..f1bfa7e9 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/load_image.sh @@ -0,0 +1,55 @@ +#!/bin/bash +# Self-restoring loader for the K3 disagg image. Ensures TAG is present on this node, +# trying sources in order: +# 1. already present locally -> done +# 2. docker pull from a registry you control (set HUB_IMAGE + DOCKER_USER/DOCKER_PAT) +# 3. docker load from a tar (offline fallback; set TAR) +# 4. build from Dockerfile.kimik3_disagg (last resort; needs base image + gh token) +# Usage: bash load_image.sh (env: TAG, HUB_IMAGE, DOCKER_USER, DOCKER_PAT, TAR) +# +# The image is built from Dockerfile.kimik3_disagg (see README). Push it to your own +# registry and point HUB_IMAGE at it, or rely on the local build in step 4. +set -euo pipefail + +TAG="${TAG:-kimik3-wideep-disagg:latest}" +HUB_IMAGE="${HUB_IMAGE:-}" # e.g. /kimik3-wideep-disagg:latest +DOCKER_USER="${DOCKER_USER:-}" +DOCKER_PAT="${DOCKER_PAT:-}" # set to enable authed pull +TAR="${TAR:-}" # optional: path to a saved image tar + +_have() { docker image inspect "$1" >/dev/null 2>&1; } + +# 1. already present (either the local tag or the hub tag) +if _have "$TAG"; then echo "[load] $(hostname -s): $TAG already present"; exit 0; fi +if [ -n "$HUB_IMAGE" ] && _have "$HUB_IMAGE"; then + docker tag "$HUB_IMAGE" "$TAG" 2>/dev/null || true + echo "[load] $(hostname -s): $HUB_IMAGE present -> tagged $TAG"; exit 0 +fi + +# 2. pull from a registry you control (only if HUB_IMAGE is set) +if [ -n "$HUB_IMAGE" ]; then + echo "[load] $(hostname -s): pulling $HUB_IMAGE ..." + if [ -n "$DOCKER_PAT" ]; then + echo "$DOCKER_PAT" | docker login -u "$DOCKER_USER" --password-stdin >/dev/null 2>&1 || true + fi + if docker pull "$HUB_IMAGE" 2>&1 | tail -1; then + if _have "$HUB_IMAGE"; then + docker tag "$HUB_IMAGE" "$TAG" 2>/dev/null || true + echo "[load] $(hostname -s): pulled + tagged $TAG"; exit 0 + fi + fi +fi + +# 3. offline tar fallback +if [ -f "$TAR" ]; then + echo "[load] $(hostname -s): pull failed; loading $TAR ..." + docker load -i "$TAR" 2>&1 | tail -1 + _have "$TAG" && { echo "[load] $(hostname -s): loaded $TAG from tar"; exit 0; } +fi + +# 4. build from source (last resort) +echo "[load] $(hostname -s): ERROR: could not obtain $TAG from hub or tar." +echo " Rebuild: docker build -f Dockerfile.kimik3_disagg \\" +echo " --build-arg MORI_REF=v1.2.2 --build-arg WITH_NIXL=0 \\" +echo " --build-arg GH_TOKEN=\$(gh auth token) -t $TAG ." +exit 1 diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/niah_probe.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/niah_probe.py new file mode 100644 index 00000000..8179e151 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/niah_probe.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Lightweight NIAH (needle-in-a-haystack) + basic accuracy probe for a served +Kimi-K3 endpoint. No deps beyond stdlib (urllib). Run from any node that can +reach the head IP:port. + +Usage: + python3 niah_probe.py --url http://:8000 --model kimi-k3 \ + --ctx 8000 --depths 0.1,0.5,0.9 +""" +import argparse, json, time, urllib.request, urllib.error + +FILLER = ("The quick brown fox jumps over the lazy dog. " + "Grand vistas of rolling hills stretched to the horizon. ") + +def chat(url, model, prompt, max_tokens=256, timeout=180): + body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, "temperature": 0, + }).encode() + req = urllib.request.Request(url.rstrip("/") + "/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + dt = time.time() - t0 + ch = d["choices"][0]["message"] + txt = (ch.get("content") or "") + " " + (ch.get("reasoning") or "") + return txt.strip(), d.get("usage", {}), dt + +def make_haystack(approx_words, needle, depth): + # Build filler to approx word count; insert needle at fractional depth. + words = [] + while len(words) < approx_words: + words.extend(FILLER.split()) + n = len(words) + pos = int(n * depth) + needle_words = needle.split() + words[pos:pos] = needle_words + return " ".join(words) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", required=True) + ap.add_argument("--model", default="kimi-k3") + ap.add_argument("--ctx", type=int, default=8000, help="approx haystack words") + ap.add_argument("--ctx-list", default=None, + help="comma list of sizes to sweep (overrides --ctx), e.g. " + "10000,50000,100000,200000,300000") + ap.add_argument("--depths", default="0.1,0.5,0.9") + ap.add_argument("--timeout", type=float, default=600, help="per-request seconds") + args = ap.parse_args() + + print(f"== sanity ==") + txt, usage, dt = chat(args.url, args.model, "Reply with only: OK", max_tokens=8) + print(f" resp={txt!r} usage={usage} {dt:.1f}s") + + secret = "The secret passcode is HELIOTROPE-7492." + question = ("\n\nQuestion: What is the secret passcode? " + "Answer with only the passcode.") + sizes = [int(x) for x in args.ctx_list.split(",")] if args.ctx_list else [args.ctx] + depths = [float(x) for x in args.depths.split(",")] + passed = 0; total = 0 + for ctx in sizes: + for depth in depths: + hay = make_haystack(ctx, secret, depth) + prompt = hay + question + try: + txt, usage, dt = chat(args.url, args.model, prompt, + max_tokens=256, timeout=args.timeout) + except Exception as e: + print(f" ctx={ctx:>7} depth={depth}: ERROR {e}", flush=True) + total += 1; continue + ok = "HELIOTROPE-7492" in txt.upper() + passed += ok; total += 1 + print(f" ctx={ctx:>7} depth={depth:>4}: {'PASS' if ok else 'FAIL'} " + f"(prompt_tok={usage.get('prompt_tokens')}, {dt:.1f}s) " + f"resp={txt[:50]!r}", flush=True) + print(f"== NIAH {passed}/{total} passed ==") + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_aiter_situv2_int4.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_aiter_situv2_int4.py new file mode 100644 index 00000000..e5849cac --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_aiter_situv2_int4.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Backport AITER PR #4471 (SiTUv2 in packed-int4 MoE stage1) SURGICALLY into the +grafted aiter's own files (keeps its imports; HEAD file-replace pulled missing +modules buffer_ops/vector). + +WHY: gfx942 requantizes K3 MoE to packed-int4; the grafted aiter's FlyDSL stage1 +kernel (compile_moe_gemm1, _abi3) hardcodes SiLU and ignores the requested SiTUv2 +-> K3 MoE silently computes SiLU -> gibberish. The upper plumbing already threads +act/situ_beta/situ_linear_beta (compile_flydsl_moe_stage1 has the params, fused_moe +passes them); only (a) the int4_bf16 branch's compile_moe_gemm1 call doesn't forward +them, and (b) the kernel codegen lacks the SiTUv2 epilogue. This applies exactly +those two things as in-place edits. + +5 anchor-based edits, idempotent, py_compile-checked. Also wipes stale _abi3 JIT. +Usage: apply_kimik3_aiter_situv2_int4.py (aiter resolved from import) +""" +import os +import shutil +import sys + + +def _edit(path, subs, tag): + src = open(path).read() + orig = src + for old, new in subs: + if new.split("\n", 1)[0] in src and old not in src: + # already applied (new present, old gone) + continue + if old not in src: + print(f"[k3-situv2] {tag}: anchor NOT found:\n {old[:70]!r}", file=sys.stderr) + return False + src = src.replace(old, new, 1) + if src == orig: + print(f"[k3-situv2] {tag}: no change (already applied).") + return True + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-situv2] {tag}: compile failed, rolled back: {e}", file=sys.stderr) + return False + print(f"[k3-situv2] {tag}: applied.") + return True + + +def main(): + try: + import aiter + A = os.path.dirname(aiter.__file__) + except Exception as e: + print(f"[k3-situv2] cannot import aiter: {e} -- skip.") + return 0 + + mg = os.path.join(A, "ops/flydsl/kernels/moe_gemm_2stage.py") + mk = os.path.join(A, "ops/flydsl/moe_kernels.py") + if not (os.path.isfile(mg) and os.path.isfile(mk)): + print("[k3-situv2] aiter flydsl files not found -- skip.") + return 0 + + # idempotency: kernel already has _abi4 + situv2 + if "_abi4" in open(mg).read() and "def situv2(" in open(mg).read(): + print("[k3-situv2] already applied (_abi4 + situv2).") + return 0 + + # ---- moe_kernels.py: forward act/situ_beta/situ_linear_beta to the int4_bf16 + # compile_moe_gemm1 call ---- + mk_subs = [( + " use_cshuffle_epilog=_use_cshuffle,\n" + " scale_is_bf16=True,\n" + " k_batch=k_batch,\n" + " )\n" + " else:\n" + " raise ValueError(\n" + " f\"Unsupported stage1 dtype combination: a_dtype={a_dtype}, b_dtype={b_dtype}\"\n", + " use_cshuffle_epilog=_use_cshuffle,\n" + " scale_is_bf16=True,\n" + " k_batch=k_batch,\n" + " act=act,\n" + " situ_beta=situ_beta,\n" + " situ_linear_beta=situ_linear_beta,\n" + " )\n" + " else:\n" + " raise ValueError(\n" + " f\"Unsupported stage1 dtype combination: a_dtype={a_dtype}, b_dtype={b_dtype}\"\n", + )] + if not _edit(mk, mk_subs, "moe_kernels.py"): + return 1 + + # ---- moe_gemm_2stage.py: 4 edits ---- + mg_subs = [] + + # (1) helper fn before compile_moe_gemm1 + the 3 new params + mg_subs.append(( + "@functools.lru_cache(maxsize=1024)\n" + "def compile_moe_gemm1(\n", + "def _stage1_activation_module_tag(\n" + " act: str, situ_beta: float, situ_linear_beta: float\n" + ") -> str:\n" + " \"\"\"Filesystem-safe cache-key suffix for stage1 activation code.\"\"\"\n" + " if act == \"silu\":\n" + " return \"_silu\"\n" + " def float_tag(value: float) -> str:\n" + " return float(value).hex().replace(\"-\", \"m\").replace(\"+\", \"p\").replace(\".\", \"d\")\n" + " return f\"_situv2_sb{float_tag(situ_beta)}_slb{float_tag(situ_linear_beta)}\"\n" + "\n" + "\n" + "@functools.lru_cache(maxsize=1024)\n" + "def compile_moe_gemm1(\n", + )) + mg_subs.append(( + " scale_is_bf16: bool = False,\n" + " k_batch: int = 1,\n" + "):\n" + " \"\"\"Compile stage1 kernel", + " scale_is_bf16: bool = False,\n" + " k_batch: int = 1,\n" + " act: str = \"silu\",\n" + " situ_beta: float = 1.0,\n" + " situ_linear_beta: float = 1.0,\n" + "):\n" + " if act not in (\"silu\", \"situv2\"):\n" + " raise ValueError(f\"act must be 'silu' or 'situv2', got {act!r}\")\n" + " if act == \"situv2\":\n" + " if situ_beta <= 0.0:\n" + " raise ValueError(f\"situ_beta must be > 0, got {situ_beta!r}\")\n" + " if situ_linear_beta <= 0.0:\n" + " raise ValueError(f\"situ_linear_beta must be > 0, got {situ_linear_beta!r}\")\n" + " \"\"\"Compile stage1 kernel", + )) + + # (2) module_name tag: _abi3 -> _abi4 + _act_tag, and bind the name + mg_subs.append(( + " _split_k_tag = f\"_splitk{k_batch}\" if _is_splitk else \"\"\n" + " (\n" + " f\"mfma_moe1_{in_dtype}_{out_dtype}_{epilog_tag}\"\n" + " f\"_t{tile_m}x{tile_n}x{tile_k}\"\n" + " f\"{_gs_tag}{scale_tag}{_split_k_tag}\"\n" + " f\"_abi3\" # also mask sentinel token ids on loads (X/scale_x) to avoid illegal address faults\n" + " ).replace(\"-\", \"_\")\n", + " _split_k_tag = f\"_splitk{k_batch}\" if _is_splitk else \"\"\n" + " _act_tag = _stage1_activation_module_tag(act, situ_beta, situ_linear_beta)\n" + " module_name = (\n" + " f\"mfma_moe1_{in_dtype}_{out_dtype}_{epilog_tag}\"\n" + " f\"_t{tile_m}x{tile_n}x{tile_k}\"\n" + " f\"{_gs_tag}{scale_tag}{_split_k_tag}{_act_tag}\"\n" + " f\"_abi4\" # also mask sentinel token ids on loads (X/scale_x) to avoid illegal address faults\n" + " ).replace(\"-\", \"_\")\n", + )) + # (2b) name the kernel (so distinct activations don't collide in the JIT cache) + mg_subs.append(( + " @flyc.kernel\n" + " def moe_gemm1(\n", + " @flyc.kernel(name=module_name)\n" + " def moe_gemm1(\n", + )) + + # (3) rewrite the silu() def into sigmoid/silu/situv2/apply_activation + mg_subs.append(( + " def silu(x):\n" + " # device fast path:\n" + " # emu = exp(-x) ~= exp2(log2e * (-x)) -> v_exp_f32\n" + " # sig = rcp(1 + emu) -> v_rcp_f32\n" + " # y = x * sig\n" + " #\n" + " # Using llvm.amdgcn intrinsics prevents lowering to the div_scale/div_fixup\n" + " # sequences that introduce extra compares/cndmasks.\n" + " t = x * (-1.4426950408889634) # -log2(e)\n" + " emu = rocdl.exp2(T.f32, t)\n" + " den = 1.0 + emu\n" + " sig = rocdl.rcp(T.f32, den)\n" + " return x * sig\n", + " def sigmoid(x):\n" + " t = x * (-1.4426950408889634) # -log2(e)\n" + " emu = rocdl.exp2(T.f32, t)\n" + " den = 1.0 + emu\n" + " return rocdl.rcp(T.f32, den)\n" + "\n" + " def silu(x):\n" + " return x * sigmoid(x)\n" + "\n" + " def situv2(gate, up):\n" + " gate_tanh = 2.0 * sigmoid(2.0 * (gate / situ_beta)) - 1.0\n" + " up_tanh = 2.0 * sigmoid(2.0 * (up / situ_linear_beta)) - 1.0\n" + " situ_gate = situ_beta * gate_tanh * sigmoid(gate)\n" + " situ_up = situ_linear_beta * up_tanh\n" + " return situ_gate * situ_up\n" + "\n" + " def apply_activation(gate, up):\n" + " if const_expr(act == \"silu\"):\n" + " return silu(gate) * up\n" + " return situv2(gate, up)\n", + )) + + # (4) both apply sites: y = silu(vg) * vu -> y = apply_activation(vg, vu) + # replace_all-style: do it manually since there are two identical occurrences + src = open(mg).read() + orig = src + src = src.replace("y = silu(vg) * vu", "y = apply_activation(vg, vu)") + if src != orig: + open(mg, "w").write(src) + + if not _edit(mg, mg_subs, "moe_gemm_2stage.py"): + return 1 + + # final compile check both + try: + import py_compile + py_compile.compile(mg, doraise=True) + py_compile.compile(mk, doraise=True) + except Exception as e: + print(f"[k3-situv2] final compile failed: {e}", file=sys.stderr) + return 1 + + # wipe stale _abi3 JIT + for cache in ("/opt/vllm_cache/aiter", "/opt/vllm_cache/aiter_jit", "/root/.aiter"): + shutil.rmtree(cache, ignore_errors=True) + + print("[k3-situv2] backported #4471 SiTUv2 packed-int4 stage1 in-place (_abi4); " + "wiped stale JIT.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunk_gate_fix.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunk_gate_fix.py new file mode 100644 index 00000000..905cbf8b --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunk_gate_fix.py @@ -0,0 +1,221 @@ +#!/usr/bin/env python3 +"""FIX chunked-prefill final-chunk detection using FRESH per-step COMPUTE progress. + +ROOT CAUSE (proven by debug): the connector's chunked-prefill gates detect the +"final chunk" by BLOCK COUNT: + num_prompt_tokens > len(block_ids) * self.block_size (defer if True) +But the effective scheduler block_size is ~5760 (mamba-page-padded attention +page), so a <=5760-token prompt occupies ONE block. Chunked prefill still +computes it in multiple passes (max_num_batched_tokens=2048 at a time), yet the +block-count gate is already satisfied after chunk 1 (1*5760 >= 2611) -> the KV +transfer fires with only ~2048 tokens computed -> tokens past chunk 1 (the +needle at the end) are never transferred -> long-context recall dies exactly at +max_num_batched_tokens. When the prompt fits in <=1 block, block-count can NEVER +detect chunk completion. + +The ONLY robust signal is COMPUTE progress. req.num_computed_tokens on the +producer Request is STALE (scheduler-process copy, =0), so we must use the FRESH +per-step data the scheduler passes into build_connector_meta: + * scheduler_output.num_scheduled_tokens[req_id] (scheduled THIS step) + * scheduled_new_reqs[].num_computed_tokens (new req, before step) + * scheduled_cached_reqs.num_computed_tokens[i] (cached req, before step) +A request's prefill is COMPLETE this step when: + computed_before + scheduled_this_step >= num_prompt_tokens - SLACK +SLACK (default 2, env K3_CHUNK_GATE_SLACK) absorbs the mamba N-1 prompt +truncation (P-side does request.num_prompt_tokens -= 1). Deferring until this is +true, then emitting the full accumulated block list, transfers ALL chunks' KV. + +Implementation: a helper builds a {req_id: (computed_before, scheduled)} map from +scheduler_output at the top of build_connector_meta, stored on self so both gates +(entry defer + accumulation final-detect) use it. Falls back to the original +block-count gate if the data is missing. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_chunk_gate_fix.py +""" +import os, sys +MARK = "k3-chunk-gate" +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + +# --- Part 1: build the fresh-progress map at the top of build_connector_meta, +# right after the WRITE/PRODUCER guard. Anchor on the comment block that opens +# the producer branch. +A_OLD = ( +" if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER:\n" +" # This is the logic for checking against chunked prefill.\n" +" # When the last chunk is identified,\n" +" # It places the request metadata into the saving queue.\n" +) +A_NEW = ( +" if self.mode == MoRIIOMode.WRITE and get_role() == ROLE.PRODUCER:\n" +" # " + MARK + ": build FRESH per-step compute-progress map. Block-count\n" +" # gates fail when the whole prompt fits in <=1 block (bs~5760);\n" +" # compute progress is the only reliable final-chunk signal.\n" +" self._k3_prog = {}\n" +" try:\n" +" _k3_nst = scheduler_output.num_scheduled_tokens\n" +" for _k3_nr in getattr(scheduler_output, 'scheduled_new_reqs', []) or []:\n" +" _k3_rid = getattr(_k3_nr, 'req_id', None)\n" +" if _k3_rid is not None:\n" +" self._k3_prog[_k3_rid] = (\n" +" int(getattr(_k3_nr, 'num_computed_tokens', 0)),\n" +" int(_k3_nst.get(_k3_rid, 0)),\n" +" )\n" +" _k3_cr = scheduler_output.scheduled_cached_reqs\n" +" for _k3_ci, _k3_crid in enumerate(_k3_cr.req_ids):\n" +" self._k3_prog[_k3_crid] = (\n" +" int(_k3_cr.num_computed_tokens[_k3_ci]),\n" +" int(_k3_nst.get(_k3_crid, 0)),\n" +" )\n" +" except Exception:\n" +" self._k3_prog = {}\n" +" # This is the logic for checking against chunked prefill.\n" +" # When the last chunk is identified,\n" +" # It places the request metadata into the saving queue.\n" +) + +# Helper appended once to compute final-ness from the map. +HELPER_ANCHOR = "class MoRIIOConnector" +HELPER_CODE = ( +"def _k3_prefill_done(self, req_id, req): # " + MARK + "\n" +" import os as _os\n" +" prog = getattr(self, '_k3_prog', None)\n" +" if not prog or req_id not in prog:\n" +" return None\n" +" cb, st = prog[req_id]\n" +" slack = int(_os.environ.get('K3_CHUNK_GATE_SLACK', '2'))\n" +" return (cb + st) >= (int(req.num_prompt_tokens) - slack)\n" +"\n\n" +) + +# --- Part 2a: ENTRY defer gate (_reqs_need_save loop) --- +B_OLD = ( +" for req_id, (req, block_ids) in self._reqs_need_save.items():\n" +" kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})\n" +" if req.num_prompt_tokens > len(block_ids) * self.block_size:\n" +" # not last chunk prefill\n" +" self._reqs_need_pending_save[req_id] = (req, block_ids)\n" +" continue\n" +) +B_NEW = ( +" for req_id, (req, block_ids) in self._reqs_need_save.items():\n" +" kv_params = self._req_kv_params.get(req_id, req.kv_transfer_params or {})\n" +" # " + MARK + ": defer by COMPUTE progress (fresh per-step), not block count.\n" +" _k3done = _k3_prefill_done(self, req_id, req)\n" +" if _k3done is None:\n" +" _k3done = not (req.num_prompt_tokens > len(block_ids) * self.block_size)\n" +" if os.environ.get('K3_CHUNK_GATE_DEBUG', '0') == '1':\n" +" import logging as _k3be_lg\n" +" _k3be_lg.getLogger(__name__).warning(\n" +" '[" + MARK + "-entry] req=%s nblk=%d npt=%d done=%s prog=%s',\n" +" req_id, len(block_ids), int(req.num_prompt_tokens), _k3done,\n" +" getattr(self, '_k3_prog', {}).get(req_id))\n" +" if not _k3done:\n" +" # not last chunk prefill\n" +" self._reqs_need_pending_save[req_id] = (req, block_ids)\n" +" continue\n" +) + +# --- Part 2b: accumulation final-detect gate (scheduled_cached_reqs loop) --- +C_OLD = ( +" if (\n" +" len(self._reqs_need_pending_save[req_id][1]) * self.block_size\n" +" >= req.num_prompt_tokens\n" +" ):\n" +) +C_NEW = ( +" _k3done2 = _k3_prefill_done(self, req_id, req) # " + MARK + "\n" +" if _k3done2 is None:\n" +" _k3done2 = (\n" +" len(self._reqs_need_pending_save[req_id][1]) * self.block_size\n" +" >= req.num_prompt_tokens\n" +" )\n" +" if os.environ.get('K3_CHUNK_GATE_DEBUG', '0') == '1':\n" +" import logging as _k3cg_lg\n" +" _k3cg_lg.getLogger(__name__).warning(\n" +" '[" + MARK + "-accum] req=%s nblk=%d npt=%d done=%s prog=%s',\n" +" req_id, len(self._reqs_need_pending_save[req_id][1]),\n" +" int(req.num_prompt_tokens), _k3done2,\n" +" getattr(self, '_k3_prog', {}).get(req_id))\n" +" if _k3done2:\n" +) + +# --- Part 2c: POST-LOOP SWEEP (gate D). The accumulation loop above is nested +# under `if new_block_ids is not None:`. When the FINAL prefill chunk allocates NO +# new block (whole prompt <=1 block, bs~5760), new_block_ids is None on that +# chunk -> the deferred req is never re-examined -> stuck in +# _reqs_need_pending_save -> unmap MISS deadlock. Sweep after the loop: emit any +# pending req whose prefill completed THIS step (per the compute-progress map). +D_OLD = ( +" del self._reqs_need_pending_save[req_id]\n" +"\n" +" # Loop through scheduled reqs and convert to ReqMeta.\n" +" for req_id, (req, block_ids) in self._reqs_need_recv.items():\n" +) +D_NEW = ( +" del self._reqs_need_pending_save[req_id]\n" +"\n" +" # " + MARK + "-sweep: emit deferred reqs whose FINAL chunk added no\n" +" # new block (accum loop skipped them) but whose prefill is now done.\n" +" try:\n" +" for _k3sw_rid in list(self._reqs_need_pending_save.keys()):\n" +" _k3sw_req, _k3sw_bl = self._reqs_need_pending_save[_k3sw_rid]\n" +" _k3sw_done = _k3_prefill_done(self, _k3sw_rid, _k3sw_req)\n" +" if os.environ.get('K3_CHUNK_GATE_DEBUG', '0') == '1':\n" +" import logging as _k3sw_lg\n" +" _k3sw_lg.getLogger(__name__).warning(\n" +" '[" + MARK + "-sweep] req=%s nblk=%d npt=%d done=%s prog=%s',\n" +" _k3sw_rid, len(_k3sw_bl), int(_k3sw_req.num_prompt_tokens),\n" +" _k3sw_done, getattr(self, '_k3_prog', {}).get(_k3sw_rid))\n" +" if not _k3sw_done:\n" +" continue\n" +" _k3sw_kv = self._req_kv_params.pop(\n" +" _k3sw_rid, _k3sw_req.kv_transfer_params or {}\n" +" )\n" +" meta.add_new_req(\n" +" request_id=_k3sw_rid,\n" +" local_block_ids=self._reqs_need_pending_save[_k3sw_rid][1],\n" +" kv_transfer_params=_k3sw_kv,\n" +" write_mode=True,\n" +" mamba_local_block_ids=self._reqs_save_mamba.get(_k3sw_rid, []),\n" +" all_group_block_ids=self._reqs_save_allgrp.get(_k3sw_rid, None),\n" +" )\n" +" del self._reqs_need_pending_save[_k3sw_rid]\n" +" except Exception as _k3sw_e:\n" +" import logging as _k3sw_lg2\n" +" _k3sw_lg2.getLogger(__name__).warning('[" + MARK + "-sweep] err %s', _k3sw_e)\n" +"\n" +" # Loop through scheduled reqs and convert to ReqMeta.\n" +" for req_id, (req, block_ids) in self._reqs_need_recv.items():\n" +) + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr); return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied."); return 0 + for old, tag in [(A_OLD, "A progress-map"), (B_OLD, "B entry-defer"), + (C_OLD, "C final-detect"), (D_OLD, "D post-loop-sweep"), + (HELPER_ANCHOR, "helper-anchor")]: + if old not in src: + print(f"[{MARK}] ANCHOR MISSING ({tag})", file=sys.stderr); return 1 + # insert module-level helper before the first class definition + src = src.replace(HELPER_ANCHOR, HELPER_CODE + HELPER_ANCHOR, 1) + src = (src.replace(A_OLD, A_NEW, 1).replace(B_OLD, B_NEW, 1) + .replace(C_OLD, C_NEW, 1).replace(D_OLD, D_NEW, 1)) + open(path, "w").write(src) + try: + import py_compile; py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr); return 1 + print(f"[{MARK}] applied (compute-progress gate).") + return 0 + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunked_allgrp.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunked_allgrp.py new file mode 100644 index 00000000..9eb57e36 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_chunked_allgrp.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""FIX: accumulate ALL kv-cache groups' block ids across CHUNKED PREFILL. + +Bug: MoRIIO connector accumulates only the ATTENTION-group block list across +prefill chunks (updated_blocks = existing + new_block_ids[0]); the per-group +all_group_block_ids (used by k3-group-routing) and mamba block ids are captured +ONCE at first chunk (update_state_after_alloc) and never grown. So for prompts +> max_num_batched_tokens (multi-chunk prefill), only the FIRST chunk's KV blocks +are advertised/transferred -> decode gets only ~first-chunk tokens of KV -> +needle beyond ~chunk1 lost -> long-context NIAH fails. + +Fix: in build_connector_meta's scheduled_cached_reqs loop, accumulate EVERY +group's new blocks (new_block_ids is a per-group tuple) into +self._reqs_save_allgrp[req_id] so the final-chunk add_new_req carries the FULL +per-group block lists. Depends on k3-group-routing (which adds _reqs_save_allgrp ++ all_group_block_ids threading). Gated by K3_GROUP_ROUTING via that patcher. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_chunked_allgrp.py +""" +import os, sys +MARK = "k3-chunked-allgrp" +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + +OLD = ( +" if new_block_ids is not None:\n" +" block_ids = new_block_ids[0]\n" +) +NEW = ( +" if new_block_ids is not None:\n" +" block_ids = new_block_ids[0]\n" +" # " + MARK + ": accumulate ALL groups' new blocks across\n" +" # chunked prefill so all_group_block_ids grows with the\n" +" # request (not frozen at chunk 1). new_block_ids is a\n" +" # per-group tuple.\n" +" try:\n" +" if req_id in getattr(self, '_reqs_save_allgrp', {}):\n" +" _k3ca_cur = self._reqs_save_allgrp[req_id]\n" +" _k3ca_new = []\n" +" for _k3ca_gi in range(len(_k3ca_cur)):\n" +" _k3ca_add = (\n" +" list(new_block_ids[_k3ca_gi])\n" +" if _k3ca_gi < len(new_block_ids) else []\n" +" )\n" +" _k3ca_new.append(\n" +" list(_k3ca_cur[_k3ca_gi]) + _k3ca_add\n" +" )\n" +" self._reqs_save_allgrp[req_id] = _k3ca_new\n" +" if req_id in getattr(self, '_reqs_save_mamba', {}) and len(_k3ca_new) > 1:\n" +" # keep mamba list (group 0..n-2 are mamba; group[-1] is MLA)\n" +" pass\n" +" except Exception:\n" +" pass\n" +) + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + if OLD not in src: + print(f"[{MARK}] ANCHOR MISSING", file=sys.stderr) + return 1 + src = src.replace(OLD, NEW, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_force_prefill_kda.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_force_prefill_kda.py new file mode 100644 index 00000000..2b2a8b56 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_force_prefill_kda.py @@ -0,0 +1,81 @@ +#!/usr/bin/env python3 +"""DIAGNOSTIC/FIX: route non-spec GDN tokens through the PREFILL kda kernel. + +Disagg boundary bug: the last prompt token (query_len=1) is classified as a +DECODE on the decode instance's first step, so its KDA output is computed by +fused_recurrent_kda_packed_decode using the TRANSFERRED conv/recurrent state. +That single-step decode kernel, fed a remote-transferred state, does not +reproduce prefill's chunk_kda output for that token -> the boundary token's +hidden state (hence its MLA retrieval QUERY) is wrong -> exact recall fails +while parametric knowledge/reasoning survive. vLLM's own GDN builder already +reclassifies decodes->prefills in the spec case ("the prefill kernel handles +1-token sequences with initial state correctly, producing identical results"). + +This forces num_decodes=0 (all non-spec -> prefill path) when env +K3_FORCE_PREFILL_KDA=1, so the boundary token uses the SAME chunk_kda kernel +prefill used. Diagnostic first: if recall becomes correct, the fix direction +(boundary-token-as-prefill) is confirmed; then we refine to boundary-only. + +Gated (default off) => byte-identical unless K3_FORCE_PREFILL_KDA=1. +Idempotent, anchor-based, two-pass, py_compile-checked. +Usage: apply_kimik3_force_prefill_kda.py +""" +import os +import sys + +MARK = "k3-force-prefill-kda" +REL = "v1/attention/backends/gdn_attn.py" + +OLD = ( + " num_decodes = (non_spec_query_lens_cpu == 1).sum().item()\n" + " # Exclude zero-length padded sequences from prefill count.\n" + " num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()\n" + " num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len\n" + " num_decode_tokens = num_decodes\n" +) +NEW = ( + " num_decodes = (non_spec_query_lens_cpu == 1).sum().item()\n" + " # Exclude zero-length padded sequences from prefill count.\n" + " num_zero_len = (non_spec_query_lens_cpu == 0).sum().item()\n" + " num_prefills = non_spec_query_lens_cpu.size(0) - num_decodes - num_zero_len\n" + " num_decode_tokens = num_decodes\n" + " import os as _k3fpk_os # " + MARK + "\n" + " if _k3fpk_os.environ.get('K3_FORCE_PREFILL_KDA', '0') == '1' and num_decodes > 0:\n" + " # " + MARK + ": route ALL non-spec 1-token seqs through the prefill\n" + " # chunk_kda kernel (with initial_state) so the disagg boundary token\n" + " # is bit-consistent with prefill. num_zero_len stays excluded.\n" + " num_prefills = non_spec_query_lens_cpu.size(0) - num_zero_len\n" + " num_decodes = 0\n" + " num_decode_tokens = 0\n" +) + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + if OLD not in src: + print(f"[{MARK}] ANCHOR MISSING", file=sys.stderr) + return 1 + src = src.replace(OLD, NEW, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_fwd_breadcrumb.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_fwd_breadcrumb.py new file mode 100644 index 00000000..48d6f1e1 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_fwd_breadcrumb.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Pre-forward breadcrumbs + cuda.synchronize to locate the disagg producer GPU fault +(no GPU debugger; sync attributes the async fault to the exact step). + +Gated by K3_FWD_BREADCRUMB=1. Active runner is v1/worker/gpu_model_runner.py. + [k3-bc] zero_block_ids SYNCED-ok / DONE+SYNCED (KV block zeroing) + [k3-bc] _model_forward START (pre-model synced-ok) / DONE+SYNCED (the model fwd) +The LAST breadcrumb before "Memory access fault" names the offending step, because +each is followed by torch.cuda.synchronize() which surfaces the async fault there. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_fwd_breadcrumb.py +""" +import os +import sys + + +def _patch(path, anchor, repl, tag): + if not os.path.isfile(path): + print(f"[k3-bc] {tag}: file not found -- skip.") + return 0 + src = open(path).read() + if repl.split("\n")[0].strip() and repl in src: + print(f"[k3-bc] {tag}: already applied.") + return 0 + if anchor not in src: + print(f"[k3-bc] {tag}: WARN anchor not found -- skip.") + return 0 + src2 = src.replace(anchor, repl, 1) + open(path, "w").write(src2) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(src) + print(f"[k3-bc] {tag}: ERROR compile: {e}", file=sys.stderr) + return 1 + print(f"[k3-bc] {tag}: applied.") + return 0 + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + root = sys.argv[1] + rc = 0 + gmr = os.path.join(root, "v1/worker/gpu_model_runner.py") + + # A) block-zeroing with syncs + a_anchor = ( + " if scheduler_output.new_block_ids_to_zero:\n" + " self._zero_block_ids(scheduler_output.new_block_ids_to_zero)\n" + ) + a_repl = ( + " if scheduler_output.new_block_ids_to_zero:\n" + " import os as _os_k3z\n" + " _k3z = _os_k3z.environ.get(\"K3_FWD_BREADCRUMB\", \"0\") == \"1\"\n" + " if _k3z:\n" + " import torch as _t_k3z, logging as _lg_k3z\n" + " _t_k3z.cuda.synchronize()\n" + " _lg_k3z.getLogger(__name__).warning(\n" + " \"[k3-bc] zero_block_ids SYNCED-ok ids=%s\",\n" + " scheduler_output.new_block_ids_to_zero)\n" + " self._zero_block_ids(scheduler_output.new_block_ids_to_zero)\n" + " if _k3z:\n" + " import torch as _t_k3z, logging as _lg_k3z\n" + " _t_k3z.cuda.synchronize()\n" + " _lg_k3z.getLogger(__name__).warning(\"[k3-bc] zero_block_ids DONE+SYNCED\")\n" + ) + rc |= _patch(gmr, a_anchor, a_repl, "zero_block_ids") + + # B) _model_forward with syncs + b_anchor = ( + " return self.model(\n" + " input_ids=input_ids,\n" + " positions=positions,\n" + " intermediate_tensors=intermediate_tensors,\n" + " inputs_embeds=inputs_embeds,\n" + " **model_kwargs,\n" + " )\n" + ) + b_repl = ( + " import os as _os_k3mf\n" + " _k3mf = _os_k3mf.environ.get(\"K3_FWD_BREADCRUMB\", \"0\") == \"1\"\n" + " if _k3mf:\n" + " import torch as _t_k3mf, logging as _lg_k3mf\n" + " _t_k3mf.cuda.synchronize()\n" + " _lg_k3mf.getLogger(__name__).warning(\"[k3-bc] _model_forward START (pre-model synced-ok)\")\n" + " _out_k3mf = self.model(\n" + " input_ids=input_ids,\n" + " positions=positions,\n" + " intermediate_tensors=intermediate_tensors,\n" + " inputs_embeds=inputs_embeds,\n" + " **model_kwargs,\n" + " )\n" + " if _k3mf:\n" + " import torch as _t_k3mf, logging as _lg_k3mf\n" + " _t_k3mf.cuda.synchronize()\n" + " _lg_k3mf.getLogger(__name__).warning(\"[k3-bc] _model_forward DONE+SYNCED\")\n" + " return _out_k3mf\n" + ) + rc |= _patch(gmr, b_anchor, b_repl, "model_forward") + + return rc + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_conv_debug.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_conv_debug.py new file mode 100644 index 00000000..2871f2b6 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_conv_debug.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Log KDA-forward entry metadata (debug) for the 2P/2D disagg producer fault. + +Gated by K3_KDA_CONV_DEBUG=1. Injected right after `m = attn_metadata_narrowed` +in kimi_gdn_linear_attn (fires for EVERY KDA layer, all sub-paths, before any +indexing). Prints conv_state/recurrent_state shapes, num_prefills/decodes, +num_actual_tokens, non_spec_state_indices min/max/count, has_initial_state +any/sum, and whether spec masks are set. The Python log flushes before the async +GPU kernel faults, so it captures the offending values on the N-1 partial prefill. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_conv_debug.py +""" +import os +import sys + +REL = "model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-kda-convdbg] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "k3-kda entry" in src: + print("[k3-kda-convdbg] already applied.") + return 0 + + anchor = ( + " m = attn_metadata_narrowed\n" + " has_initial_state = m.has_initial_state\n" + " non_spec_query_start_loc = m.non_spec_query_start_loc\n" + " non_spec_state_indices_tensor = m.non_spec_state_indices_tensor\n" + ) + inject = ( + " m = attn_metadata_narrowed\n" + " import os as _os_k3e\n" + " if _os_k3e.environ.get(\"K3_KDA_CONV_DEBUG\", \"0\") == \"1\":\n" + " try:\n" + " import logging as _lg_k3e\n" + " _cs = self.kv_cache[0] if self.kv_cache is not None else None\n" + " _rs = self.kv_cache[1] if self.kv_cache is not None else None\n" + " _nsi = m.non_spec_state_indices_tensor\n" + " _his = m.has_initial_state\n" + " _lg_k3e.getLogger(__name__).warning(\n" + " \"[k3-kda entry] layer=%s conv_state.shape=%s recur.shape=%s \"\n" + " \"n_prefills=%s n_decodes=%s num_actual=%s \"\n" + " \"nsi(min/max/n)=%s/%s/%s has_init(any/sum)=%s/%s \"\n" + " \"spec_masks=%s\",\n" + " getattr(self, \"prefix\", \"?\"),\n" + " (tuple(_cs.shape) if _cs is not None else None),\n" + " (tuple(_rs.shape) if _rs is not None else None),\n" + " getattr(m, \"num_prefills\", None), getattr(m, \"num_decodes\", None),\n" + " getattr(m, \"num_actual_tokens\", None),\n" + " (int(_nsi.min()) if _nsi is not None and _nsi.numel() else None),\n" + " (int(_nsi.max()) if _nsi is not None and _nsi.numel() else None),\n" + " (int(_nsi.numel()) if _nsi is not None else None),\n" + " (bool(_his.any()) if _his is not None else None),\n" + " (int(_his.sum()) if _his is not None else None),\n" + " (None if m.spec_sequence_masks is None else True),\n" + " )\n" + " except Exception as _e_k3e:\n" + " import logging as _lg_k3e\n" + " _lg_k3e.getLogger(__name__).warning(\"[k3-kda entry] err %s\", _e_k3e)\n" + " has_initial_state = m.has_initial_state\n" + " non_spec_query_start_loc = m.non_spec_query_start_loc\n" + " non_spec_state_indices_tensor = m.non_spec_state_indices_tensor\n" + ) + if anchor not in src: + print("[k3-kda-convdbg] WARN: entry anchor not found -- not applied.") + return 0 + src = src.replace(anchor, inject, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-kda-convdbg] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print("[k3-kda-convdbg] added KDA-entry debug logging (K3_KDA_CONV_DEBUG=1).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_fa_contiguous.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_fa_contiguous.py new file mode 100644 index 00000000..99a8cf13 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_fa_contiguous.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Make f_a contiguous before f_b_proj in the KDA forward (fixes disagg GPU fault). + +ROOT CAUSE (localized tool-free via synced breadcrumbs, top-down): + The 2P/2D disagg producer GPU-faults in KDA layer 0's forward, at + `g1 = self.f_b_proj(f_a)` (kimi_gdn_linear_attn.py). f_a is a NON-CONTIGUOUS + slice from `projected_qkvgfab.split(...)` of the padded in_proj output. Feeding + that strided view straight into f_b_proj's (bf16) GEMM faults on the small + disagg N-1 prefill shape. The model author already flagged an "Inductor + correctness issue with the row-strided G view" and padded in_proj to dodge it; + the tiny disagg batch defeats that workaround. + + Not a quant issue: mxfp4 falls all LinearBase layers back to + UnquantizedLinearMethod (only MoE experts are MXFP4); the KDA projections are + plain bf16 in the checkpoint. The fault is the strided-view GEMM input. + +FIX: `self.f_b_proj(f_a.contiguous())` -- materialize a dense f_a. Negligible cost +(f_a is [tokens, head_dim=128]). + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_fa_contiguous.py +""" +import os +import sys + +REL = "model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-fa] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "f_a.contiguous()" in src: + print("[k3-fa] already applied.") + return 0 + + anchor = " g1 = self.f_b_proj(f_a)[0]\n" + repl = ( + " # k3-kda: f_a is a non-contiguous slice from the padded in_proj\n" + " # split; the strided view faults f_b_proj's GEMM on the small disagg\n" + " # prefill shape. Materialize a dense f_a first.\n" + " g1 = self.f_b_proj(f_a.contiguous())[0]\n" + ) + if anchor not in src: + print("[k3-fa] WARN anchor not found -- not applied.") + return 0 + src = src.replace(anchor, repl, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-fa] ERROR compile: {e}", file=sys.stderr) + return 1 + print("[k3-fa] f_a.contiguous() before f_b_proj applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_guard.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_guard.py new file mode 100644 index 00000000..75f37a0c --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_guard.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Guard gather_initial_states against out-of-range KDA state indices. + +PROBLEM (the 2P/2D disagg producer GPU memory fault) + kimi_gdn_linear_attn calls gather_initial_states(recurrent_state, state_indices, + has_initial_state) at prefill. The Triton kernel forms a GPU address + state_ptr + state_idx*stride even where the value load is has_initial_state-masked; + an out-of-range state_idx therefore faults the device ("Memory access fault by GPU + ... Reason: Unknown"). Standalone prefill is fine (has_initial_state=False / valid + idx); only the disagg producer prefill faults -> it carries a state index that is + >= recurrent_state.shape[0] (or a stray value) for a request the scheduler treated + as having a prior state. + +FIX (surgical, in the op wrapper so it covers every caller) + Before launching the kernel, compute a safe index tensor: + - where has_initial_state is False -> 0 (a fresh prefill gathers nothing anyway) + - where True -> clamp into [0, n_state_blocks-1] + and log once if any index was out of range (so the upstream metadata bug is + visible). This keeps the GPU address valid; masked-off rows still contribute zeros. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_gather_guard.py +""" +import os +import sys + +REL = "model_executor/layers/mamba/ops/gather_initial_states.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-kda-gather] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + + if "k3-kda guard" in src: + print("[k3-kda-gather] already applied.") + return 0 + + anchor = ( + " row_size = state[0].numel()\n" + " # Mamba pages may pad stride(0), but each state row remains dense.\n" + " assert state[0].is_contiguous()\n" + " output = torch.empty(\n" + ) + repl = ( + " row_size = state[0].numel()\n" + " # Mamba pages may pad stride(0), but each state row remains dense.\n" + " assert state[0].is_contiguous()\n" + " # k3-kda guard: an out-of-range state index makes the kernel form an\n" + " # OOB GPU address (state_ptr + idx*stride) -> Memory access fault, even\n" + " # where the value load is has_initial_state-masked. Under 2P/2D disagg the\n" + " # producer prefill has been observed to carry indices >= state.shape[0];\n" + " # zero the effective index where has_initial_state is False and clamp any\n" + " # stray index into range so the address stays valid.\n" + " _n_state_blocks = int(state.shape[0])\n" + " _safe_idx = torch.where(\n" + " has_initial_state,\n" + " indices.to(torch.int64).clamp(0, _n_state_blocks - 1),\n" + " torch.zeros_like(indices, dtype=torch.int64),\n" + " )\n" + " if bool((indices >= _n_state_blocks).any()) or bool((indices < 0).any()):\n" + " import logging as _lg\n" + " _bad = indices[(indices >= _n_state_blocks) | (indices < 0)]\n" + " _lg.getLogger(__name__).warning(\n" + ' \"[k3-kda gather] clamped %d out-of-range state idx (n_blocks=%d, \"\n' + ' \"sample=%s); disagg producer prefill likely mis-flagged initial state.\",\n' + " int(_bad.numel()), _n_state_blocks, _bad[:8].tolist(),\n" + " )\n" + " indices = _safe_idx\n" + " output = torch.empty(\n" + ) + if anchor not in src: + print("[k3-kda-gather] WARN: anchor not found -- gather guard NOT applied.") + return 0 + src = src.replace(anchor, repl, 1) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-kda-gather] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print("[k3-kda-gather] guarded gather_initial_states against OOB state idx.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_nosync.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_nosync.py new file mode 100644 index 00000000..dfbd4ef4 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_gather_nosync.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +"""PERF/HANG FIX: make the KDA gather_initial_states OOB guard sync-free. + +The k3-kda guard in +`vllm/model_executor/layers/mamba/ops/gather_initial_states.py` clamps state +indices into range (sync-free, on-device — correct) but then runs a diagnostic +`if bool((indices >= n).any()) or bool((indices < 0).any()):` block purely to log +a warning. Each `bool(...any())` forces a device->CPU sync (`_local_scalar_dense` +-> `memcpy_and_sync`), i.e. a full-stream drain — **per KDA layer, per prefill +chunk**. With 69 KDA layers and chunked prefill (~365 chunks at 750K / +batched=2048) that is ~25k forced syncs, which stalls so hard it presents as a +hang for contexts > ~500K (py-spy'd: DP rank stuck in this sync while all other +DP ranks wait at the batch-coordination all_reduce). + +The clamp already guarantees a valid GPU address, so the diagnostic sync is pure +overhead. Gate it behind `K3_KDA_GATHER_LOG=1` (default OFF) so the hot path never +syncs. Correctness is unchanged (indices are still clamped). + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_gather_nosync.py +""" +import os, sys + +MARK = "k3-kda-nosync" +REL = "model_executor/layers/mamba/ops/gather_initial_states.py" + +OLD = ( +" if bool((indices >= _n_state_blocks).any()) or bool((indices < 0).any()):\n" +" import logging as _lg\n" +" _bad = indices[(indices >= _n_state_blocks) | (indices < 0)]\n" +" _lg.getLogger(__name__).warning(\n" +" \"[k3-kda gather] clamped %d out-of-range state idx (n_blocks=%d, \"\n" +" \"sample=%s); disagg producer prefill likely mis-flagged initial state.\",\n" +" int(_bad.numel()), _n_state_blocks, _bad[:8].tolist(),\n" +" )\n" +) +NEW = ( +" import os as _k3os # " + MARK + ": clamp above is sync-free + safe; the\n" +" # bool(...any()) diagnostic below forces a device->CPU sync EVERY call (per\n" +" # KDA layer, per chunk) -> O(layers*chunks) stalls that hang ctx > ~500K.\n" +" # Gate behind K3_KDA_GATHER_LOG=1 (default OFF); hot path never syncs.\n" +" if _k3os.environ.get('K3_KDA_GATHER_LOG', '0') == '1':\n" +" if bool((indices >= _n_state_blocks).any()) or bool((indices < 0).any()):\n" +" import logging as _lg\n" +" _bad = indices[(indices >= _n_state_blocks) | (indices < 0)]\n" +" _lg.getLogger(__name__).warning(\n" +" \"[k3-kda gather] clamped %d out-of-range state idx (n_blocks=%d, \"\n" +" \"sample=%s); disagg producer prefill likely mis-flagged initial state.\",\n" +" int(_bad.numel()), _n_state_blocks, _bad[:8].tolist(),\n" +" )\n" +) + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr); return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied."); return 0 + if OLD not in src: + print(f"[{MARK}] ANCHOR MISSING", file=sys.stderr); return 1 + src = src.replace(OLD, NEW, 1) + open(path, "w").write(src) + try: + import py_compile; py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr); return 1 + print(f"[{MARK}] applied (KDA gather sync-free).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_internal_bc.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_internal_bc.py new file mode 100644 index 00000000..4b5901af --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kda_internal_bc.py @@ -0,0 +1,149 @@ +#!/usr/bin/env python3 +"""Synced breadcrumbs INSIDE KimiGatedDeltaNetAttention.forward to localize the +fault within KDA layer 0 (the last surviving breadcrumb was 'attnres layer 0 START +type=KimiGatedDeltaNetAttention', and the KDA-entry debug at line ~412 never fired, +so the fault is in the projections (334-377) or _forward metadata setup). + +Gated by K3_FWD_BREADCRUMB=1. Prints hidden_states shape/contiguity + in_proj +weight shape before the first GEMM, and after it. Edits +vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_internal_bc.py +""" +import os +import sys + +REL = "model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-kdabc] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "[k3-bc] KDA fwd START" in src: + print("[k3-kdabc] already applied.") + return 0 + + anchor = ( + " num_tokens = hidden_states.size(0)\n" + " projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]\n" + ) + repl = ( + " num_tokens = hidden_states.size(0)\n" + " import os as _os_k3k\n" + " _k3k = _os_k3k.environ.get(\"K3_FWD_BREADCRUMB\", \"0\") == \"1\"\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\n" + " \"[k3-bc] KDA fwd START prefix=%s hs=%s contig=%s in_proj_w=%s\",\n" + " getattr(self, \"prefix\", \"?\"), tuple(hidden_states.shape),\n" + " hidden_states.is_contiguous(),\n" + " tuple(self.in_proj_qkvgfab.weight.shape)\n" + " if hasattr(self.in_proj_qkvgfab, \"weight\") else \"?\")\n" + " projected_qkvgfab = self.in_proj_qkvgfab(hidden_states)[0]\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\"[k3-bc] KDA in_proj DONE+SYNCED\")\n" + ) + if anchor not in src: + print("[k3-kdabc] WARN anchor not found -- skip.") + return 0 + src = src.replace(anchor, repl, 1) + + # around self._forward(...) + o_proj + fwd_anchor = ( + " self._forward(\n" + " mixed_qkv=mixed_qkv,\n" + " g1=g1,\n" + " g2=g2,\n" + " beta=beta,\n" + " core_attn_out=core_attn_out,\n" + " )\n" + " core_attn_out = rearrange(core_attn_out, \"1 n h d -> n (h d)\")\n" + " output[:] = self.o_proj(core_attn_out)[0]\n" + ) + fwd_repl = ( + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\n" + " \"[k3-bc] KDA projections DONE -> _forward; mixed_qkv=%s g1=%s g2=%s beta=%s\",\n" + " tuple(mixed_qkv.shape), tuple(g1.shape), tuple(g2.shape), tuple(beta.shape))\n" + " self._forward(\n" + " mixed_qkv=mixed_qkv,\n" + " g1=g1,\n" + " g2=g2,\n" + " beta=beta,\n" + " core_attn_out=core_attn_out,\n" + " )\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\"[k3-bc] KDA _forward DONE+SYNCED\")\n" + " core_attn_out = rearrange(core_attn_out, \"1 n h d -> n (h d)\")\n" + " output[:] = self.o_proj(core_attn_out)[0]\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\"[k3-bc] KDA o_proj DONE+SYNCED\")\n" + ) + if fwd_anchor in src: + src = src.replace(fwd_anchor, fwd_repl, 1) + else: + print("[k3-kdabc] WARN _forward anchor not found.") + + # sub-projection breadcrumbs (split / g_a-g_b / f_b) + gp_anchor = ( + " g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]\n" + "\n" + " g1 = self.f_b_proj(f_a)[0]\n" + " beta = beta.unsqueeze(0)\n" + ) + gp_repl = ( + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\n" + " \"[k3-bc] KDA split DONE mixed_qkv=%s f_a=%s beta=%s\",\n" + " tuple(mixed_qkv.shape), tuple(f_a.shape), tuple(beta.shape))\n" + " g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\"[k3-bc] KDA g_a/g_b_proj DONE\")\n" + "\n" + " g1 = self.f_b_proj(f_a)[0]\n" + " if _k3k:\n" + " import torch as _t_k3k, logging as _lg_k3k\n" + " _t_k3k.cuda.synchronize()\n" + " _lg_k3k.getLogger(__name__).warning(\"[k3-bc] KDA f_b_proj DONE g1=%s\", tuple(g1.shape))\n" + " beta = beta.unsqueeze(0)\n" + ) + if gp_anchor in src: + src = src.replace(gp_anchor, gp_repl, 1) + else: + print("[k3-kdabc] WARN sub-proj anchor not found.") + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-kdabc] ERROR compile: {e}", file=sys.stderr) + return 1 + print("[k3-kdabc] added KDA-internal synced breadcrumbs.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kvzero_bounds.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kvzero_bounds.py new file mode 100644 index 00000000..658cdfa6 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_kvzero_bounds.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Bounds-guard the KV block-zeroing Triton kernel (the 2P/2D disagg producer fault). + +ROOT CAUSE (found by pre-forward breadcrumbs, no GPU tools): + execute_model -> _update_states zeros freshly-allocated attention blocks via + KVBlockZeroer.zero_block_ids -> _zero_kv_blocks_kernel, BEFORE any layer forward. + The kernel writes at seg_addr + block_id*page_size_el (v1/worker/utils.py). It + never bounds-checks block_id against the segment tensor's block capacity. Under + K3 hybrid (MLA + KDA) + 2P/2D disagg the scheduler emits an attention block id + that exceeds a segment's capacity -> OOB write -> "Memory access fault" on all + ranks, before the model forward. (Breadcrumb: "zero_block_ids DONE" prints, then + the fault; the kernel launch is async so the Python line logs first.) + +FIX: + Record each segment's logical block capacity (seg_nblocks) in KVBlockZeroer.__init__, + pass it to the kernel, and skip any (block, seg) whose block_id >= that capacity. + Skipping is correct: only FullAttention/MLA managers record ids for zeroing, so an + id past an MLA segment's capacity is spurious for that segment. + +Edits v1/worker/utils.py in 4 spots (kernel sig, kernel guard, __init__ capacity, +zero_block_ids meta+call). Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kvzero_bounds.py +""" +import os +import sys + +REL = "v1/worker/utils.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-kvzero] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "_K3MLASpec" in src: + print("[k3-kvzero] already applied.") + return 0 + + subs = [] + + # 0) THE FIX: skip MLAAttentionSpec in the zeroer (MLA subclasses + # FullAttentionSpec -> wrong stride math -> OOB write -> GPU fault). + subs.append(( + " for group in attn_groups_iter:\n" + " spec = group.kv_cache_spec\n" + " if not isinstance(spec, FullAttentionSpec):\n" + " continue\n" + " if group.kv_cache_group_id >= len(kernel_block_sizes):\n" + " continue\n", + " from vllm.v1.kv_cache_interface import MLAAttentionSpec as _K3MLASpec\n" + " _os_k3skip = __import__(\"os\").environ.get(\"K3_ZERO_SKIP_MLA\", \"1\") == \"1\"\n" + " for group in attn_groups_iter:\n" + " spec = group.kv_cache_spec\n" + " if not isinstance(spec, FullAttentionSpec):\n" + " continue\n" + " if _os_k3skip and isinstance(spec, _K3MLASpec):\n" + " continue\n" + " if group.kv_cache_group_id >= len(kernel_block_sizes):\n" + " continue\n", + )) + + # 1) kernel signature: add seg_nblocks_ptr before N_SEGS + subs.append(( + " block_ids_ptr,\n" + " n_blocks,\n" + " N_SEGS: tl.constexpr,\n", + " block_ids_ptr,\n" + " n_blocks,\n" + " seg_nblocks_ptr,\n" + " N_SEGS: tl.constexpr,\n", + )) + + # 2) kernel body: guard after loading block_id + subs.append(( + " block_id = tl.load(block_ids_ptr + block_index)\n" + " seg_addr = tl.load(seg_addrs_ptr + seg_index)\n", + " block_id = tl.load(block_ids_ptr + block_index)\n" + " # k3-kda: bounds-guard OOB block_id (would fault the GPU).\n" + " seg_nblk = tl.load(seg_nblocks_ptr + seg_index)\n" + " if block_id >= seg_nblk:\n" + " return\n" + " seg_addr = tl.load(seg_addrs_ptr + seg_index)\n", + )) + + # 3a) __init__: seg_nblocks list init + subs.append(( + " seg_addrs: list[int] = []\n" + " seg_page_sizes: list[int] = []\n", + " seg_addrs: list[int] = []\n" + " seg_page_sizes: list[int] = []\n" + " seg_nblocks: list[int] = [] # k3-kda: per-segment block capacity\n", + )) + + # 3b) __init__: compute + append capacity in the outer loop + subs.append(( + " outer_strides = [kv.stride(d) * el for d in outer_dims]\n" + " for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):\n" + " off_bytes = sum(i * s for i, s in zip(outer, outer_strides))\n" + " seg_addrs.append(dp + off_bytes)\n" + " seg_page_sizes.append(cur_page_el)\n", + " outer_strides = [kv.stride(d) * el for d in outer_dims]\n" + " seg_nblk = int(kv.shape[block_dim]) // max(1, ratio)\n" + " for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):\n" + " off_bytes = sum(i * s for i, s in zip(outer, outer_strides))\n" + " seg_addrs.append(dp + off_bytes)\n" + " seg_page_sizes.append(cur_page_el)\n" + " seg_nblocks.append(seg_nblk)\n", + )) + + # 3c) __init__: add seg_nblocks tensor to _meta + subs.append(( + " max_page_size_el // blk_size,\n" + " blk_size,\n" + " len(seg_addrs),\n" + " )\n", + " max_page_size_el // blk_size,\n" + " blk_size,\n" + " len(seg_addrs),\n" + " torch.tensor(seg_nblocks, dtype=torch.int64, device=self.device),\n" + " )\n", + )) + + # 4) zero_block_ids: unpack + pass seg_nblocks + subs.append(( + " seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs = self._meta\n", + " seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs, seg_nblocks = (\n" + " self._meta\n" + " )\n", + )) + subs.append(( + " idx,\n" + " n_blocks,\n" + " N_SEGS=n_segs,\n", + " idx,\n" + " n_blocks,\n" + " seg_nblocks,\n" + " N_SEGS=n_segs,\n", + )) + + for old, new in subs: + if old not in src: + print(f"[k3-kvzero] WARN anchor not found:\n{old[:80]!r}\n-- ABORT (no partial edits).") + return 0 + src = src.replace(old, new, 1) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-kvzero] ERROR compile failed: {e}", file=sys.stderr) + return 1 + print("[k3-kvzero] bounds-guarded _zero_kv_blocks_kernel (skip OOB block_id).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_layer_breadcrumb.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_layer_breadcrumb.py new file mode 100644 index 00000000..30d5f9d2 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_layer_breadcrumb.py @@ -0,0 +1,138 @@ +#!/usr/bin/env python3 +"""Per-layer synced breadcrumbs inside the K3 (AMD) model forward. + +Gated by K3_FWD_BREADCRUMB=1. Pinpoints whether the disagg-producer GPU fault is +in embed or a specific decoder layer (and its attn type: KDA vs MLA). Each point +is followed by torch.cuda.synchronize() so the async fault surfaces at the exact +layer. Edits vllm/models/kimi_k3/amd/linear.py KimiLinearModel.forward. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_layer_breadcrumb.py +""" +import os +import sys + +RELS = [ + "models/kimi_k3/amd/linear.py", + "../vllm/models/kimi_k3/amd/linear.py", +] + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + root = sys.argv[1] + path = None + for rel in RELS: + p = os.path.join(root, rel) + if os.path.isfile(p): + path = p + break + if path is None: + # search + for base, _, files in os.walk(os.path.join(root, "models")): + if "linear.py" in files and "kimi_k3" in base and base.endswith("amd"): + path = os.path.join(base, "linear.py") + break + if path is None or not os.path.isfile(path): + print("[k3-layerbc] amd/linear.py not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "[k3-bc] model.forward START" in src: + print("[k3-layerbc] already applied.") + return 0 + + embed_anchor = ( + " if get_pp_group().is_first_rank:\n" + " if inputs_embeds is not None:\n" + " hidden_states = inputs_embeds\n" + " else:\n" + " hidden_states = self.embed_input_ids(input_ids)\n" + " residual = None\n" + " else:\n" + " assert intermediate_tensors is not None\n" + " hidden_states = intermediate_tensors[\"hidden_states\"]\n" + " residual = intermediate_tensors[\"residual\"]\n" + ) + embed_repl = ( + " import os as _os_k3ml\n" + " _k3ml = _os_k3ml.environ.get(\"K3_FWD_BREADCRUMB\", \"0\") == \"1\"\n" + " if _k3ml:\n" + " import torch as _t_k3ml, logging as _lg_k3ml\n" + " _t_k3ml.cuda.synchronize()\n" + " _lg_k3ml.getLogger(__name__).warning(\n" + " \"[k3-bc] model.forward START input_ids=%s inputs_embeds=%s\",\n" + " (None if input_ids is None else tuple(input_ids.shape)),\n" + " (None if inputs_embeds is None else tuple(inputs_embeds.shape)))\n" + + embed_anchor + + " if _k3ml:\n" + " import torch as _t_k3ml, logging as _lg_k3ml\n" + " _t_k3ml.cuda.synchronize()\n" + " _lg_k3ml.getLogger(__name__).warning(\"[k3-bc] embed DONE+SYNCED\")\n" + ) + + layer_anchor = ( + " hidden_states, residual = layer(\n" + " positions=positions,\n" + " hidden_states=hidden_states,\n" + " residual=residual,\n" + " )\n" + ) + layer_repl = ( + " if _k3ml:\n" + " import torch as _t_k3ml, logging as _lg_k3ml\n" + " _t_k3ml.cuda.synchronize()\n" + " _lg_k3ml.getLogger(__name__).warning(\n" + " \"[k3-bc] layer %d START (type=%s)\", layer_idx,\n" + " type(getattr(layer, \"self_attn\", layer)).__name__)\n" + + layer_anchor + ) + + if embed_anchor not in src: + print("[k3-layerbc] WARN embed anchor not found -- skip.") + return 0 + src = src.replace(embed_anchor, embed_repl, 1) + if layer_anchor in src: + src = src.replace(layer_anchor, layer_repl, 1) + else: + print("[k3-layerbc] WARN layer anchor not found (embed bc still added).") + + # attn_res branch loop (K3 uses this when attn_res_block_size is set) + ar_anchor = ( + " hidden_states, residual = layer(\n" + " positions=positions,\n" + " hidden_states=hidden_states,\n" + " residual=residual,\n" + " )\n" + " if (layer_idx + 1) in self.aux_hidden_state_layers:\n" + ) + ar_repl = ( + " if _k3ml:\n" + " import torch as _t_k3ml, logging as _lg_k3ml\n" + " _t_k3ml.cuda.synchronize()\n" + " _lg_k3ml.getLogger(__name__).warning(\n" + " \"[k3-bc] attnres layer %d START (type=%s)\", layer_idx,\n" + " type(getattr(layer, \"self_attn\", layer)).__name__)\n" + + ar_anchor + ) + if ar_anchor in src: + src = src.replace(ar_anchor, ar_repl, 1) + else: + print("[k3-layerbc] WARN attn_res loop anchor not found.") + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-layerbc] ERROR compile: {e}", file=sys.stderr) + return 1 + print("[k3-layerbc] added embed + per-layer synced breadcrumbs.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mla_single_split.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mla_single_split.py new file mode 100644 index 00000000..190fefa0 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mla_single_split.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""TEST/FIX: force num_kv_splits=1 in TRITON_MLA forward_mqa (decode MLA attention). + +Hypothesis (UPDATE 25): the disagg recall failure is NON-DETERMINISTIC at greedy +temp=0 because the TRITON_MLA decode kernel splits the KV sequence into +num_kv_splits>1 partials (attn_logits in a shared workspace) and merges them with a +NON-DETERMINISTIC reduction. On a remote-prefill request the paged KV block's tail +(beyond the live tokens) is UNINITIALIZED (RDMA writes only live bytes; a local +prefill inits the whole block), so some splits read garbage and the non-det merge +yields wildly varying wrong tokens. vLLM's own code sets num_kv_splits=1 under +VLLM_BATCH_INVARIANT "to ensure deterministic reduction" — but that global env is +"not supported" for our MoE/EP/MXFP4 config, so we force it locally here. + +This edits vllm/v1/attention/backends/mla/triton_mla.py forward_mqa: replace the +num_kv_splits computation with a forced 1 (gated by env K3_MLA_SINGLE_SPLIT, default +ON when the patch is applied so a simple relaunch tests it). Single split = one +contiguous KV scan per query, deterministic, and reads only [0, seq_len) so no +cross-split garbage merge. Perf: one split is slower for very long seqs (less +parallelism) but correct; if it fixes recall we then do the perf-preserving fix +(zero-init decode KV blocks before RDMA write, keep multi-split). + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_mla_single_split.py +""" +import os +import sys + +REL = "v1/attention/backends/mla/triton_mla.py" +MARK = "k3-single-split" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + old = ( + " # For batch invariance, use only 1 split to ensure deterministic reduction\n" + " if envs.VLLM_BATCH_INVARIANT:\n" + " num_kv_splits = 1\n" + " else:\n" + " num_kv_splits = _compute_num_kv_splits(\n" + " attn_metadata.max_seq_len, self._sm_count\n" + " )\n" + ) + new = ( + " # " + MARK + ": force single split on the disagg decode path so the\n" + " # split-KV reduction is deterministic AND never merges an\n" + " # uninitialized-page-tail partial (remote-prefill KV blocks have\n" + " # uninitialized tails). K3_MLA_SINGLE_SPLIT=0 restores multi-split.\n" + " import os as _k3ssos\n" + " if _k3ssos.environ.get('K3_MLA_SINGLE_SPLIT', '1') != '0':\n" + " num_kv_splits = 1\n" + " elif envs.VLLM_BATCH_INVARIANT:\n" + " num_kv_splits = 1\n" + " else:\n" + " num_kv_splits = _compute_num_kv_splits(\n" + " attn_metadata.max_seq_len, self._sm_count\n" + " )\n" + ) + if old not in src: + print(f"[{MARK}] anchor NOT found", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_block_size_fix.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_block_size_fix.py new file mode 100644 index 00000000..2bb806f6 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_block_size_fix.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python3 +"""Fix MoRIIO self.block_size reconciliation for hybrid K3 (KDA + MLA). + +PROBLEM + register_kv_caches() reconciles self.block_size from `first_layer_name`'s + geometry. `first_layer_name` is chosen (moriio_connector.py:~1707) as the first + NON-MLA 5-D KV layer, falling back to the very first layer. For hybrid Kimi-K3 + the full-attention layers are MLA (excluded), so `first_layer_name` resolves to + a KDA *mamba* layer whose block_size is 1 (indivisible conv/ssm state page). + self.block_size then becomes 1. Every real attention layer has block_size 1536 + (the mamba-page padding raises the attention page to match), so the per-layer + guard trips: + ValueError: MoRIIO KV cache block size mismatch for layer + language_model.model.layers.3.self_attn.attn: 1536 != 1 + +FIX (surgical) + Reconcile self.block_size from the first NON-mamba layer's geometry, not from + whichever layer happens to be first. Mamba layers keep block_size 1 and are + already exempted from the guard (the `_k3_is_mamba` branch a few lines below). + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_block_size_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-blocksize] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + + if "k3-kda: block_size must reflect the ATTENTION layers" in src: + print("[k3-blocksize] already applied.") + return 0 + + anchor = ( + " self.num_blocks = first_geometry.num_blocks\n" + " self.slot_size_bytes = first_geometry.slot_size_bytes\n" + " if first_geometry.block_size != self.block_size:\n" + ) + repl = ( + " self.num_blocks = first_geometry.num_blocks\n" + " self.slot_size_bytes = first_geometry.slot_size_bytes\n" + " # k3-kda: block_size must reflect the ATTENTION layers, not mamba.\n" + " # first_layer_name can be a KDA mamba layer (block_size=1); using\n" + " # it would make every real attention layer (block_size 1536 after\n" + " # mamba-page padding) trip the guard below. Reconcile from the\n" + " # first non-mamba layer's geometry instead.\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MambaSpecBS\n" + " _attn_block_size = None\n" + " for _ln in kv_caches:\n" + " if isinstance(self.layer_to_spec.get(_ln), _K3MambaSpecBS):\n" + " continue\n" + " _attn_block_size = self._get_layer_transfer_geometry(_ln).block_size\n" + " break\n" + " if _attn_block_size is None:\n" + " _attn_block_size = first_geometry.block_size\n" + " if _attn_block_size != self.block_size:\n" + ) + if anchor in src: + src = src.replace(anchor, repl, 1) + else: + print("[k3-blocksize] WARN: anchor not found -- reconcile block may still " + "use first_geometry. Review register_kv_caches().") + return 0 + + # The reconcile body still references first_geometry.block_size in the log + + # assignment; rewrite those two to _attn_block_size within the reconcile block. + old_body = ( + " logger.info(\n" + ' "KV cache block_size=%d differs from config block_size=%d; "\n' + ' "using actual tensor shape (attention backend override).",\n' + " first_geometry.block_size,\n" + " self.block_size,\n" + " )\n" + " self.block_size = first_geometry.block_size\n" + ) + new_body = ( + " logger.info(\n" + ' "KV cache block_size=%d differs from config block_size=%d; "\n' + ' "using actual tensor shape (attention backend override).",\n' + " _attn_block_size,\n" + " self.block_size,\n" + " )\n" + " self.block_size = _attn_block_size\n" + ) + if old_body in src: + src = src.replace(old_body, new_body, 1) + else: + print("[k3-blocksize] WARN: reconcile body not matched; the new block_size " + "guard var may be unused. Review manually.") + + if src != orig: + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-blocksize] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print("[k3-blocksize] reconcile self.block_size from first non-mamba " + "layer (fixes MLA vs KDA block_size guard).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_group_routing.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_group_routing.py new file mode 100644 index 00000000..0672b883 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_group_routing.py @@ -0,0 +1,446 @@ +#!/usr/bin/env python3 +"""ROOT-CAUSE FIX: per-kv-cache-group block-id routing for Kimi-K3 disagg. + +Kimi-K3 is a hybrid model with FOUR kv-cache groups (confirmed at runtime): + idx 0: MambaSpec (23 KDA layers) + idx 1: MambaSpec (23 KDA layers) + idx 2: MambaSpec (23 KDA layers) + idx 3: MLAAttentionSpec (24 full-attention MLA layers) <- ATTENTION GROUP +``blocks.get_block_ids()`` returns a tuple of 4 lists, one per group, in this +index order. + +The MoRIIO connector was written for a 2-group model. It hardcodes: + * attention blocks = get_block_ids()[0] (WRONG: idx 0 is a Mamba group) + * mamba blocks = get_block_ids()[1] (WRONG: mamba is spread over 0/1/2) +Result: the producer RDMA-writes MLA KV using group-0 (mamba) block ids while +decode's MLA attention reads group-3 blocks -> decode reads NEVER-WRITTEN blocks +-> garbage exact-recall. This is THE root cause of the disagg decode-recall bug. + +THE FIX (this patch): carry ALL groups' block-id lists end-to-end as a +list-of-lists keyed by group index, and route EVERY layer's KV transfer to ITS +OWN group's list, indexed by the layer's kv-cache-group index. Concretely, in +MoRIIOWriter._prepare_transfer_plan, per layer: + gi = worker._layer_group_idx.get(layer_name, worker._attn_group_idx) + _k3_local = task.all_group_block_ids[gi] # producer's own blocks + _k3_remote = request_info.all_group_block_ids[gi] # decode's blocks +with a safe fallback to the legacy mamba/attn ([1]/[0]) behavior when +all_group_block_ids is absent/short. + +Wire/dataclass plumbing (all trailing-defaulted -> backward-compatible): + moriio_common.py : WriteTask + ReqMeta + RemoteAllocInfo gain + ``all_group_block_ids``; add_new_req gains the kwarg. + moriio_connector.py: + - scheduler producer-save (update_state_after_alloc): stash full group tuple + in self._reqs_save_allgrp; pass it through add_new_req. + - scheduler decode-advertise: send the full list-of-lists in the notify + (key "all_group_block_notify"); send_notify_block wire dict carries it. + - worker __init__: compute self._layer_group_idx (layer_name -> group idx) + and self._attn_group_idx (first non-Mamba group) from + kv_cache_config.kv_cache_groups; self._k3_group_routing from env. + - schedule_write_blocks / WriteTask / _write_blocks_for_req: plumb the field. + moriio_engine.py: + - _handle_remote_blocks_message: store all_group_block_ids into + RemoteAllocInfo. + - _prepare_transfer_plan: the group-indexed routing (see above). + +Gated behind env K3_GROUP_ROUTING (default "1" = on; "0" restores legacy +[0]/[1] routing for A/B). Non-K3 / 2-group models NEVER populate +all_group_block_ids, so they always take the legacy fallback -> zero behavior +change. Idempotent (MARK guard per file), TWO-PASS (verify every anchor across +every file BEFORE any write; any miss => zero writes, return 1), py_compile of +each edited file. + +Usage: apply_kimik3_moriio_group_routing.py +""" +import os +import sys + +MARK = "k3-group-routing" + +COMMON = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py" +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +ENGINE = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +# --------------------------------------------------------------------------- # +# moriio_common.py +# --------------------------------------------------------------------------- # +COMMON_EDITS = [ + ( + "C1 WriteTask field", + " mamba_local_block_ids: list[int] | None = None # k3-mamba-blockids\n" + " enqueue_time: float = field(default_factory=time.perf_counter)\n", + " mamba_local_block_ids: list[int] | None = None # k3-mamba-blockids\n" + " all_group_block_ids: list[list[int]] | None = None # " + MARK + "\n" + " enqueue_time: float = field(default_factory=time.perf_counter)\n", + ), + ( + "C2 RemoteAllocInfo field", + " mamba_block_ids: list[int] | None = None # k3-mamba-blockids\n" + " writes_done: int = 0\n", + " mamba_block_ids: list[int] | None = None # k3-mamba-blockids\n" + " all_group_block_ids: list[list[int]] | None = None # " + MARK + "\n" + " writes_done: int = 0\n", + ), + ( + "C3 ReqMeta field", + " # k3-mamba-blockids: mamba KV-group [1] local slot id(s) for this req.\n" + " mamba_local_block_ids: list[int] = field(default_factory=list)\n", + " # k3-mamba-blockids: mamba KV-group [1] local slot id(s) for this req.\n" + " mamba_local_block_ids: list[int] = field(default_factory=list)\n" + " # " + MARK + ": ALL kv-cache-groups' local block-id lists (list of\n" + " # lists, indexed by group index) so each layer routes to its own group.\n" + " all_group_block_ids: list[list[int]] | None = None\n", + ), + ( + "C4 add_new_req signature", + " write_mode=False,\n" + " mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids\n" + " ):\n", + " write_mode=False,\n" + " mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids\n" + " all_group_block_ids: list[list[int]] | None = None, # " + MARK + "\n" + " ):\n", + ), + ( + "C5 add_new_req set field", + " _req.mamba_local_block_ids = list(mamba_local_block_ids or []) # k3-mamba-blockids\n" + " if write_mode:\n", + " _req.mamba_local_block_ids = list(mamba_local_block_ids or []) # k3-mamba-blockids\n" + " _req.all_group_block_ids = ( # " + MARK + "\n" + " [list(g) for g in all_group_block_ids]\n" + " if all_group_block_ids is not None else None\n" + " )\n" + " if write_mode:\n", + ), +] + + +# --------------------------------------------------------------------------- # +# moriio_connector.py +# --------------------------------------------------------------------------- # +CONN_EDITS = [ + ( + "N3 scheduler init dict", + " self._reqs_save_mamba: dict[ReqId, list[int]] = {} # k3-mamba-blockids\n", + " self._reqs_save_mamba: dict[ReqId, list[int]] = {} # k3-mamba-blockids\n" + " self._reqs_save_allgrp: dict[ReqId, list[list[int]]] = {} # " + MARK + "\n", + ), + ( + "N6 send_notify_block signature", + " mamba_block_notify_list: list[int] | None = None, # k3-mamba-blockids\n" + " ):\n" + " path = make_zmq_path(\"tcp\", host, port)\n", + " mamba_block_notify_list: list[int] | None = None, # k3-mamba-blockids\n" + " all_group_block_notify: list[list[int]] | None = None, # " + MARK + "\n" + " ):\n" + " path = make_zmq_path(\"tcp\", host, port)\n", + ), + ( + "N7 send_notify_block wire dict", + " \"mamba_block_notify_list\": mamba_block_notify_list or [], # k3-mamba-blockids\n", + " \"mamba_block_notify_list\": mamba_block_notify_list or [], # k3-mamba-blockids\n" + " \"all_group_block_notify\": all_group_block_notify or [], # " + MARK + "\n", + ), + ( + "N2 scheduler producer-save capture", + " self._reqs_save_mamba[request.request_id] = ( # k3-mamba-blockids\n" + " list(_k3_gbi[1]) if len(_k3_gbi) > 1 else []\n" + " )\n" + " self._reqs_need_save[request.request_id] = (request, local_block_ids)\n", + " self._reqs_save_mamba[request.request_id] = ( # k3-mamba-blockids\n" + " list(_k3_gbi[1]) if len(_k3_gbi) > 1 else []\n" + " )\n" + " self._reqs_save_allgrp[request.request_id] = [ # " + MARK + "\n" + " list(g) for g in _k3_gbi\n" + " ]\n" + " self._reqs_need_save[request.request_id] = (request, local_block_ids)\n", + ), + ( + "N4 scheduler decode-advertise compute", + " mamba_block_notify_list = (\n" + " (list(_k3_gbi_d[1]) if len(_k3_gbi_d) > 1 else [])\n" + " if num_external_tokens > 0 else []\n" + " )\n", + " mamba_block_notify_list = (\n" + " (list(_k3_gbi_d[1]) if len(_k3_gbi_d) > 1 else [])\n" + " if num_external_tokens > 0 else []\n" + " )\n" + " all_group_block_notify = ( # " + MARK + "\n" + " [list(g) for g in _k3_gbi_d]\n" + " if num_external_tokens > 0 else []\n" + " )\n", + ), + ( + "N5 scheduler send_notify_block call", + " self.send_notify_block(\n" + " req_id=request.request_id,\n" + " transfer_id=request.kv_transfer_params[\"transfer_id\"],\n" + " block_notify_list=block_notify_list,\n" + " host=_notify_host,\n" + " port=target_port,\n" + " mamba_block_notify_list=mamba_block_notify_list, # k3-mamba-blockids\n" + " )\n", + " self.send_notify_block(\n" + " req_id=request.request_id,\n" + " transfer_id=request.kv_transfer_params[\"transfer_id\"],\n" + " block_notify_list=block_notify_list,\n" + " host=_notify_host,\n" + " port=target_port,\n" + " mamba_block_notify_list=mamba_block_notify_list, # k3-mamba-blockids\n" + " all_group_block_notify=all_group_block_notify, # " + MARK + "\n" + " )\n", + ), + ( + "N8 build_connector_meta add_new_req (pending final chunk)", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=self._reqs_need_pending_save[req_id][1],\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids\n" + " )\n", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=self._reqs_need_pending_save[req_id][1],\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids\n" + " all_group_block_ids=self._reqs_save_allgrp.get(req_id, None), # " + MARK + "\n" + " )\n", + ), + ( + "N9 build_connector_meta add_new_req (direct save)", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=block_ids,\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids\n" + " )\n", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=block_ids,\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids\n" + " all_group_block_ids=self._reqs_save_allgrp.get(req_id, None), # " + MARK + "\n" + " )\n", + ), + ( + "N11 schedule_write_blocks signature", + " mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids\n" + " ) -> None:\n" + " \"\"\"Schedule a block write operation.\n", + " mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids\n" + " all_group_block_ids: list[list[int]] | None = None, # " + MARK + "\n" + " ) -> None:\n" + " \"\"\"Schedule a block write operation.\n", + ), + ( + "N12 WriteTask construction", + " remote_block_ids_hint=remote_block_ids,\n" + " mamba_local_block_ids=mamba_local_block_ids, # k3-mamba-blockids\n" + " layer_name=layer_name,\n", + " remote_block_ids_hint=remote_block_ids,\n" + " mamba_local_block_ids=mamba_local_block_ids, # k3-mamba-blockids\n" + " all_group_block_ids=all_group_block_ids, # " + MARK + "\n" + " layer_name=layer_name,\n", + ), + ( + "N10 _write_blocks_for_req schedule_write_blocks call", + " remote_notify_port=meta.remote_notify_port,\n" + " remote_ip=meta.remote_host,\n" + " mamba_local_block_ids=meta.mamba_local_block_ids, # k3-mamba-blockids\n" + " )\n", + " remote_notify_port=meta.remote_notify_port,\n" + " remote_ip=meta.remote_host,\n" + " mamba_local_block_ids=meta.mamba_local_block_ids, # k3-mamba-blockids\n" + " all_group_block_ids=meta.all_group_block_ids, # " + MARK + "\n" + " )\n", + ), + ( + "N1 worker __init__ compute group indices", + " self.layer_to_spec = build_layer_to_spec(kv_cache_config)\n", + " self.layer_to_spec = build_layer_to_spec(kv_cache_config)\n" + " # " + MARK + ": Kimi-K3 has 4 kv-cache groups (0/1/2 mamba, 3 MLA).\n" + " # Map each layer to ITS OWN group index and find the attention\n" + " # (non-Mamba) group so every layer routes to its own block-id list.\n" + " self._layer_group_idx: dict[str, int] = {}\n" + " self._attn_group_idx = 0\n" + " try:\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3GR_Mamba\n" + " _k3gr_groups = getattr(kv_cache_config, \"kv_cache_groups\", []) or []\n" + " for _k3gr_gi, _k3gr_grp in enumerate(_k3gr_groups):\n" + " for _k3gr_ln in getattr(_k3gr_grp, \"layer_names\", []) or []:\n" + " self._layer_group_idx[_k3gr_ln] = _k3gr_gi\n" + " for _k3gr_gi, _k3gr_grp in enumerate(_k3gr_groups):\n" + " _k3gr_lns = getattr(_k3gr_grp, \"layer_names\", []) or []\n" + " if _k3gr_lns and not isinstance(\n" + " self.layer_to_spec.get(_k3gr_lns[0]), _K3GR_Mamba\n" + " ):\n" + " self._attn_group_idx = _k3gr_gi\n" + " break\n" + " except Exception:\n" + " self._layer_group_idx = {}\n" + " self._attn_group_idx = 0\n" + " import os as _k3gros\n" + " self._k3_group_routing = (\n" + " _k3gros.environ.get(\"K3_GROUP_ROUTING\", \"1\") == \"1\"\n" + " )\n" + " logger.info(\n" + " \"[" + MARK + "] enabled=%s attn_group_idx=%s n_groups=%s \"\n" + " \"sample_layer_group_idx=%s\",\n" + " self._k3_group_routing,\n" + " self._attn_group_idx,\n" + " len(getattr(kv_cache_config, \"kv_cache_groups\", []) or []),\n" + " dict(list(self._layer_group_idx.items())[:4]),\n" + " )\n", + ), +] + + +# --------------------------------------------------------------------------- # +# moriio_engine.py +# --------------------------------------------------------------------------- # +ENGINE_EDITS = [ + ( + "E1 _handle_remote_blocks_message capture", + " block_notify_list = data.get(\"block_notify_list\", [])\n" + " mamba_block_notify_list = data.get(\"mamba_block_notify_list\", []) # k3-mamba-blockids\n" + " decode_dp_rank = data.get(\"decode_rank\", 0)\n", + " block_notify_list = data.get(\"block_notify_list\", [])\n" + " mamba_block_notify_list = data.get(\"mamba_block_notify_list\", []) # k3-mamba-blockids\n" + " all_group_block_notify = data.get(\"all_group_block_notify\", []) # " + MARK + "\n" + " decode_dp_rank = data.get(\"decode_rank\", 0)\n", + ), + ( + "E2 RemoteAllocInfo construction", + " self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(\n" + " block_ids=block_notify_list, decode_dp_rank=decode_dp_rank,\n" + " mamba_block_ids=list(mamba_block_notify_list or []), # k3-mamba-blockids\n" + " )\n", + " self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(\n" + " block_ids=block_notify_list, decode_dp_rank=decode_dp_rank,\n" + " mamba_block_ids=list(mamba_block_notify_list or []), # k3-mamba-blockids\n" + " all_group_block_ids=( # " + MARK + "\n" + " [list(g) for g in all_group_block_notify]\n" + " if all_group_block_notify else None\n" + " ),\n" + " )\n", + ), + ( + "E3 _prepare_transfer_plan group routing", + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MS_BL # k3-mamba-blockids\n" + " _k3_mamba = isinstance(\n" + " self.worker.layer_to_spec.get(task.layer_name), _K3MS_BL\n" + " )\n" + " if _k3_mamba:\n" + " # k3-mamba-blockids: mamba/KDA state lives in a SEPARATE KV-cache\n" + " # group whose slot ids differ from the attention group's block ids.\n" + " # Route the mamba-layer transfer by the mamba group's ids (falling\n" + " # back to attention ids for non-hybrid models).\n" + " _k3_local = task.mamba_local_block_ids or task.local_block_ids\n" + " _k3_remote = request_info.mamba_block_ids or request_info.block_ids\n" + " else:\n" + " _k3_local = task.local_block_ids\n" + " _k3_remote = request_info.block_ids\n", + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MS_BL # k3-mamba-blockids\n" + " # " + MARK + ": route EVERY layer by ITS OWN kv-cache-group index.\n" + " # Kimi-K3 has 4 groups (0/1/2 mamba, 3 MLA); the legacy code below\n" + " # hardcoded [0]/[1] and sent MLA KV to mamba block ids. When all\n" + " # groups' block ids are carried end-to-end (K3_GROUP_ROUTING=1) use\n" + " # the per-layer group index; otherwise fall back to legacy behavior.\n" + " _k3_gr_local = getattr(task, \"all_group_block_ids\", None)\n" + " _k3_gr_remote = getattr(request_info, \"all_group_block_ids\", None)\n" + " _k3_gr_on = getattr(self.worker, \"_k3_group_routing\", False)\n" + " if (\n" + " _k3_gr_on\n" + " and _k3_gr_local is not None\n" + " and _k3_gr_remote is not None\n" + " ):\n" + " _k3_gi = self.worker._layer_group_idx.get(\n" + " task.layer_name, self.worker._attn_group_idx\n" + " )\n" + " if _k3_gi < len(_k3_gr_local) and _k3_gi < len(_k3_gr_remote):\n" + " _k3_local = _k3_gr_local[_k3_gi]\n" + " _k3_remote = _k3_gr_remote[_k3_gi]\n" + " else:\n" + " _k3_local = task.local_block_ids\n" + " _k3_remote = request_info.block_ids\n" + " elif isinstance(\n" + " self.worker.layer_to_spec.get(task.layer_name), _K3MS_BL\n" + " ):\n" + " # k3-mamba-blockids (legacy 2-group fallback): mamba/KDA state\n" + " # lives in a SEPARATE KV-cache group whose slot ids differ from\n" + " # the attention group's block ids.\n" + " _k3_local = task.mamba_local_block_ids or task.local_block_ids\n" + " _k3_remote = request_info.mamba_block_ids or request_info.block_ids\n" + " else:\n" + " _k3_local = task.local_block_ids\n" + " _k3_remote = request_info.block_ids\n", + ), +] + + +FILES = [ + (COMMON, COMMON_EDITS), + (CONN, CONN_EDITS), + (ENGINE, ENGINE_EDITS), +] + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + base = sys.argv[1] + + resolved = [] # (path, src, edits, already) + # ---- PASS 1: verify EVERY anchor across EVERY file BEFORE any write ---- # + missing = False + for rel, edits in FILES: + path = os.path.join(base, rel) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + already = MARK in src + resolved.append((path, src, edits, already)) + if already: + print(f"[{MARK}] {rel}: MARK already present -> will skip.") + continue + for tag, old, _new in edits: + cnt = src.count(old) + if cnt != 1: + print( + f"[{MARK}] {rel}: ANCHOR MISSING/NON-UNIQUE ({cnt}) :: {tag}", + file=sys.stderr, + ) + missing = True + if missing: + print(f"[{MARK}] aborting: unmatched anchors, ZERO writes.", file=sys.stderr) + return 1 + + # ---- PASS 2: apply + py_compile ---- # + import py_compile + + for path, src, edits, already in resolved: + if already: + continue + for _tag, old, new in edits: + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {path}: {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied + compiled: {path}") + + print(f"[{MARK}] done.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_blockids.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_blockids.py new file mode 100644 index 00000000..e1c04042 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_blockids.py @@ -0,0 +1,442 @@ +#!/usr/bin/env python3 +"""ROOT-CAUSE FIX: route Kimi-K3 KDA/mamba state transfer by the MAMBA KV-cache +group's block ids, not the attention group's. + +BUG (confirmed by code + prefill-direct-correct + decode-context-free): + Kimi-K3 is HYBRID: KV-cache group [0] = MLA attention (paged blocks), + group [1] = KDA/mamba (per-request MambaSpec state slot, block_size=1). + These groups allocate slots INDEPENDENTLY -> for the same request the + attention block id != the mamba slot id. + + The MoRIIO connector treats the model as attention-only: + * decode advertises only get_block_ids()[0] (attention) as block_notify_list + (moriio_connector.py ~768); the mamba group [1] slot id is NEVER sent. + * prefill's per-layer write task uses those attention block ids as BOTH + source and dest offsets for EVERY layer, including the ~69 KDA layers. + But decode's KDA forward reads its recurrent/conv state from the MAMBA group + slot (kda_metadata.py:283 non_spec_state_indices_tensor = block_table[:,0] of + the mamba group). So the KDA state is written to an ATTENTION block offset that + decode never reads -> decode's KDA state stays zero -> 69/93 layers blank -> + fluent but CONTEXT-FREE output. (Attention KV works because its block ids DO + match on both ends -- which is why prefill-direct is correct and attention-only + looked byte-perfect.) + +FIX (thread the mamba group [1] block ids end-to-end, use them ONLY for mamba +layers; attention path unchanged -> byte-identical for attention): + 1. moriio_common: additive fields + - WriteTask.mamba_local_block_ids + - RemoteAllocInfo.mamba_block_ids + - ReqMeta.mamba_local_block_ids + - add_new_req(mamba_local_block_ids=None) sets ReqMeta field + 2. connector (decode advertise): also compute mamba_block_notify_list = + get_block_ids()[1] and send it via send_notify_block. + 3. connector send_notify_block: carry mamba_block_notify_list in the msg. + 4. engine _handle_remote_blocks_message: store mamba_block_ids in RemoteAllocInfo. + 5. connector (prefill save capture in update_state_after_alloc do_remote_decode): + capture prefill's own mamba slot get_block_ids()[1] into a side dict + self._reqs_save_mamba (init in scheduler __init__). + 6. connector build_connector_meta save add_new_req calls: pass mamba_local. + 7. connector _write_blocks_for_req -> schedule_write_blocks -> WriteTask: carry + mamba_local_block_ids. + 8. engine _prepare_transfer_plan: for a MambaSpec layer, compute offsets from + the mamba local/remote block ids instead of the attention ones. + +All hunks idempotent + anchor-based + py_compile-checked. Non-hybrid models keep +group [1] empty -> mamba lists are [] -> zero behavior change. + +Usage: apply_kimik3_moriio_mamba_blockids.py +""" +import os +import sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +ENG = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" +COMMON = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py" + +MARK = "k3-mamba-blockids" + + +def _apply(path, subs, tag): + """subs: list of (old, new, guard_substr). If guard_substr in src, skip that + sub (idempotent). Hard-error if an old anchor is missing and not applied.""" + if not os.path.isfile(path): + print(f"[{MARK}] {tag}: FILE NOT FOUND {path}", file=sys.stderr) + return False + src = open(path).read() + orig = src + for old, new, guard in subs: + if guard in src: + continue + if old not in src: + print(f"[{MARK}] {tag}: ANCHOR MISSING (guard={guard!r})", file=sys.stderr) + return False + src = src.replace(old, new, 1) + if src == orig: + print(f"[{MARK}] {tag}: already applied.") + return True + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[{MARK}] {tag}: COMPILE FAIL, rolled back: {e}", file=sys.stderr) + return False + print(f"[{MARK}] {tag}: applied.") + return True + + +def patch_common(base): + path = os.path.join(base, COMMON) + subs = [ + # WriteTask field: MUST go after the last non-default field (remote_ip) + # and before the first field with a default (enqueue_time), else + # dataclass raises "non-default argument follows default argument". + ( + " remote_ip: str\n" + " enqueue_time: float = field(default_factory=time.perf_counter)\n", + " remote_ip: str\n" + " mamba_local_block_ids: list[int] | None = None # " + MARK + "\n" + " enqueue_time: float = field(default_factory=time.perf_counter)\n", + "mamba_local_block_ids: list[int] | None = None # " + MARK, + ), + # RemoteAllocInfo field + ( + " block_ids: list[int]\n" + " writes_done: int = 0\n", + " block_ids: list[int]\n" + " mamba_block_ids: list[int] | None = None # " + MARK + "\n" + " writes_done: int = 0\n", + "mamba_block_ids: list[int] | None = None # " + MARK, + ), + # ReqMeta field (append after remote_dp_size_local default) + ( + " # Per-pod DP size; 0 means fallback to remote_dp_size.\n" + " remote_dp_size_local: int = 0\n", + " # Per-pod DP size; 0 means fallback to remote_dp_size.\n" + " remote_dp_size_local: int = 0\n" + " # " + MARK + ": mamba KV-group [1] local slot id(s) for this req.\n" + " mamba_local_block_ids: list[int] = field(default_factory=list)\n", + MARK + ": mamba KV-group [1] local slot", + ), + # add_new_req signature: add optional param + ( + " def add_new_req(\n" + " self,\n" + " request_id: ReqId,\n" + " local_block_ids: list[int],\n" + " kv_transfer_params: dict[str, Any],\n" + " write_mode=False,\n" + " ):", + " def add_new_req(\n" + " self,\n" + " request_id: ReqId,\n" + " local_block_ids: list[int],\n" + " kv_transfer_params: dict[str, Any],\n" + " write_mode=False,\n" + " mamba_local_block_ids: list[int] | None = None, # " + MARK + "\n" + " ):", + "mamba_local_block_ids: list[int] | None = None, # " + MARK, + ), + # add_new_req body: set ReqMeta field before dispatch to reqs_to_save/recv + ( + " if write_mode:\n" + " self.reqs_to_save[request_id] = _req\n" + " else:\n" + " self.reqs_to_recv[request_id] = _req\n", + " _req.mamba_local_block_ids = list(mamba_local_block_ids or []) # " + MARK + "\n" + " if write_mode:\n" + " self.reqs_to_save[request_id] = _req\n" + " else:\n" + " self.reqs_to_recv[request_id] = _req\n", + "_req.mamba_local_block_ids = list(mamba_local_block_ids or []) # " + MARK, + ), + ] + return _apply(path, subs, "moriio_common.py") + + +def patch_engine(base): + path = os.path.join(base, ENG) + subs = [ + # _handle_remote_blocks_message: capture mamba_block_notify_list + ( + " block_notify_list = data.get(\"block_notify_list\", [])\n" + " decode_dp_rank = data.get(\"decode_rank\", 0)\n", + " block_notify_list = data.get(\"block_notify_list\", [])\n" + " mamba_block_notify_list = data.get(\"mamba_block_notify_list\", []) # " + MARK + "\n" + " decode_dp_rank = data.get(\"decode_rank\", 0)\n", + "mamba_block_notify_list = data.get(\"mamba_block_notify_list\", []) # " + MARK, + ), + ( + " self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(\n" + " block_ids=block_notify_list, decode_dp_rank=decode_dp_rank\n" + " )\n", + " self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo(\n" + " block_ids=block_notify_list, decode_dp_rank=decode_dp_rank,\n" + " mamba_block_ids=list(mamba_block_notify_list or []), # " + MARK + "\n" + " )\n", + "mamba_block_ids=list(mamba_block_notify_list or []), # " + MARK, + ), + # _prepare_transfer_plan: mamba layers use mamba block ids + ( + " offsets = request_info.transfer_offsets.get(geometry_key)\n" + " if offsets is None:\n" + " offsets = self.worker._compute_block_transfer_offsets(\n" + " task.layer_name,\n" + " task.local_block_ids,\n" + " request_info.block_ids,\n" + " remote_moriio_meta,\n" + " )\n" + " request_info.transfer_offsets[geometry_key] = offsets\n", + " offsets = request_info.transfer_offsets.get(geometry_key)\n" + " if offsets is None:\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MS_BL # " + MARK + "\n" + " _k3_mamba = isinstance(\n" + " self.worker.layer_to_spec.get(task.layer_name), _K3MS_BL\n" + " )\n" + " if _k3_mamba:\n" + " _k3_local = task.mamba_local_block_ids or task.local_block_ids\n" + " _k3_remote = request_info.mamba_block_ids or request_info.block_ids\n" + " else:\n" + " _k3_local = task.local_block_ids\n" + " _k3_remote = request_info.block_ids\n" + " import os as _k3dbgos\n" + " if _k3dbgos.environ.get('K3_MAMBA_BC', '0') == '1':\n" + " _k3dbg = getattr(self, '_k3_mb_seen', None)\n" + " if _k3dbg is None:\n" + " _k3dbg = set(); self._k3_mb_seen = _k3dbg\n" + " _k3k = ('M' if _k3_mamba else 'A')\n" + " if _k3k not in _k3dbg:\n" + " _k3dbg.add(_k3k)\n" + " _k3srcnorm = -1.0; _k3el = []\n" + " try:\n" + " _k3t = self.worker.kv_caches[task.layer_name]\n" + " if _k3_local:\n" + " _k3slot = _k3t[int(_k3_local[0])].flatten()\n" + " _k3srcnorm = float(_k3slot.float().norm())\n" + " _k3el = [round(float(x),3) for x in _k3slot[:6].float().tolist()]\n" + " except Exception as _k3ee:\n" + " _k3srcnorm = -2.0\n" + " import logging as _k3dbglg\n" + " _k3dbglg.getLogger('vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine').info(\n" + " '[k3-mamba-bc] layer=%s mamba=%s used_local=%s used_remote=%s src_slot_norm=%.4e src_el=%s tshape=%s',\n" + " task.layer_name, _k3_mamba,\n" + " (_k3_local or [])[:3], (_k3_remote or [])[:3],\n" + " _k3srcnorm, _k3el, tuple(self.worker.kv_caches[task.layer_name].shape),\n" + " )\n" + " offsets = self.worker._compute_block_transfer_offsets(\n" + " task.layer_name,\n" + " _k3_local,\n" + " _k3_remote,\n" + " remote_moriio_meta,\n" + " )\n" + " request_info.transfer_offsets[geometry_key] = offsets\n", + "_k3_mamba = isinstance(", + ), + ] + return _apply(path, subs, "moriio_engine.py") + + +def patch_connector(base): + path = os.path.join(base, CONN) + subs = [ + # scheduler __init__: side dict for prefill's own mamba slot per req + ( + " self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}\n", + " self._reqs_need_save: dict[ReqId, tuple[Request, list[int]]] = {}\n" + " self._reqs_save_mamba: dict[ReqId, list[int]] = {} # " + MARK + "\n", + "self._reqs_save_mamba: dict[ReqId, list[int]] = {} # " + MARK, + ), + # prefill save capture (do_remote_decode branch) + ( + " if params.get(\"do_remote_decode\"):\n" + " local_block_ids = blocks.get_block_ids()[0]\n" + " self._reqs_need_save[request.request_id] = (request, local_block_ids)\n", + " if params.get(\"do_remote_decode\"):\n" + " _k3_gbi = blocks.get_block_ids()\n" + " local_block_ids = _k3_gbi[0]\n" + " self._reqs_save_mamba[request.request_id] = ( # " + MARK + "\n" + " list(_k3_gbi[1]) if len(_k3_gbi) > 1 else []\n" + " )\n" + " self._reqs_need_save[request.request_id] = (request, local_block_ids)\n", + "self._reqs_save_mamba[request.request_id] = ( # " + MARK, + ), + # send_notify_block signature + ( + " def send_notify_block(\n" + " self,\n" + " req_id: ReqId,\n" + " transfer_id: TransferId,\n" + " block_notify_list: list[int],\n" + " host=None,\n" + " port=None,\n" + " ):", + " def send_notify_block(\n" + " self,\n" + " req_id: ReqId,\n" + " transfer_id: TransferId,\n" + " block_notify_list: list[int],\n" + " host=None,\n" + " port=None,\n" + " mamba_block_notify_list: list[int] | None = None, # " + MARK + "\n" + " ):", + "mamba_block_notify_list: list[int] | None = None, # " + MARK, + ), + # send_notify_block body: include mamba list in message + ( + " \"block_notify_list\": block_notify_list or [],\n", + " \"block_notify_list\": block_notify_list or [],\n" + " \"mamba_block_notify_list\": mamba_block_notify_list or [], # " + MARK + "\n", + "\"mamba_block_notify_list\": mamba_block_notify_list or [], # " + MARK, + ), + # decode advertise: build mamba_block_notify_list + ( + " block_notify_list = (\n" + " blocks.get_block_ids()[0] if num_external_tokens > 0 else []\n" + " )\n", + " _k3_gbi_d = blocks.get_block_ids() # " + MARK + "\n" + " block_notify_list = (\n" + " _k3_gbi_d[0] if num_external_tokens > 0 else []\n" + " )\n" + " mamba_block_notify_list = (\n" + " (list(_k3_gbi_d[1]) if len(_k3_gbi_d) > 1 else [])\n" + " if num_external_tokens > 0 else []\n" + " )\n", + "_k3_gbi_d = blocks.get_block_ids() # " + MARK, + ), + # decode advertise: pass mamba list to send_notify_block + ( + " self.send_notify_block(\n" + " req_id=request.request_id,\n" + " transfer_id=request.kv_transfer_params[\"transfer_id\"],\n" + " block_notify_list=block_notify_list,\n" + " host=_notify_host,\n" + " port=target_port,\n" + " )\n", + " self.send_notify_block(\n" + " req_id=request.request_id,\n" + " transfer_id=request.kv_transfer_params[\"transfer_id\"],\n" + " block_notify_list=block_notify_list,\n" + " host=_notify_host,\n" + " port=target_port,\n" + " mamba_block_notify_list=mamba_block_notify_list, # " + MARK + "\n" + " )\n", + "mamba_block_notify_list=mamba_block_notify_list, # " + MARK, + ), + # build_connector_meta: chunked-prefill final-chunk save add_new_req + ( + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=self._reqs_need_pending_save[req_id][1],\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " )\n", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=self._reqs_need_pending_save[req_id][1],\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # " + MARK + "\n" + " )\n", + "mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # " + MARK + "\n" + " )\n", + ), + # build_connector_meta: single-chunk save add_new_req + ( + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=block_ids,\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " )\n", + " meta.add_new_req(\n" + " request_id=req_id,\n" + " local_block_ids=block_ids,\n" + " kv_transfer_params=kv_params,\n" + " write_mode=True,\n" + " mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # " + MARK + "\n" + " )\n", + "mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # " + MARK + "\n" + " )\n", + ), + # schedule_write_blocks signature + ( + " layer_name: str,\n" + " kv_layer: torch.Tensor,\n" + " remote_notify_port: int,\n" + " remote_ip: str,\n" + " ) -> None:\n" + " \"\"\"Schedule a block write operation.", + " layer_name: str,\n" + " kv_layer: torch.Tensor,\n" + " remote_notify_port: int,\n" + " remote_ip: str,\n" + " mamba_local_block_ids: list[int] | None = None, # " + MARK + "\n" + " ) -> None:\n" + " \"\"\"Schedule a block write operation.", + "mamba_local_block_ids: list[int] | None = None, # " + MARK, + ), + # schedule_write_blocks body: WriteTask carries mamba local + ( + " task = WriteTask(\n" + " request_id=request_id,\n" + " transfer_id=transfer_id,\n" + " dst_engine_id=dst_engine_id,\n" + " local_block_ids=local_block_ids,\n" + " remote_block_ids_hint=remote_block_ids,\n" + " layer_name=layer_name,\n", + " task = WriteTask(\n" + " request_id=request_id,\n" + " transfer_id=transfer_id,\n" + " dst_engine_id=dst_engine_id,\n" + " local_block_ids=local_block_ids,\n" + " remote_block_ids_hint=remote_block_ids,\n" + " mamba_local_block_ids=mamba_local_block_ids, # " + MARK + "\n" + " layer_name=layer_name,\n", + "mamba_local_block_ids=mamba_local_block_ids, # " + MARK, + ), + # _write_blocks_for_req: pass mamba local into schedule_write_blocks + ( + " self.schedule_write_blocks(\n" + " request_id=req_id,\n" + " transfer_id=meta.transfer_id,\n" + " dst_engine_id=meta.remote_engine_id,\n" + " local_block_ids=meta.local_block_ids,\n" + " remote_block_ids=meta.remote_block_ids,\n" + " layer_name=layer_name,\n" + " kv_layer=kv_layer,\n" + " remote_notify_port=meta.remote_notify_port,\n" + " remote_ip=meta.remote_host,\n" + " )\n", + " self.schedule_write_blocks(\n" + " request_id=req_id,\n" + " transfer_id=meta.transfer_id,\n" + " dst_engine_id=meta.remote_engine_id,\n" + " local_block_ids=meta.local_block_ids,\n" + " remote_block_ids=meta.remote_block_ids,\n" + " layer_name=layer_name,\n" + " kv_layer=kv_layer,\n" + " remote_notify_port=meta.remote_notify_port,\n" + " remote_ip=meta.remote_host,\n" + " mamba_local_block_ids=meta.mamba_local_block_ids, # " + MARK + "\n" + " )\n", + "mamba_local_block_ids=meta.mamba_local_block_ids, # " + MARK, + ), + ] + return _apply(path, subs, "moriio_connector.py") + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + base = sys.argv[1] + ok = True + ok = patch_common(base) and ok + ok = patch_engine(base) and ok + ok = patch_connector(base) and ok + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_diag.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_diag.py new file mode 100644 index 00000000..2361d6bc --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_diag.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python3 +"""Diagnostics for the KDA/mamba MoRIIO transfer GPU memory fault. + +Adds two things to compute_mamba_block_transfer_offsets in moriio_layout.py: + + 1. MORIIO_SKIP_MAMBA=1 -> return ([],[],[]) so the KDA/mamba state transfer is + a no-op (MLA KV still transfers). If the producer stops GPU-faulting under + this flag, the fault is isolated to the KDA state write path. + + 2. Bounds check: log [k3-kda OOB] when (block*stride + sub_off + size) exceeds + the local mamba cache tensor extent (num_blocks*page_size_bytes) -- i.e. the + exact offset that would form an out-of-bounds RDMA address. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_mamba_diag.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-mamba-diag] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + changed = 0 + + # --- 1) MORIIO_SKIP_MAMBA early return --- + if "MORIIO_SKIP_MAMBA" not in src: + anchor = ( + " if len(local_block_ids) > len(remote_block_ids):\n" + " raise ValueError(\n" + ' "local_block_ids longer than remote_block_ids (mamba): "\n' + ) + inject = ( + " import os as _os\n" + " if _os.environ.get(\"MORIIO_SKIP_MAMBA\", \"0\") == \"1\":\n" + " return [], [], []\n" + "\n" + " if len(local_block_ids) > len(remote_block_ids):\n" + " raise ValueError(\n" + ' "local_block_ids longer than remote_block_ids (mamba): "\n' + ) + if anchor in src: + src = src.replace(anchor, inject, 1) + changed += 1 + else: + print("[k3-mamba-diag] WARN: skip-flag anchor not found.") + + # --- 2) bounds-check logging in the write loop --- + if "[k3-kda OOB]" not in src: + loop_old = ( + " for lb, rb in zip(local_block_ids, remote_block_ids):\n" + " lbase = lb * stride\n" + " rbase = rb * stride\n" + " for off, sz in subregions:\n" + " offset_local[w] = lbase + off\n" + " offset_remote[w] = rbase + off # tp_ratio==1 -> same sub-offset\n" + " sizes[w] = sz\n" + " w += 1\n" + ) + loop_new = ( + " try:\n" + " _tensor_bytes = int(kv_cache.numel()) * int(kv_cache.element_size())\n" + " except Exception:\n" + " _tensor_bytes = -1\n" + " for lb, rb in zip(local_block_ids, remote_block_ids):\n" + " lbase = lb * stride\n" + " rbase = rb * stride\n" + " for off, sz in subregions:\n" + " offset_local[w] = lbase + off\n" + " offset_remote[w] = rbase + off # tp_ratio==1 -> same sub-offset\n" + " sizes[w] = sz\n" + " if _tensor_bytes >= 0 and (lbase + off + sz) > _tensor_bytes:\n" + " import logging as _lg\n" + " _lg.getLogger(__name__).error(\n" + ' "[k3-kda OOB] layer=%s lb=%d off=%d sz=%d end=%d > "\n' + ' "tensor_bytes=%d (stride=%d num_local=%d)",\n' + " layer_name, lb, off, sz, lbase + off + sz,\n" + " _tensor_bytes, stride, len(local_block_ids),\n" + " )\n" + " w += 1\n" + ) + if loop_old in src: + src = src.replace(loop_old, loop_new, 1) + changed += 1 + else: + print("[k3-mamba-diag] WARN: write-loop anchor not found (already " + "modified?); bounds check not added.") + + if changed == 0: + print("[k3-mamba-diag] nothing to do (already applied).") + return 0 + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-mamba-diag] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print(f"[k3-mamba-diag] applied {changed} change(s): MORIIO_SKIP_MAMBA " + "flag + KDA OOB bounds logging.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_n1.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_n1.py new file mode 100644 index 00000000..59b4fd1a --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_n1.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +"""ROOT-CAUSE FIX (the last one): mamba/KDA N-vs-N-1 prefill/decode boundary. + +Kimi-K3 is hybrid MLA + KDA(mamba). vLLM's own reference disagg connectors +(nixl/base_scheduler.py, mooncake) handle hybrid/mamba PD with a SYMMETRIC pair: + - P-side _truncate_mamba_request_for_prefill: drop the LAST prompt token so the + prefiller computes recurrent state h(N-1), not h(N). (max_tokens=1, guarded.) + - D-side _get_remote_prefill_token_count: return N-1 for mamba, so the decoder + recomputes the last prompt token and derives h(N) itself from h(N-1). +Comment verbatim from nixl base_scheduler.py:368: "Returns N-1 for Mamba models +since the decoder always recomputes the last token and must start from h(N-1)." + +The MoRIIO connector NEVER ported this. Result (proven by K3_INPUTS_PROBE): + PREFILL computes all N tokens (positions 0..N-1) -> KDA state seated AFTER token + N-1 (all N applied). DECODE gets num_computed=N-1, RECOMPUTES the last prompt + token, and re-applies it to the already-N-advanced recurrent state -> the last + token is DOUBLE-COUNTED in the KDA recurrence -> wrong logits from token 1 -> + decode echoes the last prompt token ("gold is"->"is"). Attention tolerates the + mismatch (paged, position-indexed); mamba's monolithic recurrent state does not. +This is why transferring correct KDA state is WORSE than zero state, and why the +byte-perfect transport still yields wrong output. + +THIS PATCH ports the nixl fix into MoRIIOConnectorScheduler: + H1 __init__: compute self._has_mamba from the model's hybrid config + (Kimi-K3 text_config.linear_attn_config.kda_layers) -> True for K3. + H2 add _truncate_mamba_request_for_prefill + _get_remote_prefill_token_count + (verbatim semantics from nixl base_scheduler). + H3 get_num_new_matched_tokens: + * D-side (consumer) WRITE: return N-1 (was N) for mamba so decode recomputes + the last token from h(N-1). [line ~507] + * P-side (producer) do_remote_decode + _has_mamba: truncate the prompt to + N-1 so prefill computes h(N-1). [producer branch] + +Attention path unchanged. Non-mamba models: _has_mamba False -> zero behavior +change (all N-1 -> N, no truncation). Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_mamba_n1.py +""" +import os +import sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +MARK = "k3-mamba-n1" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, CONN) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + + # --- H1: detect hybrid/mamba in MoRIIOConnectorScheduler.__init__ --- + h1_old = ( + " self.is_producer = self.kv_transfer_config.kv_role == \"kv_producer\"\n" + ) + h1_new = ( + " self.is_producer = self.kv_transfer_config.kv_role == \"kv_producer\"\n" + " # " + MARK + ": hybrid (mamba/KDA) detection. Kimi-K3 carries\n" + " # text_config.linear_attn_config.kda_layers; any linear/mamba/kda\n" + " # marker => the recurrent-state N-vs-N-1 boundary applies.\n" + " self._has_mamba = False\n" + " try:\n" + " _mc = getattr(self.vllm_config, 'model_config', None)\n" + " _hf = getattr(_mc, 'hf_config', None) if _mc is not None else None\n" + " _tc = getattr(_hf, 'text_config', None) or _hf\n" + " _la = getattr(_tc, 'linear_attn_config', None)\n" + " if _la is None and isinstance(getattr(_tc, '__dict__', None), dict):\n" + " _la = _tc.__dict__.get('linear_attn_config')\n" + " if _la:\n" + " self._has_mamba = True\n" + " except Exception:\n" + " self._has_mamba = False\n" + " import os as _k3n1os\n" + " if _k3n1os.environ.get('K3_MAMBA_N1_FORCE', '') in ('0', '1'):\n" + " self._has_mamba = (_k3n1os.environ['K3_MAMBA_N1_FORCE'] == '1')\n" + " logger.info('[" + MARK + "] _has_mamba=%s (mamba N-1 boundary %s)',\n" + " self._has_mamba, 'ON' if self._has_mamba else 'off')\n" + ) + + # --- H2: helper methods (insert right before get_num_new_matched_tokens of the + # scheduler class). Anchor on the scheduler's get_num_new_matched_tokens def + # (the one taking (self, request, num_computed_tokens) inside the Scheduler). + # There are two defs with that name; the scheduler's is the 2nd (line ~481). + h2_anchor = ( + " def get_num_new_matched_tokens(\n" + " self,\n" + " request: \"Request\",\n" + " num_computed_tokens: int,\n" + " ) -> tuple[int, bool]:\n" + ) + h2_new = ( + " def _get_remote_prefill_token_count(self, num_prompt_tokens: int) -> int:\n" + " # " + MARK + ": D-side. Mamba decoder recomputes the last prompt\n" + " # token and must start from h(N-1), so it pulls only N-1 tokens.\n" + " if getattr(self, '_has_mamba', False) and num_prompt_tokens > 1:\n" + " return num_prompt_tokens - 1\n" + " return num_prompt_tokens\n" + "\n" + " def _truncate_mamba_request_for_prefill(self, request: \"Request\") -> None:\n" + " # " + MARK + ": P-side. Drop the last prompt token so the prefiller\n" + " # computes h(N-1) not h(N); the decoder recomputes token N to get\n" + " # h(N). Guarded against repeated truncation on preempt/reschedule.\n" + " params = request.kv_transfer_params\n" + " if (\n" + " params is not None\n" + " and not params.get('_p_side_truncated')\n" + " and request.num_prompt_tokens > 1\n" + " ):\n" + " if request.prompt_token_ids is not None:\n" + " request.prompt_token_ids.pop()\n" + " elif getattr(request, 'prompt_embeds', None) is not None:\n" + " request.prompt_embeds = request.prompt_embeds[:-1]\n" + " else:\n" + " return\n" + " request._all_token_ids.pop()\n" + " request.num_prompt_tokens -= 1\n" + " request.max_tokens = 1\n" + " params['_p_side_truncated'] = True\n" + " logger.info('[" + MARK + "] P-side truncated req %s to N-1 for mamba prefill',\n" + " request.request_id)\n" + "\n" + + h2_anchor + ) + + # --- H3a: P-side truncation must run BEFORE the producer early-return + # (`if self.is_producer: return 0, False`). In MoRIIO WRITE the producer + # is the PREFILL leg (do_remote_decode) and returns early, so put the + # mamba truncation ahead of that guard. + h3a_old = ( + " if self.is_producer:\n" + " return 0, False\n" + "\n" + " token_ids = request.prompt_token_ids or []\n" + ) + h3a_new = ( + " # " + MARK + ": P-side prompt truncation for mamba runs BEFORE the\n" + " # producer early-return. In MoRIIO WRITE the producer is the prefill\n" + " # leg (do_remote_decode); drop its last prompt token so it computes\n" + " # h(N-1). The decoder (consumer) recomputes token N to derive h(N).\n" + " _k3_params = request.kv_transfer_params\n" + " if (\n" + " getattr(self, '_has_mamba', False)\n" + " and _k3_params is not None\n" + " and _k3_params.get('do_remote_decode')\n" + " ):\n" + " self._truncate_mamba_request_for_prefill(request)\n" + " if self.is_producer:\n" + " return 0, False\n" + "\n" + " token_ids = request.prompt_token_ids or []\n" + ) + # --- H3b: D-side WRITE return N-1 for mamba (was: len(token_ids) - num_computed) + h3b_old = ( + " if self.mode == MoRIIOMode.WRITE:\n" + " # MoriiO in write mode, no remote prefill\n" + "\n" + " return len(token_ids) - num_computed_tokens, True\n" + "\n" + " return len(token_ids) - 1 - num_computed_tokens, False\n" + ) + h3b_new = ( + " if self.mode == MoRIIOMode.WRITE:\n" + " # MoriiO in write mode, no remote prefill.\n" + " # " + MARK + ": D-side returns N-1 for mamba (decode recomputes the\n" + " # last prompt token from h(N-1)); non-mamba keeps N.\n" + " _k3_n = self._get_remote_prefill_token_count(len(token_ids))\n" + " return _k3_n - num_computed_tokens, True\n" + "\n" + " return len(token_ids) - 1 - num_computed_tokens, False\n" + ) + + for old, new, tag in [ + (h1_old, h1_new, "H1 _has_mamba"), + (h2_anchor, h2_new, "H2 helpers"), + (h3a_old, h3a_new, "H3a P-side truncate pre-producer"), + (h3b_old, h3b_new, "H3b D-side N-1"), + ]: + if old not in src: + print(f"[{MARK}] {tag}: ANCHOR MISSING", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_page_pad_fix.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_page_pad_fix.py new file mode 100644 index 00000000..5abcd0ab --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mamba_page_pad_fix.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Use the PADDED physical mamba page for MoRIIO block stride/geometry. + +PROBLEM (the GPU memory fault on the disagg producer path) + The hybrid KV allocator pads the mamba page so it matches the attention page + ("Padding mamba page size by N%"). The mamba KV cache tensor is therefore + [num_blocks, 1, 1, spec.page_size_bytes] (gpu_model_runner._reshape_kv_cache_tensors) + with block stride = spec.page_size_bytes (padded). + + But moriio_layout computed the mamba block stride / geometry page as the + UNPADDED live-data size (conv_bytes + ssm_bytes) from derive_mamba_conv_split: + - get_layer_transfer_geometry MambaSpec branch: _page = conv+ssm + - compute_mamba_block_transfer_offsets: page = conv+ssm; stride = page + + Since padded > unpadded, `lbase = block_id * unpadded_page` drifts below the true + block address by (padded-unpadded) per block. Low block IDs are ~ok, higher ones + land out of bounds. On the producer's first real disagg prefill the MoRIIO RDMA + write to these addresses touches unmapped GPU memory -> + "Memory access fault by GPU node-N on address 0x...". All ranks fault together. + (Standalone prefill/decode never hit this: the fault is only on the KV-transfer + write path.) + +FIX (surgical, 2 sites in moriio_layout.py) + 1) compute_mamba_block_transfer_offsets: stride = int(spec.page_size_bytes) + (keep conv/ssm sub-region offsets+sizes as the live payload within the page). + 2) get_layer_transfer_geometry MambaSpec branch: _page = int(spec.page_size_bytes) + so region_len = num_blocks * padded_page matches the real tensor extent. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_mamba_page_pad_fix.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-mamba-pad] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + + if "k3-mamba-pad: physical padded page" in src: + print("[k3-mamba-pad] already applied.") + return 0 + + changed = 0 + + # --- site 1: compute_mamba_block_transfer_offsets stride --- + # Handle the two shapes this may exist in (with or without the interim + # `page = int(conv_bytes + ssm_bytes)` line). + off_variants = [ + ( + " conv_bytes, ssm_bytes = split.ssm_sizes\n" + " page = int(conv_bytes + ssm_bytes)\n" + " # Sub-regions within one page: conv sub-projections, then ssm.\n" + " subregions = list(split.local_conv_offsets) # [(off,size), ...] for Q,K,V\n" + " subregions.append((int(conv_bytes), int(ssm_bytes))) # ssm follows conv\n" + "\n" + " # Byte stride between blocks = full page (state blocks are indivisible; one\n" + " # logical block == one physical page for mamba).\n" + " stride = page\n" + ), + ] + off_new = ( + " conv_bytes, ssm_bytes = split.ssm_sizes\n" + " # Sub-regions within one page: conv sub-projections, then ssm.\n" + " subregions = list(split.local_conv_offsets) # [(off,size), ...] for Q,K,V\n" + " subregions.append((int(conv_bytes), int(ssm_bytes))) # ssm follows conv\n" + "\n" + " # k3-mamba-pad: physical padded page as block stride (the tensor is\n" + " # [num_blocks,1,1,spec.page_size_bytes]; unpadded conv+ssm drifts OOB).\n" + " stride = int(spec.page_size_bytes)\n" + ) + for v in off_variants: + if v in src: + src = src.replace(v, off_new, 1) + changed += 1 + break + else: + # Fallback: just rewrite the bare `stride = page` if present. + if "\n stride = page\n" in src: + src = src.replace( + "\n stride = page\n", + "\n # k3-mamba-pad: physical padded page (was conv+ssm; drifted OOB)\n" + " stride = int(spec.page_size_bytes)\n", + 1, + ) + changed += 1 + + # --- site 2: geometry MambaSpec branch _page --- + geom_old = ( + " _split = _k3_derive_split(spec, local_tp=1)\n" + " _conv_bytes, _ssm_bytes = _split.ssm_sizes\n" + " _page = int(_conv_bytes + _ssm_bytes)\n" + ) + geom_new = ( + " # k3-mamba-pad: physical padded page (matches the reshaped tensor\n" + " # [num_blocks,1,1,page_size_bytes]); unpadded conv+ssm under-registers.\n" + " _page = int(spec.page_size_bytes)\n" + ) + if geom_old in src: + src = src.replace(geom_old, geom_new, 1) + changed += 1 + + if changed == 0: + print("[k3-mamba-pad] WARN: no anchors matched -- mamba page may still use " + "unpadded size. Review moriio_layout.py.") + return 0 + + if src != orig: + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-mamba-pad] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print(f"[k3-mamba-pad] applied {changed} site(s): mamba block stride/page " + "-> spec.page_size_bytes (padded).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mla_boundary_clamp.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mla_boundary_clamp.py new file mode 100644 index 00000000..1e964e65 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_mla_boundary_clamp.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""ROOT-CAUSE FIX: Kimi-K3 MLA prefill/decode last-block boundary garbage. + +Disagg (P/D) serving of Kimi-K3 uses a "mamba N-1" scheme: prefill computes +only the first N-1 of N prompt tokens (KDA recurrence needs h(N-1); decode +recomputes token N locally). The MLA/attention KV, however, is still copied +one WHOLE block at a time (full geometry.block_len bytes) by the MoRIIO RDMA +writer. Because prefill only wrote N-1 token slots, the final (partial) MLA +block's slot for the boundary token (position N-1) is UNINITIALIZED producer +HBM. The whole-block copy ships that garbage into decode's paged slot N-1, +racing + clobbering decode's own correct recompute -> nondeterministic wrong +exact-recall. + +FIX: thread valid_tokens (= req.num_prompt_tokens) from the scheduler write +leg down to the layout offset builder, and clamp the FINAL MLA block's copy +size to only the VALID slots. The stale boundary slot is then never +transferred and decode's local recompute is the sole writer. The MambaSpec +branch returns BEFORE the clamp, so KDA/mamba transfers are untouched. When +valid_tokens is None (non-mamba / missing) there is no clamp -> byte-identical +to current behavior. Idempotent, anchor-based, two-pass (verify-all-then-write), +py_compile-checked. Usage: apply_kimik3_moriio_mla_boundary_clamp.py +""" +import os +import sys + +MARK = "k3-mla-boundary" + +_MORIIO = "distributed/kv_transfer/kv_connector/v1/moriio" +COMMON = _MORIIO + "/moriio_common.py" +CONN = _MORIIO + "/moriio_connector.py" +ENGINE = _MORIIO + "/moriio_engine.py" +LAYOUT = _MORIIO + "/moriio_layout.py" + +# ---------------------------------------------------------------------------- +# moriio_common.py +# ---------------------------------------------------------------------------- +C1_old = ''' mamba_local_block_ids: list[int] | None = None # k3-mamba-blockids + enqueue_time: float = field(default_factory=time.perf_counter) + retried: int = 0''' +C1_new = ''' mamba_local_block_ids: list[int] | None = None # k3-mamba-blockids + valid_tokens: int | None = None # k3-mla-boundary + enqueue_time: float = field(default_factory=time.perf_counter) + retried: int = 0''' + +C2_old = ''' # k3-mamba-blockids: mamba KV-group [1] local slot id(s) for this req. + mamba_local_block_ids: list[int] = field(default_factory=list)''' +C2_new = ''' # k3-mamba-blockids: mamba KV-group [1] local slot id(s) for this req. + mamba_local_block_ids: list[int] = field(default_factory=list) + # k3-mla-boundary: valid prompt-token count (num_prompt_tokens) for the + # write leg; clamps the final MLA block's RDMA copy to valid slots only. + valid_tokens: int | None = None''' + +C3_old = ''' write_mode=False, + mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids + ):''' +C3_new = ''' write_mode=False, + mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids + valid_tokens: int | None = None, # k3-mla-boundary + ):''' + +C4_old = ''' _req.mamba_local_block_ids = list(mamba_local_block_ids or []) # k3-mamba-blockids + if write_mode:''' +C4_new = ''' _req.mamba_local_block_ids = list(mamba_local_block_ids or []) # k3-mamba-blockids + _req.valid_tokens = valid_tokens # k3-mla-boundary + if write_mode:''' + +# ---------------------------------------------------------------------------- +# moriio_connector.py +# ---------------------------------------------------------------------------- +N1_old = ''' meta.add_new_req( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=kv_params, + write_mode=True, + mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids + )''' +N1_new = ''' meta.add_new_req( + request_id=req_id, + local_block_ids=block_ids, + kv_transfer_params=kv_params, + write_mode=True, + mamba_local_block_ids=self._reqs_save_mamba.get(req_id, []), # k3-mamba-blockids + valid_tokens=req.num_prompt_tokens, # k3-mla-boundary + )''' + +N2_old = ''' remote_ip: str, + mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids + ) -> None: + """Schedule a block write operation.''' +N2_new = ''' remote_ip: str, + mamba_local_block_ids: list[int] | None = None, # k3-mamba-blockids + valid_tokens: int | None = None, # k3-mla-boundary + ) -> None: + """Schedule a block write operation.''' + +N3_old = ''' mamba_local_block_ids=mamba_local_block_ids, # k3-mamba-blockids + layer_name=layer_name,''' +N3_new = ''' mamba_local_block_ids=mamba_local_block_ids, # k3-mamba-blockids + valid_tokens=valid_tokens, # k3-mla-boundary + layer_name=layer_name,''' + +N4_old = ''' remote_ip=meta.remote_host, + mamba_local_block_ids=meta.mamba_local_block_ids, # k3-mamba-blockids + )''' +N4_new = ''' remote_ip=meta.remote_host, + mamba_local_block_ids=meta.mamba_local_block_ids, # k3-mamba-blockids + valid_tokens=getattr(meta, "valid_tokens", None), # k3-mla-boundary + )''' + +N5_old = ''' remote_moriio_meta: MoRIIOAgentMetadata, + remote_tp_size: int | None = None, + ) -> tuple[list[int], list[int], list[int]]:''' +N5_new = ''' remote_moriio_meta: MoRIIOAgentMetadata, + remote_tp_size: int | None = None, + valid_tokens: int | None = None, # k3-mla-boundary + ) -> tuple[list[int], list[int], list[int]]:''' + +N6_old = ''' remote_num_blocks=remote_moriio_meta.num_blocks, + merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks(''' +N6_new = ''' remote_num_blocks=remote_moriio_meta.num_blocks, + valid_tokens=valid_tokens, # k3-mla-boundary + merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks(''' + +# ---------------------------------------------------------------------------- +# moriio_engine.py +# ---------------------------------------------------------------------------- +E1_old = ''' offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + _k3_local, + _k3_remote, + remote_moriio_meta, + )''' +E1_new = ''' offsets = self.worker._compute_block_transfer_offsets( + task.layer_name, + _k3_local, + _k3_remote, + remote_moriio_meta, + valid_tokens=getattr(task, "valid_tokens", None), # k3-mla-boundary + )''' + +# ---------------------------------------------------------------------------- +# moriio_layout.py +# ---------------------------------------------------------------------------- +L1_old = ''' remote_num_blocks: int, + merge_fn: Callable[ + [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] + ] = merge_contiguous_offsets, +) -> tuple[list[int], list[int], list[int]]:''' +L1_new = ''' remote_num_blocks: int, + merge_fn: Callable[ + [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] + ] = merge_contiguous_offsets, + valid_tokens: int | None = None, # k3-mla-boundary +) -> tuple[list[int], list[int], list[int]]:''' + +L2_old = ''' sizes = [transfer_size_byte] * total + + w = 0''' +L2_new = ''' sizes = [transfer_size_byte] * total + # k3-mla-boundary: prefill wrote only `valid_tokens` slots; the last partial + # block's tail (past the last valid token) is uninitialized producer HBM. + # Clamp the final block's copy to valid slots so the stale boundary slot is + # never transferred (decode's local recompute is then the sole writer). + if valid_tokens is not None and local_block_ids: + _bs = geometry.block_size + _valid_in_last = valid_tokens - (len(local_block_ids) - 1) * _bs + if 0 < _valid_in_last < _bs: + _clamped = _valid_in_last * geometry.slot_size_bytes + _last = (len(local_block_ids) - 1) * per_block + for _j in range(per_block): + sizes[_last + _j] = _clamped + import logging as _k3lg + _k3lg.getLogger(__name__).info( + "[k3-mla-boundary] clamped last block to %d/%d slots (%d B)", + _valid_in_last, _bs, _clamped, + ) + + w = 0''' + +EDITS = { + COMMON: [ + (C1_old, C1_new, "C1 WriteTask.valid_tokens field"), + (C2_old, C2_new, "C2 ReqMeta.valid_tokens field"), + (C3_old, C3_new, "C3 add_new_req signature"), + (C4_old, C4_new, "C4 add_new_req body setter"), + ], + CONN: [ + (N1_old, N1_new, "N1 build_connector_meta write-mode call"), + (N2_old, N2_new, "N2 schedule_write_blocks signature"), + (N3_old, N3_new, "N3 WriteTask construction"), + (N4_old, N4_new, "N4 _write_blocks_for_req call"), + (N5_old, N5_new, "N5 _compute_block_transfer_offsets signature"), + (N6_old, N6_new, "N6 forward into compute_block_transfer_offsets"), + ], + ENGINE: [ + (E1_old, E1_new, "E1 _prepare_transfer_plan call"), + ], + LAYOUT: [ + (L1_old, L1_new, "L1 compute_block_transfer_offsets signature"), + (L2_old, L2_new, "L2 MLA last-block clamp"), + ], +} + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + base = sys.argv[1] + + # Pass 1: load every target, skip already-applied files, and verify that + # EVERY anchor is present BEFORE writing anything. Any missing anchor aborts + # with return 1 and zero partial edits (safe-by-construction). + plans = [] + for rel, edits in EDITS.items(): + path = os.path.join(base, rel) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] {rel}: already applied, skipping.") + continue + for old, new, tag in edits: + if old not in src: + print(f"[{MARK}] {rel}: {tag}: ANCHOR MISSING", file=sys.stderr) + return 1 + plans.append((path, rel, src, edits)) + + if not plans: + print(f"[{MARK}] all target files already applied; nothing to do.") + return 0 + + # Pass 2: all anchors verified -> apply + write + py_compile each file. + for path, rel, src, edits in plans: + for old, new, tag in edits: + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {path}: {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied {rel}") + + print(f"[{MARK}] applied to {len(plans)} file(s).") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_pod_hosts.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_pod_hosts.py new file mode 100644 index 00000000..5b7e4497 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_pod_hosts.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +"""Advertise the peer pool's per-DP-pod node IPs so multi-NODE decode/prefill DP works. + +ROOT CAUSE (multi-node disagg KV transfer, traced 2026-08-04): + In 2P/2D wide-EP the decode (and prefill) pool spans >1 physical node: a master + (rank 0, has the api-server) + a `--headless` worker (rank 1) that joins the + master's DP group via --data-parallel-address. The headless worker is INVISIBLE + to router service-discovery (only the master's api-server registers) and its IP + is never conveyed to the peer pool. + + prefill's all-to-all handshake (_background_moriio_handshake) resolves each remote + decode rank's host via `pod_hosts[pod_index]`, where + pod_hosts = meta.multi_pod_hosts if set else [meta.remote_host] (line ~1624) + and multi_pod_hosts is populated ONLY from kv_transfer_params["remote_hosts"], + which the router never sets in WRITE mode. So pod_hosts collapses to the single + master host, and prefill dials EVERY decode rank at the master IP. The rank(s) + living on the worker node are never reached -> their KV is never written -> those + decode ranks generate WITHOUT the prompt context. Symptom: clean alternation + (DP2: 50% wrong; DP8: ~88% wrong) while standalone forward is perfectly coherent. + +FIX (static, launcher-advertised): + The launcher already knows every node IP. It passes the peer pool's per-pod host + list (ordered by pod index = global_dp_rank // dp_size_local) into + kv_connector_extra_config["moriio_pod_hosts"] (comma-separated). This patcher makes + the handshake fall back to THAT list when meta.multi_pod_hosts is empty, instead of + the single master host. Topology-independent: fixes TP8/DP2 AND TP2/DP8. + + Precedence unchanged: an explicit meta.multi_pod_hosts (router-provided) still wins; + the env list is only the fallback that replaces the single-host default. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_pod_hosts.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-podhosts] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + + if "k3-podhosts" in src: + print("[k3-podhosts] already applied.") + return 0 + + anchor = ( + " pod_hosts = list(meta.multi_pod_hosts) if meta.multi_pod_hosts else [host]\n" + " remote_dp_size_local = int(meta.remote_dp_size_local) or remote_dp_size\n" + ) + repl = ( + " # k3-podhosts: fall back to the launcher-advertised peer pod-host\n" + " # list (kv_connector_extra_config['moriio_pod_hosts'], ordered by\n" + " # pod index) instead of the single master host, so prefill can reach\n" + " # decode ranks that live on a HEADLESS worker node (invisible to the\n" + " # router). meta.multi_pod_hosts is [remote_host] (a single-host DEFAULT)\n" + " # when the router omits remote_hosts, so we must prefer the launcher list\n" + " # whenever it advertises MORE pods than meta -- not only when meta is empty.\n" + " _k3_meta_hosts = list(meta.multi_pod_hosts) if meta.multi_pod_hosts else []\n" + " _k3_ph = (\n" + " self.kv_transfer_config.kv_connector_extra_config.get(\n" + " \"moriio_pod_hosts\", \"\"\n" + " )\n" + " if hasattr(self, \"kv_transfer_config\")\n" + " else \"\"\n" + " )\n" + " _k3_ph = [h.strip() for h in str(_k3_ph).split(\",\") if h.strip()]\n" + " if len(_k3_ph) > len(_k3_meta_hosts):\n" + " pod_hosts = _k3_ph\n" + " logger.info(\n" + " \"[k3-podhosts] using launcher pod_hosts=%s (meta had %s)\",\n" + " pod_hosts, _k3_meta_hosts,\n" + " )\n" + " elif _k3_meta_hosts:\n" + " pod_hosts = _k3_meta_hosts\n" + " else:\n" + " pod_hosts = [host]\n" + " # k3-podhosts: remote_dp_size_local (DP ranks PER POD) drives\n" + " # pod_index = global_dp_rank // dp_local. The router never sends it, so\n" + " # meta.remote_dp_size_local defaults to remote_dp_size (the GLOBAL size)\n" + " # -> pod_index collapses to 0 -> every rank resolves to pod 0 (master),\n" + " # so rank1's KV is written to the master not the worker. Derive the true\n" + " # per-pod local size from len(pod_hosts): dp_local = remote_dp_size //\n" + " # num_pods. With 2 pods and remote_dp_size 2 -> dp_local 1 -> pod_index\n" + " # (1,1)=1 -> pod_hosts[1]=worker. Honor an explicit meta value if the\n" + " # router ever sets one (< remote_dp_size).\n" + " _k3_npods = max(1, len(pod_hosts))\n" + " if int(meta.remote_dp_size_local) and int(meta.remote_dp_size_local) < remote_dp_size:\n" + " remote_dp_size_local = int(meta.remote_dp_size_local)\n" + " elif remote_dp_size % _k3_npods == 0:\n" + " remote_dp_size_local = remote_dp_size // _k3_npods\n" + " logger.info(\n" + " \"[k3-podhosts] derived remote_dp_size_local=%d (dp_size=%d, pods=%d)\",\n" + " remote_dp_size_local, remote_dp_size, _k3_npods,\n" + " )\n" + ) + if anchor not in src: + print("[k3-podhosts] WARN anchor not found -- not applied.", file=sys.stderr) + return 0 + src = src.replace(anchor, repl, 1) + + # --- SYMMETRIC FIX on the decode->prefill NOTIFY path (send_notify_block) --- + # Same two bugs: _dp_local and _remote_hosts come from kv_transfer_params, which + # the router never sets -> _dp_local=0 (no per-pod host resolution) and + # _notify_host stays the single prefill master -> notify for prefill rank1 lands + # on the master -> prefill rank1's write_ready_flags never set -> prefill times + # out ("Timed out waiting for write_ready_flags") -> EngineDead. Feed the + # launcher-advertised prefill pod-hosts (extra_config['moriio_pod_hosts']) and a + # derived _dp_local here too. + notify_anchor = ( + " _notify_host = remote_host\n" + " _kvp = request.kv_transfer_params or {}\n" + " _remote_hosts = _kvp.get(\"remote_hosts\") or []\n" + " if _dp_local > 0 and _remote_hosts:\n" + " _pod_idx = pod_index(remote_dp_rank, _dp_local)\n" + " if 0 <= _pod_idx < len(_remote_hosts):\n" + " _notify_host = _remote_hosts[_pod_idx]\n" + ) + notify_repl = ( + " _notify_host = remote_host\n" + " _kvp = request.kv_transfer_params or {}\n" + " _remote_hosts = _kvp.get(\"remote_hosts\") or []\n" + " # k3-podhosts: fall back to launcher pod-hosts + derived dp_local\n" + " if not _remote_hosts and hasattr(self, \"kv_transfer_config\"):\n" + " _k3n = self.kv_transfer_config.kv_connector_extra_config.get(\n" + " \"moriio_pod_hosts\", \"\"\n" + " )\n" + " _remote_hosts = [h.strip() for h in str(_k3n).split(\",\") if h.strip()]\n" + " _k3_dpl = _dp_local\n" + " if (not _k3_dpl) and _remote_hosts and _dp_size % len(_remote_hosts) == 0:\n" + " _k3_dpl = _dp_size // len(_remote_hosts)\n" + " if _k3_dpl > 0 and _remote_hosts:\n" + " _pod_idx = pod_index(remote_dp_rank, _k3_dpl)\n" + " if 0 <= _pod_idx < len(_remote_hosts):\n" + " _notify_host = _remote_hosts[_pod_idx]\n" + " _remote_dp_rank_for_port = fold_local_rank(remote_dp_rank, _k3_dpl)\n" + " logger.info(\n" + " \"[k3-podhosts] notify prefill rank=%d -> host=%s (dp_local=%d)\",\n" + " remote_dp_rank, _notify_host, _k3_dpl,\n" + " )\n" + ) + if notify_anchor in src: + src = src.replace(notify_anchor, notify_repl, 1) + else: + print("[k3-podhosts] WARN notify anchor not found -- notify path unpatched.", + file=sys.stderr) + + # --- 3rd site: _write_blocks_for_req stashes multi_pod_hosts + remote_dp_size_local + # on the WORKER; MoRIIOEngine._execute_write_task/_finalize_if_complete read these + # to target the "write_done" completion notify back to the correct PREFILL pod. + # Same wrong meta defaults ([single host] / global dp_size) -> completion notify + # for prefill rank1 lands on the master -> prefill rank1 write_ready never set -> + # "Timed out waiting for write_ready_flags" -> EngineDead. Prefer the launcher + # list here too (worker HAS self.kv_transfer_config.kv_connector_extra_config). + stash_anchor = ( + " if meta.multi_pod_hosts:\n" + " self.multi_pod_hosts = list(meta.multi_pod_hosts)\n" + " else:\n" + " self.multi_pod_hosts = [meta.remote_host]\n" + " if meta.remote_dp_size_local:\n" + " self.remote_dp_size_local = int(meta.remote_dp_size_local)\n" + " else:\n" + " self.remote_dp_size_local = int(meta.remote_dp_size)\n" + ) + stash_repl = ( + " # k3-podhosts: prefer launcher-advertised peer pod-hosts + derived\n" + " # dp_local over the single-host / global-dp_size meta defaults, so the\n" + " # write_done completion notify targets the right prefill pod/rank.\n" + " _k3_meta_ph = list(meta.multi_pod_hosts) if meta.multi_pod_hosts else []\n" + " _k3_lp = (\n" + " self.kv_transfer_config.kv_connector_extra_config.get(\n" + " \"moriio_pod_hosts\", \"\"\n" + " )\n" + " if hasattr(self, \"kv_transfer_config\")\n" + " else \"\"\n" + " )\n" + " _k3_lp = [h.strip() for h in str(_k3_lp).split(\",\") if h.strip()]\n" + " if len(_k3_lp) > len(_k3_meta_ph):\n" + " self.multi_pod_hosts = _k3_lp\n" + " elif _k3_meta_ph:\n" + " self.multi_pod_hosts = _k3_meta_ph\n" + " else:\n" + " self.multi_pod_hosts = [meta.remote_host]\n" + " _k3_gdp = int(meta.remote_dp_size)\n" + " _k3_npods = max(1, len(self.multi_pod_hosts))\n" + " if 0 < int(meta.remote_dp_size_local) < _k3_gdp:\n" + " self.remote_dp_size_local = int(meta.remote_dp_size_local)\n" + " elif _k3_gdp % _k3_npods == 0:\n" + " self.remote_dp_size_local = _k3_gdp // _k3_npods\n" + " else:\n" + " self.remote_dp_size_local = _k3_gdp\n" + " logger.info(\n" + " \"[k3-podhosts] write-stash multi_pod_hosts=%s remote_dp_size_local=%d\",\n" + " self.multi_pod_hosts, self.remote_dp_size_local,\n" + " )\n" + ) + if stash_anchor in src: + src = src.replace(stash_anchor, stash_repl, 1) + else: + print("[k3-podhosts] WARN write-stash anchor not found -- completion path unpatched.", + file=sys.stderr) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-podhosts] ERROR compile: {e}", file=sys.stderr) + return 1 + print("[k3-podhosts] handshake falls back to launcher-advertised peer pod_hosts.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_port_offset_tpsize.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_port_offset_tpsize.py new file mode 100644 index 00000000..c19a4e6a --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_port_offset_tpsize.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +"""Fix MoRIIO per-rank port collision when DP-local > 1 AND TP > 1 (wide-EP). + +ROOT CAUSE + MoRIIO derives each rank's side-channel/notify/handshake port as + base_port + get_port_offset(dp_rank, tp_rank[, tp_size]) + where get_port_offset(dp, tp, tp_size=1) == dp * tp_size + tp. The offset is + UNIQUE per (dp, tp) only when tp_size is the REAL tp size. Most call sites omit + it, so tp_size defaults to 1 and the offset collapses to (dp + tp). That is + collision-free only when at most one of dp/tp varies per node -- i.e. the old + TP8 x DP-local-1 shape. In the wide-EP TP2 x DP-local-4 shape the local ranks + are (dp0..3, tp0..1) and (dp0,tp1) and (dp1,tp0) both map to offset 1: + handshake_port 8405: dp0/tp1 -> 8406 AND dp1/tp0 -> 8406 (COLLISION) + => the second listener's `socket.bind('tcp://*:8406')` raises + "zmq.error.ZMQError: Address already in use" and the moriio_handshake_listener + thread dies; all engines then idle-spin on "No available shared memory + broadcast block" and the pool never reaches startup. + +FIX + Pass the real tp_size to the 5 bare get_port_offset() calls so the offset is + dp*tp_size + tp (a proper 2-D -> 1-D flattening, always unique). tp_size is + already available at every site (a local `tp_size`, self.tp_size, + self.moriio_config.tp_size, or self.worker.moriio_config.tp_size). The two + call sites that ALREADY pass a tp_size (moriio_connector.py:1538 dial + :2621 + recv-callback) are left untouched. Producer and consumer both use the same + formula so handshake/notify addressing stays matched. + + NOTE: base ports must be spaced >= dp_local*tp_size apart to avoid cross-family + overlap. Defaults handshake=8405 / notify=61005 / local_ping=61555 are spaced + by >=8, and TP2 x DP-local-4 needs only 8 slots -> no cross-family overlap. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_port_offset_tpsize.py +""" +import os +import sys + +MORIIO = "distributed/kv_transfer/kv_connector/v1/moriio" + + +def _patch(path, subs, tag): + if not os.path.isfile(path): + print(f"[k3-portoff] {tag}: {path} not found -- skip.") + return True + src = open(path).read() + orig = src + for old, new, note in subs: + if new in src and old not in src: + continue # already applied + if old not in src: + print(f"[k3-portoff] {tag}: anchor NOT found ({note}):\n {old[:78]!r}", + file=sys.stderr) + return False + src = src.replace(old, new, 1) + if src == orig: + print(f"[k3-portoff] {tag}: no change (already applied).") + return True + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-portoff] {tag}: compile failed, rolled back: {e}", file=sys.stderr) + return False + print(f"[k3-portoff] {tag}: applied.") + return True + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + base = os.path.join(sys.argv[1], MORIIO) + + common_ok = _patch( + os.path.join(base, "moriio_common.py"), + [( + " port_offset = get_port_offset(dp_rank, tp_rank)\n", + " port_offset = get_port_offset(dp_rank, tp_rank, tp_size) # k3-portoff\n", + "common notify base offset", + )], + "moriio_common.py", + ) + + conn_ok = _patch( + os.path.join(base, "moriio_connector.py"), + [ + ( + " target_port = remote_notify_port + get_port_offset(remote_dp_rank, tp_index)\n", + " target_port = remote_notify_port + get_port_offset(remote_dp_rank, tp_index, self.tp_size) # k3-portoff\n", + "release-notify target", + ), + ( + " target_port = remote_notify_port + get_port_offset(\n" + " _remote_dp_rank_for_port, tp_index\n" + " )\n", + " target_port = remote_notify_port + get_port_offset(\n" + " _remote_dp_rank_for_port, tp_index, self.tp_size # k3-portoff\n" + " )\n", + "block-notify target", + ), + ( + " self.side_channel_port: int = (\n" + " self.moriio_config.handshake_port\n" + " + get_port_offset(self.dp_rank, self.tp_rank)\n" + " )\n", + " self.side_channel_port: int = (\n" + " self.moriio_config.handshake_port\n" + " + get_port_offset(self.dp_rank, self.tp_rank, self.moriio_config.tp_size) # k3-portoff\n" + " )\n", + "worker side_channel_port bind", + ), + ], + "moriio_connector.py", + ) + + eng_ok = _patch( + os.path.join(base, "moriio_engine.py"), + [( + " remote_port = remote_notify_port + get_port_offset(\n" + " _decode_dp_rank_for_port, self.worker.tp_rank\n" + " )\n", + " remote_port = remote_notify_port + get_port_offset(\n" + " _decode_dp_rank_for_port, self.worker.tp_rank, self.worker.moriio_config.tp_size # k3-portoff\n" + " )\n", + "engine notify target", + )], + "moriio_engine.py", + ) + + return 0 if (common_ok and conn_ok and eng_ok) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_remote_tp_fix.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_remote_tp_fix.py new file mode 100644 index 00000000..a61f9a70 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_remote_tp_fix.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""ROOT-CAUSE FIX: remote_tp_size=1 collapses all prefill ranks to decode tp0. + +Proven by K3_HS_BC=1: every prefill worker (self_tp=0..7) resolves dial_tp=0 +because meta.tp_size (remote_tp_size) arrives as 1 (the router does not advertise +the decode pool's TP size for WRITE producer requests). get_moriio_remote_tp_rank( +k, local=8, remote=1) = k // (8//1) = 0 for all k -> all prefill ranks dial decode +tp0's handshake port (8405), fetch tp0's base, and RDMA-write ONLY into decode +tp0. Decode tp1..7 stay zero -> 7/8 of every attention (and KDA) shard is empty on +decode -> fluent-but-context-free output. + +FIX: our P/D deployment is SYMMETRIC TP (prefill TP == decode TP == self.world_size). +When the advertised remote_tp_size is unknown/degenerate (<= 1) but we run multi-way +TP locally, treat the remote pool as the SAME TP size (self.world_size). Then +get_moriio_remote_tp_rank(k, 8, 8) = k, so prefill rank k dials decode rank k +(port 8405+k), writes to decode rank k's base, and all 8 decode shards get their KV. + +Two hunks (both in moriio_connector.py): + H1 _remote_tp_rank: extend the "0/unknown" normalization to also cover + remote_tp_size == 1 when self.world_size > 1 (symmetric-TP assumption). + H2 add_new_req is in moriio_common.py; instead of touching wire parsing we fix + at the single choke point (_remote_tp_rank) AND at the port-offset call in + _moriio_handshake, which independently uses the raw remote_tp_size for the + port math. Normalize there too so port_offset = get_port_offset(dp, k, 8). + +Idempotent, anchor-based, py_compile-checked. Symmetric-TP only (our target); +genuinely heterogeneous TP would need the router to advertise remote_tp_size. +Usage: apply_kimik3_moriio_remote_tp_fix.py +""" +import os +import sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +MARK = "k3-remote-tp-fix" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, CONN) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + + # H1: _remote_tp_rank normalization (0 OR 1 -> world_size when local TP > 1). + h1_old = ( + " def _remote_tp_rank(self, remote_tp_size: int) -> int:\n" + " # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0).\n" + " if remote_tp_size == 0:\n" + " remote_tp_size = self.world_size\n" + " return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size)\n" + ) + h1_new = ( + " def _remote_tp_rank(self, remote_tp_size: int) -> int:\n" + " # 0/unknown remote TP == homogeneous (avoids collapsing all ranks to 0).\n" + " # " + MARK + ": remote_tp_size==1 from an un-advertising router ALSO\n" + " # collapses every prefill rank to decode tp0 (k//local = 0). For our\n" + " # symmetric-TP P/D, normalize any degenerate (<=1) remote size to the\n" + " # local world_size so rank k -> decode rank k.\n" + " if remote_tp_size <= 1 and self.world_size > 1:\n" + " remote_tp_size = self.world_size\n" + " elif remote_tp_size == 0:\n" + " remote_tp_size = self.world_size\n" + " return get_moriio_remote_tp_rank(self.tp_rank, self.world_size, remote_tp_size)\n" + ) + + # H2: _moriio_handshake port-offset uses the raw remote_tp_size arg; normalize + # it the same way so the dialed port = 8405 + rank (not 8405 + 0). + h2_old = ( + " dial_tp_rank = (\n" + " self._remote_tp_rank(remote_tp_size)\n" + " if remote_tp_rank is None\n" + " else int(remote_tp_rank)\n" + " )\n" + " port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size)\n" + ) + h2_new = ( + " dial_tp_rank = (\n" + " self._remote_tp_rank(remote_tp_size)\n" + " if remote_tp_rank is None\n" + " else int(remote_tp_rank)\n" + " )\n" + " # " + MARK + ": normalize degenerate remote_tp_size for the port math\n" + " # too, so prefill rank k dials decode port base+k (symmetric TP).\n" + " _k3_rts = remote_tp_size\n" + " if _k3_rts <= 1 and self.world_size > 1:\n" + " _k3_rts = self.world_size\n" + " port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, _k3_rts)\n" + ) + + for old, new, tag in [(h1_old, h1_new, "H1 _remote_tp_rank"), + (h2_old, h2_new, "H2 handshake port_offset")]: + if old not in src: + print(f"[{MARK}] {tag}: ANCHOR MISSING", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_save_skip_mamba.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_save_skip_mamba.py new file mode 100644 index 00000000..69d8402d --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_save_skip_mamba.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Skip the MoRIIO per-layer save hook for KDA/mamba layers under MORIIO_SKIP_MAMBA=1. + +The attention forward calls connector.save_kv_layer(layer_name, kv_cache, ...) after +EACH attention-family layer (kv_transfer_utils.maybe_transfer_kv_layer). K3 layers 0-3 +are KDA (mamba), so the very first save_kv_layer of the forward runs on a mamba cache: +it does the remote handshake and schedules an INLINE RDMA write (schedule_write_blocks +-> torch.cuda.Event().record on the layer tensor) on the compute stream, mid-forward. +Under 2P/2D that collides with the MoRI-EP all2all and faults all ranks before any KDA +compute (observed: fault precedes the KDA conv debug hook). + +This makes save_kv_layer a no-op for mamba layers when MORIIO_SKIP_MAMBA=1, isolating +that path (the mamba-offset skip alone did NOT cover the save-hook entry/handshake). + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_moriio_save_skip_mamba.py +""" +import os +import sys + +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-save-skip] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + if "k3-kda: MORIIO_SKIP_MAMBA" in src and "save hook" in src: + print("[k3-save-skip] already applied.") + return 0 + + # Anchor on the worker save_kv_layer body prologue (the version that takes + # metadata as first arg -> has the `if self.mode == MoRIIOMode.READR` guard + # followed by `remote_engine_id = None`). + anchor = ( + " if not self.is_producer:\n" + " return\n" + " if self.mode == MoRIIOMode.READ:\n" + " return\n" + " remote_engine_id = None\n" + ) + inject = ( + " if not self.is_producer:\n" + " return\n" + " if self.mode == MoRIIOMode.READ:\n" + " return\n" + " # k3-kda: MORIIO_SKIP_MAMBA=1 -> skip the per-layer save hook for\n" + " # KDA/mamba layers (inline RDMA write on the compute stream at the\n" + " # first layers races the MoRI-EP all2all -> all-rank GPU fault).\n" + " import os as _os_k3s\n" + " if _os_k3s.environ.get(\"MORIIO_SKIP_MAMBA\", \"0\") == \"1\":\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MambaSpecSK\n" + " if isinstance(self.layer_to_spec.get(layer_name), _K3MambaSpecSK):\n" + " return\n" + " remote_engine_id = None\n" + ) + n = src.count(anchor) + if n == 0: + print("[k3-save-skip] WARN: save_kv_layer anchor not found -- not applied.") + return 0 + # Only the worker save_kv_layer has this exact prologue; replace the first. + src = src.replace(anchor, inject, 1) + + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-save-skip] ERROR: compile failed: {e}", file=sys.stderr) + open(path, "w").write(orig) + return 1 + print("[k3-save-skip] save_kv_layer now skips mamba layers under MORIIO_SKIP_MAMBA=1.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_write_fence.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_write_fence.py new file mode 100644 index 00000000..67887fb2 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_moriio_write_fence.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""FIX/DIAGNOSTIC for the RDMA-write-then-notify ORDERING HAZARD (UPDATE 20). + +Symptom: disagg recall is NON-DETERMINISTIC at greedy temp=0 (same prompt -> different +garbage each run), while colocated is deterministic+correct. Root cause: prefill +RDMA-writes KV into decode HBM, waits only for SENDER-side completion +(waiting_for_transfer_complete polls status.Succeeded()), then sends write_done over +a SEPARATE ZMQ/TCP path. Sender-local RDMA completion does not order against the +RECEIVER's HBM visibility, and TCP notify races the data landing -> decode reads +stale/partial HBM -> non-deterministic wrong recall. + +This patch inserts a sender-side ORDERING FENCE in _finalize_if_complete, right +before send_notify(write_done). Modes via env K3_WRITE_FENCE: + - 'delay' (default when enabled): sleep K3_WRITE_FENCE_MS milliseconds (default 20) + before sending write_done. DIAGNOSTIC: if recall becomes correct+deterministic, + the race is confirmed. Cheap, no MoRI API dependency. + - 'off' / unset: no change (baseline). +(A read-back RDMA fence is the proper production fix but needs the per-transfer +session+offsets; staged separately once the delay confirms the race.) + +Gated: only active when K3_WRITE_FENCE is set. Idempotent, anchor-based, +py_compile-checked. +Usage: apply_kimik3_moriio_write_fence.py +""" +import os +import sys + +ENG = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" +MARK = "k3-write-fence" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, ENG) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + # H2: replace the per-event sync with a FULL device sync before the RDMA read, + # so KV inserts on ANY stream (K3 MLA uses an aux_stream) are complete before + # the transfer reads local cache. Gated by K3_WRITE_DEVSYNC=1. This targets the + # intermittent per-position corruption (multi-stream write vs RDMA read race). + h2_old = ( + " # This event is used to synchronize the kv transfer and computation tasks.\n" + " task.event.synchronize()\n" + ) + h2_new = ( + " # This event is used to synchronize the kv transfer and computation tasks.\n" + " task.event.synchronize()\n" + " import os as _k3dsos\n" + " if _k3dsos.environ.get('K3_WRITE_DEVSYNC', '') in ('1', 'true', 'on'):\n" + " # " + MARK + ": full-device sync so aux-stream KV inserts finish\n" + " # before RDMA reads local cache (event only covers one stream).\n" + " try:\n" + " import torch as _k3t\n" + " _k3t.cuda.synchronize()\n" + " except Exception:\n" + " pass\n" + ) + old = ( + " # Send completion notification\n" + " self.worker.moriio_wrapper.send_notify(\n" + " transfer_id, remote_ip, remote_port, message_type=\"write_done\"\n" + " )\n" + ) + new = ( + " # " + MARK + ": ordering fence before write_done. The RDMA write and\n" + " # the ZMQ/TCP write_done travel different paths; sender-local RDMA\n" + " # completion does not guarantee the data is visible in the RECEIVER's\n" + " # HBM. Without a fence decode can read stale HBM (non-deterministic\n" + " # recall). 'delay' mode is the diagnostic; 'readback' the real fix.\n" + " import os as _k3wfos, time as _k3wftime\n" + " _k3wf = _k3wfos.environ.get('K3_WRITE_FENCE', '').lower()\n" + " if _k3wf in ('delay', '1', 'true', 'on'):\n" + " try:\n" + " _k3ms = float(_k3wfos.environ.get('K3_WRITE_FENCE_MS', '20'))\n" + " except Exception:\n" + " _k3ms = 20.0\n" + " _k3wftime.sleep(_k3ms / 1000.0)\n" + " # Send completion notification\n" + " self.worker.moriio_wrapper.send_notify(\n" + " transfer_id, remote_ip, remote_port, message_type=\"write_done\"\n" + " )\n" + ) + if old not in src: + print(f"[{MARK}] anchor NOT found", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + # H2: full-device-sync before RDMA read (targets multi-stream write race) + if h2_old in src: + src = src.replace(h2_old, h2_new, 1) + else: + print(f"[{MARK}] H2 devsync anchor NOT found (continuing with H1 only)", file=sys.stderr) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mxfp4_int4_guard_relax.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mxfp4_int4_guard_relax.py new file mode 100644 index 00000000..3f53c7bb --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/apply_kimik3_mxfp4_int4_guard_relax.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +"""Relax the over-strict aiter#4471 guard on K3's packed-int4 gfx942 MoE path. + +PROBLEM + On gfx942 (MI300X) there is no scaled-MXFP4 MFMA, so vLLM's mxfp4 quant layer + requantizes Kimi-K3's SiTU MXFP4 experts to groupwise int4 and serves them via + AITER's bf16 x int4 FlyDSL path (mxfp4.py::_setup_kernel_k3_situ_gfx942, enabled + by --quantization-config.moe.weight int4_per_group_32). + + That method guards on whether aiter's build-time helper + `aiter.ops.flydsl.kernels.moe_gemm_2stage.compile_moe_gemm1` exposes an `act` + parameter (ROCm/aiter#4471). If not, it RAISES: + "This AITER build ignores the SiTUv2 activation on the packed-int4 MoE path + and would silently compute SiLU. Rebuild with an AITER that includes + ROCm/aiter#4471." + + BUT the guard is a false positive on this stack: + - It inspects a build-time *compile* helper's signature, not the runtime call. + - The runtime MoE dispatch (Mxfp4MoeMethod.apply -> moe_kernel.apply) passes + `activation=layer.activation` through the standard AiterExperts fused_moe path + (mxfp4.py:~460), so SiTUv2 DOES reach the kernel at runtime. + - The PROVEN, validated colocated image (amdsiloai/vllm:kimi-k3-mi325x-release-v2) + ships this EXACT aiter (compile_moe_gemm1 with NO `act` param) and serves K3 + SiTU int4 on gfx942 with correct output. Its (older) vLLM simply lacks this + newer guard. Our from-source vLLM added the stricter check. + +FIX (surgical) + Neutralize only the guard's `raise` so the packed-int4 setup proceeds, matching + the proven image's behavior. We do NOT touch the conversion/kernel logic. + + Correctness note: this trusts the runtime activation plumbing (validated by the + proven image producing correct K3 output on identical aiter). If a future AITER + regresses the runtime SiTUv2 handling, re-enable the guard. + +Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_mxfp4_int4_guard_relax.py +""" +import os +import sys + +REL = "model_executor/layers/quantization/mxfp4.py" + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + path = os.path.join(sys.argv[1], REL) + if not os.path.isfile(path): + print(f"[k3-int4-guard] {REL} not found -- skip.") + return 0 + src = open(path).read() + orig = src + + if "k3-int4-guard: relaxed" in src: + print("[k3-int4-guard] already applied.") + return 0 + + anchor = ( + ' if "act" not in inspect.signature(compile_moe_gemm1).parameters:\n' + " raise RuntimeError(\n" + ) + repl = ( + " # k3-int4-guard: relaxed -- the proven colocated image runs this\n" + " # exact aiter (compile_moe_gemm1 without `act`) and serves K3 SiTU\n" + " # int4 correctly; the runtime dispatch passes activation through the\n" + " # AiterExperts path regardless. Warn instead of aborting.\n" + ' if "act" not in inspect.signature(compile_moe_gemm1).parameters:\n' + " import logging as _lg\n" + ' _lg.getLogger(__name__).warning(\n' + ) + if anchor in src: + src = src.replace(anchor, repl, 1) + # The original raise(...) body now feeds logger.warning(...) instead; + # that is valid Python (warning takes the same string args). Leave the + # message text and closing paren as-is. + else: + print("[k3-int4-guard] WARN: guard anchor not found -- the strict " + "aiter#4471 check may still abort. Review " + "mxfp4.py::_setup_kernel_k3_situ_gfx942.") + return 0 + + if src != orig: + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-int4-guard] ERROR: compile failed: {e}", file=sys.stderr) + # restore to avoid leaving a broken file + open(path, "w").write(orig) + return 1 + print("[k3-int4-guard] relaxed the aiter#4471 packed-int4 guard " + "(raise -> warning) in mxfp4.py") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_inputs_probe.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_inputs_probe.py new file mode 100644 index 00000000..9f77684f --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_inputs_probe.py @@ -0,0 +1,85 @@ +#!/usr/bin/env python3 +"""Log decode's actual input prep (positions / num_computed_tokens / block_table / +seq_lens) for the first few requests, to find why WRITE-mode decode produces wrong +logits from byte-perfect KV. Runs in gpu_model_runner._prepare_inputs (plain +Python, NOT the @eager_break_during_capture forward -> logging works). + +Gated K3_INPUTS_PROBE=1. Logs once per ~request via a small counter. Idempotent, +anchor-based, py_compile-checked. +Usage: apply_kimik3_decode_inputs_probe.py +""" +import os, sys + +REL = "v1/worker/gpu_model_runner.py" +MARK = "k3-inputsprobe" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, REL) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + anchor = ( + " token_indices_tensor = torch.from_numpy(token_indices)\n" + ) + inject = ( + " import os as _k3ipos\n" + " if _k3ipos.environ.get('K3_INPUTS_PROBE','0')=='1':\n" + " try:\n" + " _c=getattr(self,'_k3_ip_n',0)\n" + " if _c < 12:\n" + " self._k3_ip_n=_c+1\n" + " _nct=self.input_batch.num_computed_tokens_cpu[:num_reqs].tolist()\n" + " _pos=positions_np[:min(8,len(positions_np))].tolist()\n" + " _bt=None\n" + " try:\n" + " _bt=self.input_batch.block_table.block_table[0].get_cpu_tensor()[:num_reqs,:4].tolist()\n" + " except Exception:\n" + " try: _bt=self.input_batch.block_table[0].get_cpu_tensor()[:num_reqs,:4].tolist()\n" + " except Exception as _e2: _bt='ERR:'+repr(_e2)[:40]\n" + " _tokfed=None; _ntok=None; _nprompt=None; _nout=None; _lastprompt=None\n" + " try:\n" + " _tidx=token_indices[:min(4,len(token_indices))]\n" + " _tokfed=self.input_batch.token_ids_cpu.reshape(-1)[_tidx].tolist()\n" + " except Exception as _te: _tokfed='ERR:'+repr(_te)[:30]\n" + " try:\n" + " _rid=self.input_batch.req_ids[0]\n" + " _rs=self.requests.get(_rid) if hasattr(self,'requests') else None\n" + " if _rs is not None:\n" + " _ntok=getattr(_rs,'num_tokens',None); _nprompt=getattr(_rs,'num_prompt_tokens',None)\n" + " _ao=getattr(_rs,'_all_token_ids',None) or getattr(_rs,'all_token_ids',None)\n" + " if _ao is not None:\n" + " _nout=len(_ao)-(_nprompt or 0)\n" + " _lastprompt=list(_ao[max(0,(_nprompt or 1)-1):(_nprompt or 0)+2])\n" + " except Exception as _re: _ntok='ERR:'+repr(_re)[:30]\n" + " import logging as _k3iplg\n" + " _k3iplg.getLogger('vllm.v1.worker.gpu_model_runner').info(\n" + " '[" + MARK + "] num_reqs=%s num_sched=%s num_computed=%s pos=%s token_fed=%s num_tokens=%s num_prompt=%s num_out=%s around_boundary=%s bt=%s',\n" + " num_reqs, list(num_scheduled_tokens[:num_reqs]), _nct, _pos, _tokfed, _ntok, _nprompt, _nout, _lastprompt, _bt)\n" + " except Exception as _k3ipe:\n" + " import logging as _k3iplg2\n" + " _k3iplg2.getLogger('vllm.v1.worker.gpu_model_runner').info('[" + MARK + "] EXC %r', _k3ipe)\n" + " token_indices_tensor = torch.from_numpy(token_indices)\n" + ) + if anchor not in src: + print(f"[{MARK}] anchor NOT found", file=sys.stderr) + return 1 + src = src.replace(anchor, inject, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_recv_probe.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_recv_probe.py new file mode 100644 index 00000000..2ce300a2 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_decode_recv_probe.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Decode-side ground truth: when a WRITE completes (write_done seen), read the +norm of decode's OWN KV slots to prove the RDMA bytes actually landed. + +Write side is proven correct (non-zero source, right slots, no RDMA errors) yet +decode is context-free. This probe runs in the connector's get_finished thread +(NOT the @eager_break_during_capture forward, so logging works) and, on the first +write completion, logs the L2 norm of decode's attention slot and mamba slot. If +those are ~0, the RDMA write never reached decode's memory (transport delivered to +a wrong address / decode registered a different tensor). If non-zero, the bytes +ARE present and the bug is in how the KDA/attention kernels READ them. + +Gated K3_DECODE_RECV_PROBE=1. Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_decode_recv_probe.py +""" +import os +import sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +MARK = "k3-recvprobe" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, CONN) + if not os.path.isfile(path): + print(f"[{MARK}] not found {path}", file=sys.stderr) + return 1 + src = open(path).read() + if MARK in src: + print(f"[{MARK}] already applied.") + return 0 + old = ( + " if self.mode == MoRIIOMode.WRITE:\n" + " fresh = self.moriio_wrapper.pop_finished_write_req_ids()\n" + " # Accumulate with any completions that arrived before their\n" + " # transfer_id was registered in transfer_id_to_request_id.\n" + " self._unmatched_write_completions |= fresh\n" + " done_recving = self._unmatched_write_completions\n" + ) + new = ( + " if self.mode == MoRIIOMode.WRITE:\n" + " fresh = self.moriio_wrapper.pop_finished_write_req_ids()\n" + " # Accumulate with any completions that arrived before their\n" + " # transfer_id was registered in transfer_id_to_request_id.\n" + " self._unmatched_write_completions |= fresh\n" + " done_recving = self._unmatched_write_completions\n" + " import os as _k3rpos # " + MARK + "\n" + " if fresh and _k3rpos.environ.get('K3_DECODE_RECV_PROBE','0')=='1' and not getattr(self,'_k3_rp_done',False):\n" + " try:\n" + " self._k3_rp_done = True\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MS_RP\n" + " _att=_mam=None\n" + " for _ln,_t in self.kv_caches.items():\n" + " _ism=isinstance(self.layer_to_spec.get(_ln),_K3MS_RP)\n" + " if _ism and _mam is None: _mam=(_ln,_t)\n" + " if (not _ism) and _att is None: _att=(_ln,_t)\n" + " def _nz(t,slot):\n" + " try:\n" + " s=t[int(slot)].flatten(); \n" + " _el=[round(float(x),3) for x in s[:6].float().tolist()]\n" + " return (float(s.float().norm()), _el)\n" + " except Exception as e: return (-2.0,[])\n" + " _an=_nz(_att[1],1) if _att else (-1,[])\n" + " _mn=_nz(_mam[1],2) if _mam else (-1,[])\n" + " logger.info('[" + MARK + "] attn_slot1_norm=%.4e attn_el=%s mamba_slot2_norm=%.4e mamba_el=%s',\n" + " _an[0], _an[1], _mn[0], _mn[1])\n" + " except Exception as _e:\n" + " logger.info('[" + MARK + "] EXC %r', _e)\n" + ) + if old not in src: + print(f"[{MARK}] anchor NOT found", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] compile FAIL {e}", file=sys.stderr) + return 1 + print(f"[{MARK}] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_handshake_dial_bc.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_handshake_dial_bc.py new file mode 100644 index 00000000..9d1a8003 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_handshake_dial_bc.py @@ -0,0 +1,47 @@ +#!/usr/bin/env python3 +"""Log the ACTUAL handshake dial: self.tp_rank, dial_tp_rank, port, port_offset, +final path -- to find why all 8 prefill ranks fetch the SAME decode base (only +decode tp0's KV lands). Gated K3_HS_BC=1. Edits _moriio_handshake in +moriio_connector.py. +""" +import os, sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + + +def main(): + base = sys.argv[1] + path = os.path.join(base, CONN) + src = open(path).read() + if "k3-hsbc" in src: + print("[k3-hsbc] already applied.") + return 0 + old = ( + " port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size)\n" + " path = make_zmq_path(\"tcp\", host, port + port_offset)\n" + ) + new = ( + " port_offset = get_port_offset(remote_dp_rank, dial_tp_rank, remote_tp_size)\n" + " path = make_zmq_path(\"tcp\", host, port + port_offset)\n" + " import os as _hsos # k3-hsbc\n" + " if _hsos.environ.get('K3_HS_BC','0')=='1':\n" + " logger.info('[k3-hsbc] self_tp=%s dial_tp=%s remote_tp_size=%s remote_dp_rank=%s port=%s off=%s path=%s eid=%s',\n" + " getattr(self,'tp_rank','?'), dial_tp_rank, remote_tp_size, remote_dp_rank, port, port_offset, path, expected_engine_id)\n" + ) + if old not in src: + print("[k3-hsbc] anchor NOT found", file=sys.stderr) + return 1 + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-hsbc] compile FAIL {e}", file=sys.stderr) + return 1 + print("[k3-hsbc] applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_kda_state_probe.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_kda_state_probe.py new file mode 100644 index 00000000..01902f38 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_kda_state_probe.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Ground-truth probe: is the transferred KDA state actually present in the +slot decode reads? + +Decode generates fluent-but-context-free with disagg. Attention KV + mamba +state both *transfer* (proven by write breadcrumbs), yet output ignores the +prompt. This probe logs, on the DECODE pure-decode path for KDA layer 0, the +L2 norm of the recurrent_state and conv_state rows at the exact indices the +kernel is about to read (decode_conv_indices). If those norms are ~0, the +transferred state did NOT land in the slot decode reads (block-table / slot +mapping bug), which is the smoking gun. If they're clearly non-zero, the state +IS present and the bug is elsewhere (e.g. offset content garbled, or the gate). + +Gated by K3_KDA_STATE_PROBE=1. Logs once per ~forward via a step counter so it +doesn't spam. Idempotent, anchor-based, py_compile-checked. +Usage: apply_kimik3_kda_state_probe.py +""" +import os +import sys + +KDA = "models/kimi_k3/nvidia/kda.py" + + +def _edit(path, old, new, tag): + if not os.path.isfile(path): + print(f"[k3-stateprobe] {tag}: not found {path}", file=sys.stderr) + return False + src = open(path).read() + if "k3-stateprobe" in src: + print(f"[k3-stateprobe] {tag}: already applied.") + return True + if old not in src: + print(f"[k3-stateprobe] {tag}: anchor NOT found", file=sys.stderr) + return False + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[k3-stateprobe] {tag}: compile FAIL {e}", file=sys.stderr) + return False + print(f"[k3-stateprobe] {tag}: applied.") + return True + + +def main(): + base = sys.argv[1] + path = os.path.join(base, KDA) + # Anchor at the common point after the KDA cache tensors are bound -- covers + # BOTH decode kernels (fused_kda_decode fast path AND + # fused_recurrent_kda_packed_decode). m + non_spec_state_indices_tensor are + # already in scope here. Only probe pure-decode (num_prefills==0) so the + # norm reflects the transferred state, not a freshly-computed prefill state. + old = ( + " conv_state, recurrent_state = self.kv_cache\n" + ) + new = ( + " conv_state, recurrent_state = self.kv_cache\n" + " import os as _k3os\n" + " if _k3os.environ.get('K3_KDA_STATE_PROBE', '0') == '1':\n" + " try:\n" + " _ln = getattr(self, 'prefix', '?')\n" + " if ('.layers.1.' in str(_ln)) or ('.layers.5.' in str(_ln)):\n" + " _idx = non_spec_state_indices_tensor\n" + " if _idx is not None:\n" + " _ii = _idx[:max(1,int(num_actual_tokens))].long()\n" + " _rs = recurrent_state.index_select(0, _ii)\n" + " _cs = conv_state.index_select(0, _ii)\n" + " _rn, _ra = float(_rs.float().norm()), float(_rs.float().abs().max())\n" + " _cn, _ca = float(_cs.float().norm()), float(_cs.float().abs().max())\n" + " _idl = _idx[:4].tolist()\n" + " else:\n" + " _rn=_ra=_cn=_ca=-1.0; _idl=None\n" + " import sys as _k3sys\n" + " _hi = (None if has_initial_state is None else (has_initial_state.tolist() if hasattr(has_initial_state,'tolist') else has_initial_state))\n" + " print('[k3-stateprobe] layer=%s nprefill=%s ndecode=%s nact=%s hasinit=%s idx=%s rs_norm=%.4e rs_absmax=%.4e cs_norm=%.4e cs_absmax=%.4e' % (\n" + " _ln, m.num_prefills, m.num_decodes, num_actual_tokens, _hi, _idl, _rn, _ra, _cn, _ca), file=_k3sys.stderr, flush=True)\n" + " except Exception as _k3e:\n" + " import sys as _k3sys2\n" + " print('[k3-stateprobe] EXC %r' % (_k3e,), file=_k3sys2.stderr, flush=True)\n" + ) + return 0 if _edit(path, old, new, "kda decode state probe") else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_moriio_write_bc.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_moriio_write_bc.py new file mode 100644 index 00000000..62254731 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_moriio_write_bc.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Breadcrumbs for the MoRIIO WRITE-mode KV delivery (diagnose multi-node context-loss). + +The multi-node transport now completes (handshake OK, ack OK, no timeout) but decode +generates context-free even with MORIIO_SKIP_MAMBA=1 (so it's the MLA-KV write itself, +not KDA). Suspicion: prefill RDMA-writes to the wrong decode rank's base addr / block +offsets, so decode reads zeros for the prompt. These logs expose the actual runtime +values per request so we can see the divergence in ONE run. + +Logs (all one-line, INFO, gated so they don't spam every layer): + PREFILL write: decode_dp_rank, dst_engine_id (per-rank), first remote base addr, + #local_blocks, #remote_blocks, first few block ids. + DECODE alloc: in update_state_after_alloc WRITE branch -- remote_dp_rank, the + local block_ids decode allocated + count (what it tells prefill to + write into), num_external_tokens. + +Idempotent, anchor-based, py_compile-checked. Remove after diagnosis. +Usage: apply_kimik3_moriio_write_bc.py +""" +import os +import sys + +CONN = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" +ENG = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py" + + +def _edit(path, subs, tag): + if not os.path.isfile(path): + print(f"[k3-writebc] {tag}: not found -- skip.") + return True + src = open(path).read() + orig = src + for old, new, note in subs: + if new.split("\n")[0].strip() and new.split("\n")[1].strip() in src and old not in src: + continue + if old not in src: + print(f"[k3-writebc] {tag}: anchor NOT found ({note})", file=sys.stderr) + return False + src = src.replace(old, new, 1) + if src == orig: + print(f"[k3-writebc] {tag}: no change (already applied).") + return True + open(path, "w").write(src) + try: + import py_compile + py_compile.compile(path, doraise=True) + except Exception as e: + open(path, "w").write(orig) + print(f"[k3-writebc] {tag}: compile failed, rolled back: {e}", file=sys.stderr) + return False + print(f"[k3-writebc] {tag}: applied.") + return True + + +def main(): + if len(sys.argv) != 2: + print(f"usage: {sys.argv[0]} ", file=sys.stderr) + return 2 + base = sys.argv[1] + + # ENGINE: log the write target right after dst_engine_id gets the dp suffix + + # the session/meta are fetched. + eng_old = ( + " # Get or create sessions\n" + " sessions, remote_moriio_meta = self.worker._get_built_session(\n" + " task.dst_engine_id\n" + " )\n" + ) + eng_new = ( + " # Get or create sessions\n" + " sessions, remote_moriio_meta = self.worker._get_built_session(\n" + " task.dst_engine_id\n" + " )\n" + " try:\n" + " _k3_ba = list(getattr(remote_moriio_meta, 'kv_caches_base_addr', []) or [])[:1]\n" + " import logging as _k3lg\n" + " _k3lg.getLogger('vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine').info(\n" + " '[k3-writebc] WRITE ptp=%s decode_dp_rank=%s dst=%s base0=%s remote_ip=%s',\n" + " getattr(self.worker, 'tp_rank', '?'),\n" + " request_info.decode_dp_rank, task.dst_engine_id, _k3_ba,\n" + " getattr(request_info, 'completion_remote_ip', None),\n" + " )\n" + " except Exception:\n" + " pass\n" + ) + # ENGINE 2: log the actual transfer plan (block-id counts + first offsets + + # remote_num_blocks) so an offset/num_blocks mismatch is visible. + eng2_old = ( + " local_off, remote_off, sizes = offsets\n" + ) + eng2_new = ( + " local_off, remote_off, sizes = offsets\n" + " try:\n" + " from vllm.v1.kv_cache_interface import MambaSpec as _K3MS_BC\n" + " _k3_spec = self.worker.layer_to_spec.get(task.layer_name)\n" + " _k3_ismamba = isinstance(_k3_spec, _K3MS_BC)\n" + " _k3_seen = getattr(self, '_k3_bc_seen', None)\n" + " if _k3_seen is None:\n" + " _k3_seen = set(); self._k3_bc_seen = _k3_seen\n" + " _k3_key = ('M' if _k3_ismamba else 'A')\n" + " if _k3_key not in _k3_seen:\n" + " _k3_seen.add(_k3_key)\n" + " import logging as _k3lg2\n" + " _k3lg2.getLogger('vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine').info(\n" + " '[k3-writebc] PLAN layer=%s mamba=%s spec=%s n_local_blk=%s n_remote_blk=%s '\n" + " 'remote_num_blocks=%s loff0=%s roff0=%s sz0=%s nseg=%s',\n" + " task.layer_name, _k3_ismamba, type(_k3_spec).__name__,\n" + " len(task.local_block_ids), len(request_info.block_ids or []),\n" + " getattr(remote_moriio_meta, 'num_blocks', '?'),\n" + " (local_off[:1] if local_off else []), (remote_off[:1] if remote_off else []),\n" + " (sizes[:1] if sizes else []), len(sizes or []),\n" + " )\n" + " except Exception as _e:\n" + " pass\n" + ) + ok_eng = _edit(os.path.join(base, ENG), + [(eng_old, eng_new, "engine write target"), + (eng2_old, eng2_new, "engine transfer plan")], + "moriio_engine.py") + + # CONNECTOR: decode-side alloc in WRITE branch -- log what blocks decode allocated. + conn_old = ( + " else:\n" + " # WRITE mode, decode side: notify P that blocks are ready\n" + " assert request.kv_transfer_params is not None, (\n" + " \"kv_transfer_params should not be None\"\n" + " )\n" + "\n" + " remote_dp_rank = request.kv_transfer_params.get(\"remote_dp_rank\", 0)\n" + ) + conn_new = ( + " else:\n" + " # WRITE mode, decode side: notify P that blocks are ready\n" + " assert request.kv_transfer_params is not None, (\n" + " \"kv_transfer_params should not be None\"\n" + " )\n" + "\n" + " remote_dp_rank = request.kv_transfer_params.get(\"remote_dp_rank\", 0)\n" + " try:\n" + " _k3_lb = blocks.get_block_ids()[0]\n" + " logger.info(\n" + " '[k3-writebc] DECODE-alloc self_dp=%s remote_dp_rank=%s " + "n_local_blocks=%s first=%s num_external=%s',\n" + " getattr(self, '_global_dp_rank', '?'), remote_dp_rank,\n" + " len(_k3_lb), _k3_lb[:4], num_external_tokens,\n" + " )\n" + " except Exception:\n" + " pass\n" + ) + ok_conn = _edit(os.path.join(base, CONN), [(conn_old, conn_new, "decode alloc")], + "moriio_connector.py") + + return 0 if (ok_eng and ok_conn) else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_xfer_probe.py b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_xfer_probe.py new file mode 100644 index 00000000..7124e953 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/patchers/diagnostics/apply_kimik3_xfer_probe.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""RUNTIME PROBE (K3_XFER_PROBE=1): log MoRIIO transfer offsets + source KV checksum +on the PRODUCER (compute_block_transfer_offsets, MLA branch) so we can see the +ACTUAL local/remote byte offsets and a checksum of the local KV being sent. + +Gated by K3_XFER_PROBE=1 (default off => byte-identical). Logs at most the FIRST +few transfers per process to avoid flooding. Correlate producer log (this) with +decode-side KV checksum to determine if offsets are wrong vs decode reads wrong +blocks. Anchor-based, idempotent, py_compile-checked. +Usage: apply_kimik3_xfer_probe.py +""" +import os, sys + +MARK = "k3-xfer-probe" +REL = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py" +REL2 = "v1/attention/backends/mla/triton_mla.py" +REL3 = "distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py" + +OLD = " return merge_fn(offset_local, offset_remote, sizes)\n" +# NOTE: this OLD appears twice (mamba fn + MLA fn). We target the MLA one via a +# larger unique anchor including the preceding zip loop tail. +OLD_CTX = ( + " offset_remote[w] = element_size * (\n" + " geometry.remote_kv_stride + rb * geometry.block_stride\n" + " )\n" + " w += 1\n" + "\n" + " return merge_fn(offset_local, offset_remote, sizes)\n" +) +NEW_CTX = ( + " offset_remote[w] = element_size * (\n" + " geometry.remote_kv_stride + rb * geometry.block_stride\n" + " )\n" + " w += 1\n" + "\n" + " import os as _k3xp_os # " + MARK + "\n" + " if _k3xp_os.environ.get('K3_XFER_PROBE', '0') == '1':\n" + " try:\n" + " import logging as _k3xp_lg\n" + " _k3xp_L = _k3xp_lg.getLogger(__name__)\n" + " _k3xp_n = getattr(compute_block_transfer_offsets, '_k3xp_count', 0)\n" + " if _k3xp_n < 6 and local_block_ids:\n" + " compute_block_transfer_offsets._k3xp_count = _k3xp_n + 1\n" + " _k3xp_lb0 = int(local_block_ids[0])\n" + " _k3xp_rb0 = int(remote_block_ids[0])\n" + " _k3xp_flat = kv_cache.flatten()\n" + " _k3xp_o = offset_local[0] // element_size\n" + " _k3xp_slice = _k3xp_flat[_k3xp_o:_k3xp_o + 64]\n" + " _k3xp_cs = float(_k3xp_slice.detach().to('cpu').double().sum().item())\n" + " _k3xp_L.warning(\n" + " '[" + MARK + "] layer=%s nblk=%d lb0=%d rb0=%d block_stride=%d elt=%d '\n" + " 'off_local0=%d off_remote0=%d size0=%d src_cs=%.6f',\n" + " layer_name, len(local_block_ids), _k3xp_lb0, _k3xp_rb0,\n" + " int(geometry.block_stride), int(element_size),\n" + " int(offset_local[0]), int(offset_remote[0]), int(sizes[0]), _k3xp_cs,\n" + " )\n" + " except Exception as _k3xp_e:\n" + " import logging as _k3xp_lg2\n" + " _k3xp_lg2.getLogger(__name__).warning('[" + MARK + "] err %s', _k3xp_e)\n" + "\n" + " return merge_fn(offset_local, offset_remote, sizes)\n" +) + + +DEC_OLD = ' block_table = attn_metadata.decode.block_table\n seq_lens = attn_metadata.decode.seq_lens\n' +DEC_NEW = " block_table = attn_metadata.decode.block_table\n seq_lens = attn_metadata.decode.seq_lens\n import os as _k3dp_os # k3-xfer-probe\n if _k3dp_os.environ.get('K3_XFER_PROBE', '0') == '1':\n try:\n import logging as _k3dp_lg\n _k3dp_n = getattr(decode_attention_fwd, '_k3dp_count', 0)\n if _k3dp_n < 8 and block_table is not None and block_table.numel() > 0:\n decode_attention_fwd._k3dp_count = _k3dp_n + 1\n _k3dp_b0 = int(block_table.flatten()[0].item())\n _k3dp_sl = int(seq_lens.flatten()[0].item()) if seq_lens is not None and seq_lens.numel() else -1\n _k3dp_blk = kv_c_and_k_pe_cache[_k3dp_b0].flatten()\n _k3dp_cs = float(_k3dp_blk[:64].detach().to('cpu').double().sum().item())\n _k3dp_nz = int((_k3dp_blk != 0).sum().item())\n _k3dp_lg.getLogger(__name__).warning(\n '[k3-xfer-probe DECODE] read_block0=%d seq_len0=%d dst_cs=%.6f nonzero=%d/%d',\n _k3dp_b0, _k3dp_sl, _k3dp_cs, _k3dp_nz, int(_k3dp_blk.numel()),\n )\n except Exception as _k3dp_e:\n import logging as _k3dp_lg2\n _k3dp_lg2.getLogger(__name__).warning('[k3-xfer-probe DECODE] err %s', _k3dp_e)\n" + + +GBI_OLD = ' _k3_gbi_d = blocks.get_block_ids() # k3-mamba-blockids\n block_notify_list = (\n _k3_gbi_d[0] if num_external_tokens > 0 else []\n )\n' +GBI_NEW = " _k3_gbi_d = blocks.get_block_ids() # k3-mamba-blockids\n import os as _k3gp_os # k3-xfer-probe\n if _k3gp_os.environ.get('K3_XFER_PROBE','0')=='1':\n try:\n import logging as _k3gp_lg\n _k3gp_lg.getLogger(__name__).warning(\n '[k3-xfer-probe GBI] ngroups=%d group0=%s group1=%s num_ext=%d',\n len(_k3_gbi_d),\n str([int(x) for x in _k3_gbi_d[0][:6]]) if len(_k3_gbi_d)>0 else '[]',\n str([int(x) for x in _k3_gbi_d[1][:6]]) if len(_k3_gbi_d)>1 else '[]',\n int(num_external_tokens),\n )\n except Exception as _k3gp_e:\n import logging as _k3gp_lg2\n _k3gp_lg2.getLogger(__name__).warning('[k3-xfer-probe GBI] err %s', _k3gp_e)\n block_notify_list = (\n _k3_gbi_d[0] if num_external_tokens > 0 else []\n )\n" + + +GD_OLD = ' self.layer_to_spec = build_layer_to_spec(kv_cache_config)\n' +GD_NEW = " self.layer_to_spec = build_layer_to_spec(kv_cache_config)\n import os as _k3grp_os # k3-xfer-probe\n if _k3grp_os.environ.get('K3_XFER_PROBE','0')=='1':\n try:\n import logging as _k3grp_lg\n for _gi, _g in enumerate(kv_cache_config.kv_cache_groups):\n _sp = _g.kv_cache_spec\n _lns = list(getattr(_g, 'layer_names', []) or [])\n _k3grp_lg.getLogger(__name__).warning(\n '[k3-xfer-probe GROUP] idx=%d spec=%s nlayers=%d layers0_3=%s',\n _gi, type(_sp).__name__, len(_lns), str(_lns[:4]),\n )\n except Exception as _k3grp_e:\n import logging as _k3grp_lg2\n _k3grp_lg2.getLogger(__name__).warning('[k3-xfer-probe GROUP] err %s', _k3grp_e)\n" + + +MB_OLD = ' w += 1\n return merge_fn(offset_local, offset_remote, sizes)\n' +MB_NEW = " w += 1\n import os as _k3mp_os # k3-xfer-probe MAMBA\n if _k3mp_os.environ.get('K3_XFER_PROBE','0')=='1':\n try:\n import logging as _k3mp_lg\n _k3mp_n = getattr(compute_mamba_block_transfer_offsets,'_k3mp_c',0)\n if _k3mp_n < 6 and local_block_ids:\n compute_mamba_block_transfer_offsets._k3mp_c = _k3mp_n+1\n _k3mp_lg.getLogger(__name__).warning(\n '[k3-xfer-probe MAMBA] layer=%s nblk=%d lb=%s rb=%s stride=%d nsub=%d off0=%d sz0=%d',\n layer_name, len(local_block_ids),\n str([int(x) for x in local_block_ids[:6]]),\n str([int(x) for x in remote_block_ids[:6]]),\n int(stride), len(subregions), int(offset_local[0]) if offset_local else -1,\n int(sizes[0]) if sizes else -1,\n )\n except Exception as _k3mp_e:\n import logging as _k3mp_l2\n _k3mp_l2.getLogger(__name__).warning('[k3-xfer-probe MAMBA] err %s', _k3mp_e)\n return merge_fn(offset_local, offset_remote, sizes)\n" + + +def main(): + if len(sys.argv) < 2: + print(f"[{MARK}] usage: {sys.argv[0]} ", file=sys.stderr) + return 1 + base = sys.argv[1] + edits = [(REL, OLD_CTX, NEW_CTX, "producer-offsets"), + (REL2, DEC_OLD, DEC_NEW, "decode-read"), + (REL3, GBI_OLD, GBI_NEW, "gbi-groups"), + (REL3, GD_OLD, GD_NEW, "group-dump"), + (REL, MB_OLD, MB_NEW, "mamba-offsets")] + for rel, old, new, tag in edits: + path = os.path.join(base, rel) + if not os.path.isfile(path): + print(f"[{MARK}] {tag}: not found {path}", file=sys.stderr); return 1 + src = open(path).read() + # per-edit idempotency: skip only if THIS edit's unique tag marker present + _tagmark = "k3-xfer-probe " + tag.upper().split("-")[0] + if old not in src: + if _tagmark in src or new[:60] in src: + print(f"[{MARK}] {tag}: already applied."); continue + print(f"[{MARK}] {tag}: ANCHOR MISSING", file=sys.stderr); return 1 + src = src.replace(old, new, 1) + open(path, "w").write(src) + try: + import py_compile; py_compile.compile(path, doraise=True) + except Exception as e: + print(f"[{MARK}] {tag}: compile FAIL {e}", file=sys.stderr); return 1 + print(f"[{MARK}] {tag}: applied.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d.sh b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d.sh new file mode 100755 index 00000000..4a792eb6 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d.sh @@ -0,0 +1,393 @@ +#!/bin/bash +# Kimi-K3 MXFP4 2P/2D wide-EP DISAGGREGATED serve: DP/EP16 per role, MoRI-EP +# (mori all2all) + MoRIIO connector (prefill->decode KV + KDA state transfer). +# +# Topology (4 nodes, 8 GPU each = EP16 per pool): +# Prefill pool: P-master (rank0, proxy+kv_producer) + P-worker (rank8, headless) +# Decode pool: D-master (kv_consumer) + D-worker (headless) +# Run this per node with ROLE + the shared *_ADDR env set (see run_2p2d_launch.sh). +# +# Applies the KDA MoRIIO patchers at container start (idempotent) so the MoRIIO +# connector carries K3's ~69 KDA (GDN) recurrent+conv state, not just MLA KV. +set -euo pipefail + +IMAGE="${IMAGE:-kimik3-wideep-disagg:latest}" +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR to your Kimi-K3-MXFP4 weights path}" +ROLE="${ROLE:?ROLE=prefill_master|prefill_worker|decode_master|decode_worker}" +PMASTER="${PMASTER:?prefill master eth0 IP}" +DMASTER="${DMASTER:?decode master eth0 IP}" +PROXY_IP="${PROXY_IP:-$PMASTER}" +PATCHER_DIR="${PATCHER_DIR:-$HOME/k3disagg/patchers}" +LOGHOST="${LOGHOST:-$HOME/k3disagg/logs}"; mkdir -p "$LOGHOST" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +# Bound the K3 MoE profiling M. On gfx942 the tuned FlyDSL a8w4 configs are +# gfx950-only + unsharded (896 experts), so our EP16 (56-expert) shard ALWAYS +# falls to the heuristic FlyDSL kernel. That heuristic crashes LLVM codegen +# ("Do not know how to expand this operator's operand") at the giant profiling +# shape sorted-M=131072 (= max_num_batched_tokens 16384 x topk 8). Shrinking +# max-num-batched-tokens shrinks the profiling M so the heuristic kernel compiles. +MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-2048}" +# K3 MoE requant path. gfx942 has NO scaled-MXFP4 MFMA and the a16w4 SiTUv2 +# heuristic FlyDSL kernel CANNOT codegen on gfx942 (LLVM ExpandIntegerOperand on +# a 128-bit buffer->LDS async load; all tuned a8w4 configs are gfx950-only). The +# PROVEN-coherent colocated path (logbook: NIAH 3/3 @9600) requants MoE to +# packed-int4 and runs it through Situv2 (dtype torch.int4, per_1x32) with +# AITER_SITUV2_A8W4=1 + AITER PR#4471 (SiTUv2 in the int4 stage1 epilogue). +QUANT_CONFIG="${QUANT_CONFIG:-{\"moe\":{\"weight\":\"int4_per_group_32\"}}}" +GPU_UTIL="${GPU_UTIL:-0.88}" +# --- RDMA fabric (OVERRIDABLE) ----------------------------------------------- +# Defaults are validated Broadcom Thor2 (bnxt RoCE) values: +# ibv device names rdma0..rdma7, host NIC eno0, GID index 3. On a DIFFERENT +# fabric (e.g. Mellanox mlx5) override these, e.g.: +# NCCL_IB_HCA=mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9 \ +# RDMA_DEVICES=mlx5_0,mlx5_2,... SOCKET_IFNAME=eth0 IB_GID_INDEX=3 THOR2_BNXT_FIX=0 +SOCKET_IFNAME="${SOCKET_IFNAME:-eno0}" +NCCL_IB_HCA_VAL="${NCCL_IB_HCA:-rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7}" +RDMA_DEVICES="${RDMA_DEVICES:-rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7}" +IB_GID_INDEX="${IB_GID_INDEX:-3}" +# Thor2 bnxt libibverbs ABI fix (host v34 driver vs image v59). Needed on Thor2; +# harmless elsewhere but set THOR2_BNXT_FIX=0 to skip the two -v mounts. +THOR2_BNXT_FIX="${THOR2_BNXT_FIX:-1}" +THOR2_LIBIBVERBS_HOST="${THOR2_LIBIBVERBS_HOST:-/usr/lib/x86_64-linux-gnu/libibverbs.so.1.14.39.0}" +THOR2_LIBIBVERBS_IMG="${THOR2_LIBIBVERBS_IMG:-/usr/lib/x86_64-linux-gnu/libibverbs.so.1.16.62.0}" +THOR2_BNXT_HOST="${THOR2_BNXT_HOST:-/usr/local/lib/libbnxt_re-rdmav34.so}" +THOR2_BNXT_IMG="${THOR2_BNXT_IMG:-/usr/lib/x86_64-linux-gnu/libibverbs/libbnxt_re-rdmav34.so}" +if [ "${THOR2_BNXT_FIX}" = "1" ]; then + BNXT_MOUNTS="-v ${THOR2_LIBIBVERBS_HOST}:${THOR2_LIBIBVERBS_IMG}:ro -v ${THOR2_BNXT_HOST}:${THOR2_BNXT_IMG}:ro" +else + BNXT_MOUNTS="" +fi +# ----------------------------------------------------------------------------- +# Skip the boot memory-profiling forward (profile_run) entirely by pinning KV +# cache size — mirrors DeepSeek #181 (which used it for an AITER fp8 profiling +# bug). profile_run's dummy forward hangs under tp8xDP2 + mori all2all (sampler +# gather / _sync_device all2all deadlocks); pinning KV bytes bypasses all of it. +# At TP2/DP8 weights are ~137.5 GiB/GPU + 16 GiB MoRI heap, so KV room is tighter +# (~10 GiB) than at TP8 (was 14e9). Pin 8 GiB — ample for K3's tiny MLA KV +# (kv_lora_rank=512, only 24/93 full-attn layers; ~13.5 KiB/tok fp8 => ~600k tok). +KV_CACHE_MEMORY_BYTES="${KV_CACHE_MEMORY_BYTES:-8000000000}" +# Wide-EP shape per pool: TP2 x DP8 -> EP16, no PP (per user: "No PP — disagg via +# MoRIIO, DP/EP via MoRI-EP"). 16 GPUs/pool (2 nodes x 8). world=16 => EP16 (896/16 +# =56 experts/GPU via MoRI-EP all2all). DP8 = 8 independent attention streams (the +# 69/93 KDA recurrent layers stay DP-local => best decode throughput); TP2 shards +# the 106.5 GiB replicated attn+shared-expert weight to 53.3 GiB/GPU so it FITS +# (pure TP1/DP16 = 190.7 GiB weights > 192 GiB HBM; see Confluence 1830010189). +# Per-GPU: experts 84.2 + repl 53.3 = 137.5 GiB weights + 16 GiB MoRI heap. DP_SIZE +# is the GLOBAL dp count (=world/TP=16/2=8); DP_LOCAL=4 dp ranks/node (8 GPUs/2). +TP_SIZE="${TP_SIZE:-2}" +DP_SIZE="${DP_SIZE:-8}" +DP_LOCAL="${DP_LOCAL:-4}" +SERVE_PORT=20005; RPC_PORT=13345 +# JIT cache persistence (best-practice): host dir keyed per image AND per role +# CLASS. Prefill (mori_high_throughput + cudagraph NONE) and decode +# (mori_low_latency + PIECEWISE) compile DIFFERENT kernel variants under the SAME +# aiter .so filenames (module_moe_asm/moe_sorting_opus/...). Node-local nvme keeps +# prefill (025/043) and decode (047/048) caches physically separate; the per-CLASS +# subdir also prevents collision if a node ever hosts both roles. Do NOT put this +# on a SHARED FS (would let prefill/decode clobber each other's same-named .so). +ROLE_CLASS="${ROLE%%_*}" # prefill | decode +JIT_HOST="${JIT_HOST:-/tmp/$USER/vllm_jit_cache/k3disagg_${ROLE_CLASS}}"; mkdir -p "$JIT_HOST" + +VLLM_SP=/usr/local/lib/python3.12/dist-packages/vllm + +# Per-pool topology: TP2 x DP8 -> EP16 (see the TP_SIZE/DP_SIZE block above). +# Master node hosts DP ranks 0..DP_LOCAL-1; worker node hosts DP_LOCAL..2*DP_LOCAL-1 +# (=> --data-parallel-start-rank ${DP_LOCAL}). gpu-util 0.88 for the 16 GiB MoRI +# shmem heap reserved before the vLLM snapshot. +case "$ROLE" in + prefill_master) DP_ADDR=$PMASTER; KV_ROLE=kv_producer; BACKEND=mori_high_throughput; CG=NONE; HEADLESS=""; START="" ;; + prefill_worker) DP_ADDR=$PMASTER; KV_ROLE=kv_producer; BACKEND=mori_high_throughput; CG=NONE; HEADLESS="--headless"; START="--data-parallel-start-rank ${DP_LOCAL}" ;; + decode_master) DP_ADDR=$DMASTER; KV_ROLE=kv_consumer; BACKEND=mori_low_latency; CG=PIECEWISE; HEADLESS=""; START="" ;; + decode_worker) DP_ADDR=$DMASTER; KV_ROLE=kv_consumer; BACKEND=mori_low_latency; CG=PIECEWISE; HEADLESS="--headless"; START="--data-parallel-start-rank ${DP_LOCAL}" ;; + *) echo "bad ROLE=$ROLE"; exit 1 ;; +esac +# Optional cudagraph override (DECODE_CG=NONE bypasses PIECEWISE capture to isolate +# a capture-time GPU fault; PREFILL_CG likewise). Decode capture of the KDA conv +# path can fault on some builds; NONE trades decode-graph perf for stability. +case "$ROLE" in + prefill_*) [ -n "${PREFILL_CG:-}" ] && CG="$PREFILL_CG" ;; + decode_*) [ -n "${DECODE_CG:-}" ] && CG="$DECODE_CG" ;; +esac +# Optional all2all-backend override (test: mori_low_latency on both pools to rule +# out mori_high_throughput/InterNodeV1 as the profile-forward all2all deadlock). +if [[ "$ROLE" == prefill_* && -n "${PREFILL_BACKEND:-}" ]]; then BACKEND="$PREFILL_BACKEND"; fi +if [[ "$ROLE" == decode_* && -n "${DECODE_BACKEND:-}" ]]; then BACKEND="$DECODE_BACKEND"; fi + +IS_MASTER=0; [[ "$ROLE" == *_master ]] && IS_MASTER=1 +CONTAINER="k3disagg_${ROLE}" +docker rm -f "$CONTAINER" 2>/dev/null || true + +# kv-transfer-config on ALL ranks (masters AND headless workers). A headless +# worker hosts real DP ranks (e.g. decode rank1 on the worker node); WITHOUT +# --kv-transfer-config its engine never instantiates the MoRIIO connector, so it +# binds NO handshake listener -> prefill can't transfer KV to those ranks -> +# "Timed out waiting for write_ready_flags" -> EngineDead. (Was gated on +# IS_MASTER, which silently made every worker-node DP rank a transfer black hole.) +# To survive host-bash -> ssh -> docker -c quoting, pass JSON as base64 via env. +KVCFG_B64="" +if true; then + # Peer pool's per-DP-pod node IPs (ordered by pod index = global_dp_rank//dp_local). + # A prefill (kv_producer) handshakes the DECODE pool -> needs decode hosts; a decode + # (kv_consumer) notifies the PREFILL pool -> needs prefill hosts. Without this the + # connector falls back to a single peer host (the master), so KV writes/notifies to + # ranks on the peer's WORKER node silently miss -> that node's decode ranks generate + # context-free (the 50%/DP2, ~88%/DP8 wrong-answer alternation). Consumed by the + # apply_kimik3_moriio_pod_hosts patcher as multi_pod_hosts in the handshake. + if [[ "$ROLE" == prefill_* ]]; then PEER_POD_HOSTS="${DECODE_POD_HOSTS:-}"; else PEER_POD_HOSTS="${PREFILL_POD_HOSTS:-}"; fi + # MORIIO_READ_MODE=1 selects the connector's READ path (decode pulls KV, sync, + # returns N-1) instead of WRITE (prefill pushes, async, returns N). READ is the + # more-tested vLLM disagg path; toggled for A/B against the WRITE decode-consume bug. + READ_MODE_JSON=""; if [ "${MORIIO_READ_MODE:-0}" = "1" ]; then READ_MODE_JSON=",\"read_mode\":\"true\""; fi + KVCFG_JSON="{\"kv_connector\":\"MoRIIOConnector\",\"kv_role\":\"${KV_ROLE}\",\"kv_port\":\"9711\",\"kv_connector_extra_config\":{\"proxy_ip\":\"${PROXY_IP}\",\"proxy_port\":\"30000\",\"proxy_ping_port\":\"36367\",\"http_port\":\"${SERVE_PORT}\",\"local_ping_port\":\"61555\",\"handshake_port\":\"8405\",\"notify_port\":\"61005\",\"moriio_pod_hosts\":\"${PEER_POD_HOSTS}\",\"post_batch_size\":${MORIIO_POST_BATCH_SIZE:--1},\"qp_per_transfer\":${MORIIO_QP_PER_TRANSFER:-1},\"num_workers\":${MORIIO_NUM_WORKERS:-1}${READ_MODE_JSON}}}" + KVCFG_B64=$(printf '%s' "$KVCFG_JSON" | base64 -w0) +fi +# api-server ONLY on masters (headless workers must NOT bind an api-server). +# api-server-count MUST be <= data-parallel-size: the frontend DP load-balancer +# round-robins requests across data_parallel_rank [0, api_server_count). Match it +# to DP_SIZE so every request lands on a live DP rank. +if [[ $IS_MASTER -eq 1 ]]; then + APISERVERS="--api-server-count ${DP_SIZE} --port ${SERVE_PORT}" +else + APISERVERS="" +fi + +echo "[disagg] node=$(hostname -s) role=$ROLE dp_addr=$DP_ADDR kv_role=${KV_ROLE:-none} backend=$BACKEND" + +docker run -d --name "$CONTAINER" \ + --network host --ipc host \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband --group-add video \ + --cap-add SYS_PTRACE --cap-add IPC_LOCK --security-opt seccomp=unconfined \ + --shm-size 128g --ulimit memlock=-1 --ulimit stack=67108864 --ulimit nofile=524288:524288 \ + -e VLLM_ROCM_USE_AITER_MLA=0 \ + -e AITER_SITUV2_A8W4=1 \ + -e VLLM_ROCM_USE_AITER=1 -e VLLM_ROCM_USE_AITER_MOE=1 \ + -e VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 \ + -e VLLM_USE_AITER_TRITON_SILU_MUL=0 -e VLLM_ROCM_USE_AITER_RMSNORM=1 \ + -e VLLM_ENGINE_READY_TIMEOUT_S=3600 \ + -e VLLM_SSM_CONV_STATE_LAYOUT=DS \ + -e NCCL_SOCKET_IFNAME=${SOCKET_IFNAME} -e GLOO_SOCKET_IFNAME=${SOCKET_IFNAME} \ + -e NCCL_IB_DISABLE=0 -e NCCL_IB_HCA=${NCCL_IB_HCA_VAL} -e MORI_RDMA_DEVICES=${RDMA_DEVICES} -e MORI_SOCKET_IFNAME=${SOCKET_IFNAME} \ + -e NCCL_IB_GID_INDEX=${IB_GID_INDEX} -e NCCL_IGNORE_CPU_AFFINITY=1 \ + -e HSA_ENABLE_IPC_MODE_LEGACY=0 -e HSA_NO_SCRATCH_RECLAIM=1 \ + -e PYTORCH_ALLOC_CONF=expandable_segments:False -e PYTORCH_HIP_ALLOC_CONF=expandable_segments:False \ + -e MORIIO_SKIP_MAMBA="${MORIIO_SKIP_MAMBA:-0}" \ + -e K3_MAMBA_N1_FORCE="${K3_MAMBA_N1_FORCE:-}" \ + -e K3_WRITE_FENCE="${K3_WRITE_FENCE:-}" \ + -e VLLM_BATCH_INVARIANT="${VLLM_BATCH_INVARIANT:-0}" \ + -e K3_MLA_SINGLE_SPLIT="${K3_MLA_SINGLE_SPLIT:-1}" \ + -e K3_GROUP_ROUTING="${K3_GROUP_ROUTING:-1}" \ + -e K3_EXTRA_FIXES="${K3_EXTRA_FIXES:-0}" \ + -e K3_CHUNK_GATE_SLACK="${K3_CHUNK_GATE_SLACK:-2}" \ + -e K3_CHUNK_GATE_DEBUG="${K3_CHUNK_GATE_DEBUG:-0}" \ + -e K3_XFER_PROBE="${K3_XFER_PROBE:-0}" \ + -e K3_MLA_FULL_PREFILL="${K3_MLA_FULL_PREFILL:-1}" \ + -e K3_FORCE_PREFILL_KDA="${K3_FORCE_PREFILL_KDA:-0}" \ + -e K3_WRITE_FENCE_MS="${K3_WRITE_FENCE_MS:-20}" \ + -e K3_WRITE_DEVSYNC="${K3_WRITE_DEVSYNC:-0}" \ + -e K3_KDA_CONV_DEBUG="${K3_KDA_CONV_DEBUG:-0}" \ + -e K3_FWD_BREADCRUMB="${K3_FWD_BREADCRUMB:-0}" \ + -e K3_WRITE_BC="${K3_WRITE_BC:-0}" \ + -e K3_KDA_STATE_PROBE="${K3_KDA_STATE_PROBE:-0}" \ + -e K3_MAMBA_BC="${K3_MAMBA_BC:-0}" \ + -e K3_DECODE_RECV_PROBE="${K3_DECODE_RECV_PROBE:-0}" \ + -e K3_HS_BC="${K3_HS_BC:-0}" \ + -e K3_INPUTS_PROBE="${K3_INPUTS_PROBE:-0}" \ + -e AMD_SERIALIZE_KERNEL="${AMD_SERIALIZE_KERNEL:-0}" -e AMD_LOG_LEVEL="${AMD_LOG_LEVEL:-0}" \ + -e MORI_GPU_ARCHS=gfx942 -e MORI_IB_GID_INDEX=${IB_GID_INDEX} -e MORI_IB_ENABLE_RELAXED_ORDERING=1 \ + -e MORI_NUM_QP_PER_PE=8 -e MORI_SHMEM_HEAP_SIZE=17179869184 \ + -e MORI_RDMA_TC=41 -e MORI_RDMA_SL=0 -e MORI_IO_SL=1 \ + -e VLLM_MORIIO_QP_PER_TRANSFER="${VLLM_MORIIO_QP_PER_TRANSFER:-2}" -e VLLM_MORIIO_NUM_WORKERS="${VLLM_MORIIO_NUM_WORKERS:-4}" \ + -e AITER_JIT_DIR=/opt/vllm_cache/aiter -e TRITON_CACHE_DIR=/opt/vllm_cache/triton \ + -e VLLM_CACHE_ROOT=/opt/vllm_cache/vllm \ + -e KVCFG_B64="$KVCFG_B64" \ + -e QUANT_CONFIG="$QUANT_CONFIG" \ + -v "$MODEL_DIR":/model:ro -v "$LOGHOST":/logs -v "$PATCHER_DIR":/patchers:ro \ + -v "$JIT_HOST":/opt/vllm_cache \ + ${BNXT_MOUNTS} \ + --entrypoint bash \ + "$IMAGE" -c " + set -e + mkdir -p /opt/vllm_cache/aiter /opt/vllm_cache/triton /opt/vllm_cache/vllm + # The NFS/DockerHub image tar is a STALE build shipping flydsl 0.2.2, but the + # K3 int4 SiTUv2 MoE path (_setup_kernel_k3_situ_gfx942 -> compile_moe_gemm1) + # hard-requires flydsl>=0.2.4 (ImportError otherwise -> WorkerProc init fails -> + # pool never starts). Dockerfile.kimik3_disagg installs 0.2.4 (line 148) but this + # tar predates that. Bump at container start (once per node, pure-python, fast). + FLYDSL_VER=\$(python3 -c 'import flydsl,importlib.metadata as m; print(m.version(\"flydsl\"))' 2>/dev/null || echo 0) + if [ \"\$FLYDSL_VER\" != \"0.2.4\" ]; then + echo \"[disagg] flydsl \$FLYDSL_VER -> upgrading to 0.2.4 (K3 int4 SiTUv2 requires >=0.2.4)\" + pip install --no-cache-dir 'flydsl==0.2.4' 2>&1 | tail -1 + fi + echo '[disagg] applying KDA MoRIIO patchers...' + # KDA/HMA/sampler patchers baked into image source (kimik3-wideep-disagg). + # Runtime-relax the over-strict aiter#4471 packed-int4 guard (the grafted + # K3-aware AITER lacks compile_moe_gemm1(act=); the proven image runs the + # same aiter fine -> guard is a false positive on this stack). + VLLM_SP=\$(python3 -c 'import vllm,os;print(os.path.dirname(vllm.__file__))' 2>/dev/null) + if [ -f /patchers/apply_kimik3_aiter_situv2_int4.py ]; then + python3 /patchers/apply_kimik3_aiter_situv2_int4.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_situ_aiter_gfx942.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_situ_aiter_gfx942.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_mxfp4_int4_guard_relax.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_mxfp4_int4_guard_relax.py \"\$VLLM_SP\" || true + fi + # Reconcile MoRIIO self.block_size from an attention (non-mamba) layer so the + # MLA block_size (1536) doesn't trip the guard against KDA's block_size=1. + if [ -f /patchers/apply_kimik3_moriio_block_size_fix.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_block_size_fix.py \"\$VLLM_SP\" || true + fi + # Use the PADDED physical mamba page for KDA block stride/geometry (fixes the + # producer-side GPU memory fault: unpadded conv+ssm page drifts RDMA offsets OOB). + if [ -f /patchers/apply_kimik3_moriio_mamba_page_pad_fix.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_mamba_page_pad_fix.py \"\$VLLM_SP\" || true + fi + # Thread real tp_size into get_port_offset so per-rank ports don't collide when + # DP-local>1 AND TP>1 (wide-EP TP2xDP-local-4: dp0/tp1 and dp1/tp0 both bound + # handshake port 8406 -> ZMQError Address already in use -> listener dies -> pool + # hangs on 'No available shared memory broadcast block'). offset = dp*tp_size+tp. + if [ -f /patchers/apply_kimik3_moriio_port_offset_tpsize.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_port_offset_tpsize.py \"\$VLLM_SP\" || true + fi + # Multi-NODE disagg: advertise the peer pool's per-pod node IPs so prefill can + # reach decode ranks on the HEADLESS worker node (else KV writes to that node's + # ranks miss -> context-free generation -> 50%/DP2 88%/DP8 wrong-answer alternation). + if [ -f /patchers/apply_kimik3_moriio_pod_hosts.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_pod_hosts.py \"\$VLLM_SP\" || true + fi + # ROOT-CAUSE FIX: route KDA/mamba state transfer by the MAMBA KV-cache group's + # block ids (group [1]), not attention's (group [0]). Without this, decode reads + # zero KDA state -> fluent but context-free output. Load-bearing; always on. + if [ -f /patchers/apply_kimik3_moriio_mamba_blockids.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_mamba_blockids.py \"\$VLLM_SP\" || true + fi + # ROOT-CAUSE FIX: remote_tp_size=1 (un-advertising router) collapses ALL prefill + # ranks to decode tp0 -> only 1/8 decode shards get KV -> context-free. Normalize + # degenerate remote TP to local world_size (symmetric TP). Load-bearing; always on. + if [ -f /patchers/apply_kimik3_moriio_remote_tp_fix.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_remote_tp_fix.py \"\$VLLM_SP\" || true + fi + # ROOT-CAUSE FIX: mamba/KDA N-vs-N-1 boundary. Prefill computes h(N-1) (drop last + # prompt token), decode recomputes token N from h(N-1). Without this, decode + # double-counts the last token in the recurrent state -> echoes it -> wrong output. + # Ports vLLM's own nixl/mooncake hybrid-PD handling. Load-bearing; always on. + if [ -f /patchers/apply_kimik3_moriio_mamba_n1.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_mamba_n1.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_moriio_group_routing.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_group_routing.py \"\$VLLM_SP\" || true + fi + if [ \"\${K3_EXTRA_FIXES:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_chunked_allgrp.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_chunked_allgrp.py \"\$VLLM_SP\" || true + fi + if [ \"\${K3_EXTRA_FIXES:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_chunk_gate_fix.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_chunk_gate_fix.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_xfer_probe.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_xfer_probe.py \"\$VLLM_SP\" || true + fi + # k3-mla-boundary: clamp final MLA block RDMA copy to valid slots (fix decode recall) + if [ \"\${K3_ENABLE_CLAMP:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_moriio_mla_boundary_clamp.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_mla_boundary_clamp.py \"\$VLLM_SP\" || true + fi + # k3-mla-full: the real fix (prefill computes N MLA / KDA stays N-1). Clamp above OFF by default. + if [ -f /patchers/apply_kimik3_mla_full_prefill.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_mla_full_prefill.py \"\$VLLM_SP\" || true + fi + # k3-force-prefill-kda: route disagg boundary token through prefill KDA kernel + if [ -f /patchers/apply_kimik3_force_prefill_kda.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_force_prefill_kda.py \"\$VLLM_SP\" || true + fi + # Force single-split TRITON_MLA decode (deterministic; avoids uninit-tail merge) + if [ -f /patchers/apply_kimik3_mla_single_split.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_mla_single_split.py \"\$VLLM_SP\" || true + fi + # RDMA write-then-notify ordering fence (K3_WRITE_FENCE=delay): settle before + # write_done so decode does not read stale HBM (non-deterministic recall). + if [ -f /patchers/apply_kimik3_moriio_write_fence.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_write_fence.py \"\$VLLM_SP\" || true + fi + # Diagnostic breadcrumbs for the WRITE KV delivery (K3_WRITE_BC=1). + if [ \"\${K3_WRITE_BC:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_moriio_write_bc.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_write_bc.py \"\$VLLM_SP\" || true + fi + # Ground-truth KDA state probe (K3_KDA_STATE_PROBE=1): norm of recurrent/conv + # state at the slot decode reads -- ~0 means transferred state didn't land. + if [ \"\${K3_KDA_STATE_PROBE:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_kda_state_probe.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_state_probe.py \"\$VLLM_SP\" || true + fi + # Decode-side receive probe (K3_DECODE_RECV_PROBE=1): on write completion, + # read decode's OWN KV slot norm -- ~0 means RDMA bytes never landed on decode. + if [ \"\${K3_DECODE_RECV_PROBE:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_decode_recv_probe.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_decode_recv_probe.py \"\$VLLM_SP\" || true + fi + # Handshake dial breadcrumb (K3_HS_BC=1): logs self_tp/dial_tp/port/path so we + # can see if all prefill ranks wrongly dial the same decode rank. + if [ \"\${K3_HS_BC:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_handshake_dial_bc.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_handshake_dial_bc.py \"\$VLLM_SP\" || true + fi + # Decode inputs probe (K3_INPUTS_PROBE=1): logs num_computed_tokens/positions/ + # block_table in _prepare_inputs -- to find the WRITE-mode decode-consume bug. + if [ \"\${K3_INPUTS_PROBE:-0}\" = \"1\" ] && [ -f /patchers/apply_kimik3_decode_inputs_probe.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_decode_inputs_probe.py \"\$VLLM_SP\" || true + fi + # Diagnostics: MORIIO_SKIP_MAMBA=1 isolates the KDA transfer; KDA OOB bounds log. + if [ -f /patchers/apply_kimik3_moriio_mamba_diag.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_mamba_diag.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_moriio_save_skip_mamba.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_moriio_save_skip_mamba.py \"\$VLLM_SP\" || true + fi + # Guard gather_initial_states against OOB KDA state idx (the disagg producer + # prefill GPU memory fault: mis-flagged has_initial_state -> reads a stale/OOB block). + if [ -f /patchers/apply_kimik3_kda_gather_guard.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_gather_guard.py \"\$VLLM_SP\" || true + fi + # k3-kda-nosync: make the KDA gather OOB guard sync-free (unblocks + speeds up + # long context >500K; removes a per-KDA-layer device->CPU sync). No-op on a v3+ + # image where it is already baked in. + if [ -f /patchers/apply_kimik3_kda_gather_nosync.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_gather_nosync.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_kda_conv_debug.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_conv_debug.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_fwd_breadcrumb.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_fwd_breadcrumb.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_layer_breadcrumb.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_layer_breadcrumb.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_kda_fa_contiguous.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_fa_contiguous.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_kda_internal_bc.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kda_internal_bc.py \"\$VLLM_SP\" || true + fi + if [ -f /patchers/apply_kimik3_kvzero_bounds.py ] && [ -n \"\$VLLM_SP\" ]; then + python3 /patchers/apply_kimik3_kvzero_bounds.py \"\$VLLM_SP\" || true + fi + QUANTARG=() + if [ -n \"\$QUANT_CONFIG\" ]; then + QUANTARG=(--quantization-config \"\$QUANT_CONFIG\") + fi + KVARG=() + if [ -n \"\$KVCFG_B64\" ]; then + KVJSON=\$(printf '%s' \"\$KVCFG_B64\" | base64 -d) + KVARG=(--kv-transfer-config \"\$KVJSON\") + fi + echo '[disagg] launching vllm serve...' + vllm serve /model --served-model-name kimi-k3 --tensor-parallel-size ${TP_SIZE} \ + --data-parallel-size ${DP_SIZE} --data-parallel-size-local ${DP_LOCAL} \ + --data-parallel-address ${DP_ADDR} --data-parallel-rpc-port ${RPC_PORT} ${START} ${HEADLESS} \ + --enable-expert-parallel --all2all-backend ${BACKEND} \ + --trust-remote-code --reasoning-parser kimi_k3 --mm-encoder-tp-mode data --safetensors-load-strategy ${LOAD_STRATEGY:-prefetch} \ + --no-enable-prefix-caching --kv-cache-dtype ${KV_CACHE_DTYPE:-fp8} --block-size ${BLOCK_SIZE:-16} \ + --kv-cache-memory-bytes ${KV_CACHE_MEMORY_BYTES} \ + --max-model-len ${MAX_MODEL_LEN} --max-num-seqs 8 --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS} --gpu-memory-utilization ${GPU_UTIL} \ + --distributed-timeout-seconds 7200 \ + --compilation-config '{\"cudagraph_mode\":\"${CG}\",\"custom_ops\":[\"+quant_fp8\"]}' \ + ${APISERVERS} \"\${QUANTARG[@]}\" \"\${KVARG[@]}\" 2>&1 | tee /logs/vllm_${ROLE}.log + " +echo "[disagg] $ROLE started -> $LOGHOST/vllm_${ROLE}.log" diff --git a/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d_launch.sh b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d_launch.sh new file mode 100755 index 00000000..877fd3fa --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/run_2p2d_launch.sh @@ -0,0 +1,73 @@ +#!/bin/bash +# Orchestrate the K3 2P/2D disagg bring-up from the head/control host. +# Prefill pool: PM (master) + PW (worker). Decode pool: DM (master) + DW (worker). +# Router runs on PM (proxy_ip). Launch order: workers first, then masters, then router. +set -euo pipefail + +# --- node eth0 IPs (EDIT for your 4-node allocation) --- +# Prefill pool = 2 nodes (master+worker); decode pool = 2 nodes (master+worker). +# PM_NODE/DM_NODE are the ssh hostnames; *_IP are the eth0 IPs the peers dial. +PM_NODE=${PM_NODE:-}; PM_IP=${PM_IP:-} # prefill master + proxy/router +PW_NODE=${PW_NODE:-}; PW_IP=${PW_IP:-} # prefill worker +DM_NODE=${DM_NODE:-}; DM_IP=${DM_IP:-} # decode master +DW_NODE=${DW_NODE:-}; DW_IP=${DW_IP:-} # decode worker + +# Path to THIS recipe folder on the control host (scp'd to each node). Override REPO=. +REPO=${REPO:-$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)} +# Kimi-K3 MXFP4 weights path (must exist on every node; local nvme recommended). +MODEL_DIR=${MODEL_DIR:?set MODEL_DIR to your Kimi-K3-MXFP4 path} +SSH="ssh -o StrictHostKeyChecking=no" +# Load-bearing env only. TP2xDP8 -> EP16 per pool (no PP). See README for knobs. +COMMON="IMAGE=${IMAGE:-kimik3-wideep-disagg:latest} TP_SIZE=${TP_SIZE:-2} DP_SIZE=${DP_SIZE:-8} DP_LOCAL=${DP_LOCAL:-4} KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-8000000000} PREFILL_BACKEND=mori_low_latency DECODE_CG=${DECODE_CG:-NONE} MODEL_DIR=${MODEL_DIR} PMASTER=$PM_IP DMASTER=$DM_IP PROXY_IP=$PM_IP DECODE_POD_HOSTS=$DM_IP,$DW_IP PREFILL_POD_HOSTS=$PM_IP,$PW_IP" + +deploy() { # $1=node + $SSH "$1" 'mkdir -p ~/k3disagg/patchers ~/k3disagg/logs' 2>/dev/null + scp -o StrictHostKeyChecking=no "$REPO/run_2p2d.sh" "$REPO/load_image.sh" "$1:~/k3disagg/" >/dev/null + scp -o StrictHostKeyChecking=no "$REPO/patchers/"*.py "$1:~/k3disagg/patchers/" >/dev/null + # Self-restoring: ensure the disagg image is present (pull from DockerHub if not). + $SSH "$1" "cd ~/k3disagg && TAG='${IMAGE:-kimik3-wideep-disagg:latest}' HUB_IMAGE='${HUB_IMAGE:-}' DOCKER_USER='${DOCKER_USER:-}' DOCKER_PAT='${DOCKER_PAT:-}' bash load_image.sh" 2>&1 | tail -1 +} + +echo "=== deploy scripts+patchers+image to 4 nodes ===" +for n in $PM_NODE $PW_NODE $DM_NODE $DW_NODE; do deploy "$n"; echo " $n ok"; done + +echo "=== start WORKERS first ===" +$SSH $PW_NODE "cd ~/k3disagg && $COMMON ROLE=prefill_worker bash run_2p2d.sh" 2>&1 | tail -1 +$SSH $DW_NODE "cd ~/k3disagg && $COMMON ROLE=decode_worker bash run_2p2d.sh" 2>&1 | tail -1 +sleep 5 +echo "=== start MASTERS ===" +$SSH $PM_NODE "cd ~/k3disagg && $COMMON ROLE=prefill_master bash run_2p2d.sh" 2>&1 | tail -1 +$SSH $DM_NODE "cd ~/k3disagg && $COMMON ROLE=decode_master bash run_2p2d.sh" 2>&1 | tail -1 + +# Auto-start the router in the discovery window so the engines' _ping threads +# connect before exhausting MAX_PING_RETRIES (a late router misses discovery -> +# "0 prefill 0 decode"). AUTO_ROUTER=1 waits for both masters' /v1/models, then +# launches the router on PM. +if [ "${AUTO_ROUTER:-0}" = "1" ]; then + echo "=== AUTO_ROUTER: waiting for both masters, then starting router ===" + ( for t in $(seq 1 60); do + pm=$($SSH $PM_NODE "curl -s -m 5 http://$PM_IP:20005/v1/models 2>/dev/null | grep -c kimi-k3" 2>/dev/null) + dm=$($SSH $DM_NODE "curl -s -m 5 http://$DM_IP:20005/v1/models 2>/dev/null | grep -c kimi-k3" 2>/dev/null) + if [ "${pm:-0}" -ge 1 ] && [ "${dm:-0}" -ge 1 ]; then + echo "[auto-router] both masters ready; starting router" + $SSH $PM_NODE "docker exec k3disagg_prefill_master bash -c 'setsid nohup vllm-router --host 0.0.0.0 --port 30000 --vllm-pd-disaggregation --kv-connector moriio --prefill http://$PM_IP:20005 --decode http://$DM_IP:20005 --vllm-discovery-address 0.0.0.0:36367 --intra-node-data-parallel-size ${DP_LOCAL:-4} --moriio-dp-size ${DP_SIZE:-8} --policy round_robin --prefill-policy round_robin --decode-policy round_robin --log-level info > /logs/router.log 2>&1 < /dev/null &'" + break + fi + sleep 20 + done ) & + echo "[auto-router] watcher started (pid $!)" +fi + +echo "" +echo "=== bring-up started. Watch for 'Application startup complete' in: ===" +echo " PM: $SSH $PM_NODE 'docker logs -f k3disagg_prefill_master'" +echo " DM: $SSH $DM_NODE 'docker logs -f k3disagg_decode_master'" +echo "" +echo "=== once BOTH masters are up, start the router on $PM_NODE: ===" +cat < /logs/router.log 2>&1"' +EOF +echo "" +echo "=== test (after router shows 'Add Prefill'+'Add Decode'): ===" +echo " curl http://$PM_IP:30000/v1/models" +echo " python3 niah_probe.py --url http://$PM_IP:30000 --model kimi-k3 --ctx 6000 --depths 0.1,0.5,0.9" diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/README.md b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/README.md new file mode 100644 index 00000000..75f69421 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/README.md @@ -0,0 +1,48 @@ +# Kimi-K3 (MXFP4) on MI300X / gfx942 — Wide-EP (generic all2all) + a8w4 + +Serves Kimi-K3 (MXFP4) across **2 nodes** (16× MI300X) with **expert parallelism**: +PP2×TP8 for weight fit (~102 GB/GPU) plus `--enable-expert-parallel`, so the 896 experts are split 8-way across each node's 8 GPUs (112/GPU), replicated per PP stage. Expert all-to-all uses the generic +`--all2all-backend allgather_reducescatter` (see `../wideep_int4_moriep` for the true +MoRI-EP kernels). `AITER_SITUV2_A8W4=1` selects the a8w4 (fp8-activation × int4-weight) +SiTU MoE path. Colocated (single instance; no prefill/decode disaggregation). + +- Image: `amdsiloai/vllm:kimi-k3-mi325x-release-v2` (gfx942 K3 vLLM build; public, anonymous pull) +- Model: `moonshotai/Kimi-K3` (MXFP4), on local NVMe or NFS + +## Run (worker FIRST, then head) + +```bash +# on the worker node (rank 1): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=worker MASTER= bash run.sh + +# on the head node (rank 0, serves the API on :8000): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=head MASTER= bash run.sh +``` + +- `MASTER` = the head node's **eth0** IP. `MODEL_DIR` = Kimi-K3-MXFP4 weights (prefer local NVMe). +- First start compiles the AITER EP MoE asm + a8w4 kernels (a few minutes), then cached. + +## Verify + +```bash +curl http://:8000/v1/models +python3 niah_probe.py --url http://:8000 --model kimi-k3 --ctx 8500 --depths 0.1,0.5,0.9 +``` +The response `system_fingerprint` contains `-ep-` when expert parallelism is active. + +## Toggles + +- Drop a8w4: remove `-e AITER_SITUV2_A8W4=1` from `run.sh` → default a16w4 MoE path. +- MoE dispatch: add `-e VLLM_ROCM_AITER_MOE_DISPATCH_POLICY=2` (multi-pass; may help MoE-heavy at higher concurrency). + +## Key env (set in run.sh) + +| Var | Value | Why | +|-----|-------|-----| +| `VLLM_ROCM_USE_AITER_MLA` | `0` | Required on gfx942 (AITER MLA is gfx950-only). | +| `AITER_SITUV2_A8W4` | `1` | Route K3 SiTU MXFP4 MoE through the a8w4 interleaved flydsl kernels. | +| `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` | 8× mlx5 / `3` | RDMA fabric (**cluster-specific — override for yours**). | +| `HSA_ENABLE_IPC_MODE_LEGACY` / `PYTORCH_(HIP_)ALLOC_CONF` | `0` / `expandable_segments:False` | ROCm 7.2.x requirements. | + +Overridable env: `IMAGE`, `MODEL_DIR`, `MASTER`, `PORT` (8000), `MAX_MODEL_LEN` (10240), +`MAX_NUM_SEQS` (8), `GPU_UTIL` (0.90). diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/niah_probe.py b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/niah_probe.py new file mode 100644 index 00000000..0fed8766 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/niah_probe.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Lightweight NIAH (needle-in-a-haystack) + basic accuracy probe for a served +Kimi-K3 endpoint. No deps beyond stdlib (urllib). Run from any node that can +reach the head IP:port. + +Usage: + python3 niah_probe.py --url http://:8000 --model kimi-k3 \ + --ctx 8000 --depths 0.1,0.5,0.9 +""" +import argparse, json, time, urllib.request, urllib.error + +FILLER = ("The quick brown fox jumps over the lazy dog. " + "Grand vistas of rolling hills stretched to the horizon. ") + +def chat(url, model, prompt, max_tokens=256, timeout=180): + body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, "temperature": 0, + }).encode() + req = urllib.request.Request(url.rstrip("/") + "/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + dt = time.time() - t0 + ch = d["choices"][0]["message"] + txt = (ch.get("content") or "") + " " + (ch.get("reasoning") or "") + return txt.strip(), d.get("usage", {}), dt + +def make_haystack(approx_words, needle, depth): + # Build filler to approx word count; insert needle at fractional depth. + words = [] + while len(words) < approx_words: + words.extend(FILLER.split()) + n = len(words) + pos = int(n * depth) + needle_words = needle.split() + words[pos:pos] = needle_words + return " ".join(words) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", required=True) + ap.add_argument("--model", default="kimi-k3") + ap.add_argument("--ctx", type=int, default=8000, help="approx haystack words") + ap.add_argument("--depths", default="0.1,0.5,0.9") + args = ap.parse_args() + + print(f"== sanity ==") + txt, usage, dt = chat(args.url, args.model, "Reply with only: OK", max_tokens=8) + print(f" resp={txt!r} usage={usage} {dt:.1f}s") + + secret = "The secret passcode is HELIOTROPE-7492." + question = ("\n\nQuestion: What is the secret passcode? " + "Answer with only the passcode.") + passed = 0; total = 0 + for depth in [float(x) for x in args.depths.split(",")]: + hay = make_haystack(args.ctx, secret, depth) + prompt = hay + question + try: + txt, usage, dt = chat(args.url, args.model, prompt, max_tokens=256) + except Exception as e: + print(f" depth={depth}: ERROR {e}") + total += 1; continue + ok = "HELIOTROPE-7492" in txt.upper() + passed += ok; total += 1 + print(f" depth={depth:>4}: {'PASS' if ok else 'FAIL'} " + f"(prompt_tok={usage.get('prompt_tokens')}, {dt:.1f}s) resp={txt[:60]!r}") + print(f"== NIAH {passed}/{total} passed ==") + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/run.sh b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/run.sh new file mode 100755 index 00000000..fe52d399 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_allgather/run.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# Kimi-K3 (MXFP4) on MI300X / gfx942 -- Wide expert-parallel, generic all2all. +# PP2xTP8 across 2 nodes for weight fit (~102 GB/GPU) PLUS --enable-expert-parallel +# so the 896 experts split 8-way across each node's 8 GPUs (112/GPU), replicated per PP +# stage. Uses the generic --all2all-backend allgather_reducescatter (see ../wideep_int4_moriep +# MoRI-EP kernels). AITER_SITUV2_A8W4=1 selects the a8w4 (fp8-act x int4-wt) SiTU +# MoE path. Colocated (single instance; no prefill/decode disaggregation). +# +# Usage (worker FIRST, then head): +# ROLE=worker MASTER= bash run.sh # on node1 (rank1) +# ROLE=head MASTER= bash run.sh # on node0 (rank0, serves API) +set -euo pipefail + +IMAGE="${IMAGE:-amdsiloai/vllm:kimi-k3-mi325x-release-v2}" +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR=}" +ROLE="${ROLE:?set ROLE=head|worker}" +MASTER="${MASTER:?set MASTER=}" +PORT="${PORT:-8000}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-8}" +GPU_UTIL="${GPU_UTIL:-0.90}" +CONTAINER="k3_wideepint4_${ROLE}" +LOGHOST="${LOGHOST:-$HOME/k3run/logs}"; mkdir -p "$LOGHOST" + +# Proven fabric env (from cluster_rdma_env_recommender.py) +BOOT_NIC="eth0" +IB_HCA="mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9" +GID_INDEX=3 + +if [ "$ROLE" = "head" ]; then + NODE_RANK=0; SERVE_EXTRA="--port ${PORT}"; HEADLESS="" +else + NODE_RANK=1; SERVE_EXTRA=""; HEADLESS="--headless" +fi + +echo "[pp2tp8] node=$(hostname -s) role=$ROLE rank=$NODE_RANK master=$MASTER" +[ -f "$MODEL_DIR/model.safetensors.index.json" ] || { echo "ERROR: model missing at $MODEL_DIR"; exit 1; } +docker rm -f "$CONTAINER" 2>/dev/null || true + +docker run -d --name "$CONTAINER" \ + --network host --ipc host \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband --group-add video \ + --cap-add SYS_PTRACE --security-opt seccomp=unconfined --shm-size 128g \ + --ulimit memlock=-1 --ulimit stack=67108864 \ + -e VLLM_ROCM_USE_AITER_MLA=0 \ + -e AITER_SITUV2_A8W4=1 \ + -e NCCL_SOCKET_IFNAME=$BOOT_NIC -e GLOO_SOCKET_IFNAME=$BOOT_NIC \ + -e NCCL_IB_DISABLE=0 -e NCCL_IB_HCA=$IB_HCA -e NCCL_IB_GID_INDEX=$GID_INDEX \ + -e NCCL_IGNORE_CPU_AFFINITY=1 -e NCCL_DEBUG=WARN \ + -e HSA_ENABLE_IPC_MODE_LEGACY=0 -e HSA_NO_SCRATCH_RECLAIM=1 \ + -e PYTORCH_ALLOC_CONF=expandable_segments:False \ + -e PYTORCH_HIP_ALLOC_CONF=expandable_segments:False \ + -v "$MODEL_DIR":/model:ro -v "$LOGHOST":/logs \ + --entrypoint bash \ + "$IMAGE" -c " + vllm serve /model --served-model-name kimi-k3 \ + --tensor-parallel-size 8 --pipeline-parallel-size 2 \ + --distributed-executor-backend mp \ + --enable-expert-parallel --all2all-backend allgather_reducescatter \ + --nnodes 2 --node-rank ${NODE_RANK} --master-addr ${MASTER} --master-port 29500 ${HEADLESS} \ + --trust-remote-code --reasoning-parser kimi_k3 --mm-encoder-tp-mode data \ + --safetensors-load-strategy prefetch \ + --max-model-len ${MAX_MODEL_LEN} --max-num-seqs ${MAX_NUM_SEQS} \ + --gpu-memory-utilization ${GPU_UTIL} ${SERVE_EXTRA} 2>&1 | tee /logs/vllm_wideepint4_${ROLE}.log + " +echo "[pp2tp8] $ROLE started. log: $LOGHOST/vllm_wideepint4_${ROLE}.log" +[ "$ROLE" = head ] && echo "[pp2tp8] health: curl http://${MASTER}:${PORT}/v1/models" diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/README.md b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/README.md new file mode 100644 index 00000000..0f5cd8fd --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/README.md @@ -0,0 +1,51 @@ +# Kimi-K3 (MXFP4) on MI300X / gfx942 — Wide-EP via MoRI-EP + a8w4 + +Serves Kimi-K3 (MXFP4) across **2 nodes** (16× MI300X) with **expert parallelism** over +true **MoRI-EP** all-to-all kernels: PP2×TP8 for weight fit (~102 GB/GPU) plus +`--enable-expert-parallel --all2all-backend mori_low_latency`, so the 896 experts are +split **8-way across each node's 8 GPUs** (112 experts/GPU, `[EP Rank x/8]`), replicated on +each of the 2 pipeline stages, and dispatched via MoRI all2all. The EP group is intra-node, +so MoRI-EP dispatch/combine runs across a node's 8 GPUs; the only cross-node traffic is the +PP activation hand-off (NCCL). +`AITER_SITUV2_A8W4=1` selects the a8w4 (fp8-activation × int4-weight) SiTU MoE path. +This is the "MoRI-EP + a8w4" path (vs the generic all2all in `../wideep_int4_allgather`). +Colocated (single instance; no prefill/decode disaggregation). + +- Image: `amdsiloai/vllm:kimi-k3-mi325x-release-v2` (gfx942 K3 vLLM build; public, anonymous pull) +- Model: `moonshotai/Kimi-K3` (MXFP4), on local NVMe or NFS + +## Run (worker FIRST, then head) + +```bash +# on the worker node (rank 1): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=worker MASTER= bash run.sh + +# on the head node (rank 0, serves the API on :8000): +MODEL_DIR=/path/to/Kimi-K3-MXFP4 ROLE=head MASTER= bash run.sh +``` + +- `MASTER` = the head node's **eth0** IP. `MODEL_DIR` = Kimi-K3-MXFP4 weights (prefer local NVMe). +- First start compiles the MoRI-EP dispatch kernels + a8w4 MoE (a few minutes), then cached. + +## Verify + +```bash +curl http://:8000/v1/models +python3 niah_probe.py --url http://:8000 --model kimi-k3 --ctx 8500 --depths 0.1,0.5,0.9 +``` +The response `system_fingerprint` contains `-ep-` when expert parallelism is active. + +## Key env (set in run.sh) + +| Var | Value | Why | +|-----|-------|-----| +| `VLLM_ROCM_USE_AITER_MLA` | `0` | Required on gfx942 (AITER MLA is gfx950-only). | +| `AITER_SITUV2_A8W4` | `1` | Route K3 SiTU MXFP4 MoE through the a8w4 interleaved flydsl kernels. | +| `MORI_GPU_ARCHS` | `gfx942` | MoRI-EP target arch. | +| `MORI_IB_GID_INDEX` / `MORI_IB_ENABLE_RELAXED_ORDERING` / `MORI_NUM_QP_PER_PE` | `3` / `1` / `8` | MoRI RDMA fabric tuning (**cluster-specific**). | +| `NCCL_IB_HCA` / `NCCL_IB_GID_INDEX` | 8× mlx5 / `3` | NCCL RDMA fabric (**cluster-specific — override for yours**). | +| `HSA_ENABLE_IPC_MODE_LEGACY` / `PYTORCH_(HIP_)ALLOC_CONF` | `0` / `expandable_segments:False` | ROCm 7.2.x requirements. | + +`--all2all-backend mori_low_latency` is used for both roles (low-latency MoRI-EP dispatch). +Overridable env: `IMAGE`, `MODEL_DIR`, `MASTER`, `PORT` (8000), `MAX_MODEL_LEN` (10240), +`MAX_NUM_SEQS` (8), `GPU_UTIL` (0.90). diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/niah_probe.py b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/niah_probe.py new file mode 100644 index 00000000..0fed8766 --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/niah_probe.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""Lightweight NIAH (needle-in-a-haystack) + basic accuracy probe for a served +Kimi-K3 endpoint. No deps beyond stdlib (urllib). Run from any node that can +reach the head IP:port. + +Usage: + python3 niah_probe.py --url http://:8000 --model kimi-k3 \ + --ctx 8000 --depths 0.1,0.5,0.9 +""" +import argparse, json, time, urllib.request, urllib.error + +FILLER = ("The quick brown fox jumps over the lazy dog. " + "Grand vistas of rolling hills stretched to the horizon. ") + +def chat(url, model, prompt, max_tokens=256, timeout=180): + body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, "temperature": 0, + }).encode() + req = urllib.request.Request(url.rstrip("/") + "/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + dt = time.time() - t0 + ch = d["choices"][0]["message"] + txt = (ch.get("content") or "") + " " + (ch.get("reasoning") or "") + return txt.strip(), d.get("usage", {}), dt + +def make_haystack(approx_words, needle, depth): + # Build filler to approx word count; insert needle at fractional depth. + words = [] + while len(words) < approx_words: + words.extend(FILLER.split()) + n = len(words) + pos = int(n * depth) + needle_words = needle.split() + words[pos:pos] = needle_words + return " ".join(words) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--url", required=True) + ap.add_argument("--model", default="kimi-k3") + ap.add_argument("--ctx", type=int, default=8000, help="approx haystack words") + ap.add_argument("--depths", default="0.1,0.5,0.9") + args = ap.parse_args() + + print(f"== sanity ==") + txt, usage, dt = chat(args.url, args.model, "Reply with only: OK", max_tokens=8) + print(f" resp={txt!r} usage={usage} {dt:.1f}s") + + secret = "The secret passcode is HELIOTROPE-7492." + question = ("\n\nQuestion: What is the secret passcode? " + "Answer with only the passcode.") + passed = 0; total = 0 + for depth in [float(x) for x in args.depths.split(",")]: + hay = make_haystack(args.ctx, secret, depth) + prompt = hay + question + try: + txt, usage, dt = chat(args.url, args.model, prompt, max_tokens=256) + except Exception as e: + print(f" depth={depth}: ERROR {e}") + total += 1; continue + ok = "HELIOTROPE-7492" in txt.upper() + passed += ok; total += 1 + print(f" depth={depth:>4}: {'PASS' if ok else 'FAIL'} " + f"(prompt_tok={usage.get('prompt_tokens')}, {dt:.1f}s) resp={txt[:60]!r}") + print(f"== NIAH {passed}/{total} passed ==") + +if __name__ == "__main__": + main() diff --git a/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/run.sh b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/run.sh new file mode 100755 index 00000000..05a35aed --- /dev/null +++ b/scripts/vllm/kimik3_mi300x/wideep_int4_moriep/run.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# Kimi-K3 (MXFP4) on MI300X / gfx942 -- Wide expert-parallel via true MoRI-EP. +# PP2xTP8 across 2 nodes for weight fit (~102 GB/GPU) PLUS --enable-expert-parallel +# with --all2all-backend mori_low_latency (MoRI-EP all2all kernels) + MoRI RDMA env. +# AITER_SITUV2_A8W4=1 selects the a8w4 (fp8-act x int4-wt) SiTU MoE path. Colocated +# (single instance; no prefill/decode disaggregation). +# +# Usage (worker FIRST, then head): +# ROLE=worker MASTER= bash run.sh # on node1 (rank1) +# ROLE=head MASTER= bash run.sh # on node0 (rank0, serves API) +set -euo pipefail + +IMAGE="${IMAGE:-amdsiloai/vllm:kimi-k3-mi325x-release-v2}" +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR=}" +ROLE="${ROLE:?set ROLE=head|worker}" +MASTER="${MASTER:?set MASTER=}" +PORT="${PORT:-8000}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-8}" +GPU_UTIL="${GPU_UTIL:-0.90}" +CONTAINER="k3_moriep_${ROLE}" +LOGHOST="${LOGHOST:-$HOME/k3run/logs}"; mkdir -p "$LOGHOST" + +# Proven fabric env (from cluster_rdma_env_recommender.py) +BOOT_NIC="eth0" +IB_HCA="mlx5_0,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_7,mlx5_8,mlx5_9" +GID_INDEX=3 + +if [ "$ROLE" = "head" ]; then + NODE_RANK=0; SERVE_EXTRA="--port ${PORT}"; HEADLESS="" +else + NODE_RANK=1; SERVE_EXTRA=""; HEADLESS="--headless" +fi + +echo "[pp2tp8] node=$(hostname -s) role=$ROLE rank=$NODE_RANK master=$MASTER" +[ -f "$MODEL_DIR/model.safetensors.index.json" ] || { echo "ERROR: model missing at $MODEL_DIR"; exit 1; } +docker rm -f "$CONTAINER" 2>/dev/null || true + +docker run -d --name "$CONTAINER" \ + --network host --ipc host \ + --device /dev/kfd --device /dev/dri --device /dev/infiniband --group-add video \ + --cap-add SYS_PTRACE --security-opt seccomp=unconfined --shm-size 128g \ + --ulimit memlock=-1 --ulimit stack=67108864 \ + -e VLLM_ROCM_USE_AITER_MLA=0 \ + -e AITER_SITUV2_A8W4=1 \ + -e MORI_GPU_ARCHS=gfx942 -e MORI_IB_GID_INDEX=3 -e MORI_IB_ENABLE_RELAXED_ORDERING=1 -e MORI_NUM_QP_PER_PE=8 \ + -e NCCL_SOCKET_IFNAME=$BOOT_NIC -e GLOO_SOCKET_IFNAME=$BOOT_NIC \ + -e NCCL_IB_DISABLE=0 -e NCCL_IB_HCA=$IB_HCA -e NCCL_IB_GID_INDEX=$GID_INDEX \ + -e NCCL_IGNORE_CPU_AFFINITY=1 -e NCCL_DEBUG=WARN \ + -e HSA_ENABLE_IPC_MODE_LEGACY=0 -e HSA_NO_SCRATCH_RECLAIM=1 \ + -e PYTORCH_ALLOC_CONF=expandable_segments:False \ + -e PYTORCH_HIP_ALLOC_CONF=expandable_segments:False \ + -v "$MODEL_DIR":/model:ro -v "$LOGHOST":/logs \ + --entrypoint bash \ + "$IMAGE" -c " + vllm serve /model --served-model-name kimi-k3 \ + --tensor-parallel-size 8 --pipeline-parallel-size 2 \ + --distributed-executor-backend mp \ + --enable-expert-parallel --all2all-backend mori_low_latency \ + --nnodes 2 --node-rank ${NODE_RANK} --master-addr ${MASTER} --master-port 29500 ${HEADLESS} \ + --trust-remote-code --reasoning-parser kimi_k3 --mm-encoder-tp-mode data \ + --safetensors-load-strategy prefetch \ + --max-model-len ${MAX_MODEL_LEN} --max-num-seqs ${MAX_NUM_SEQS} \ + --gpu-memory-utilization ${GPU_UTIL} ${SERVE_EXTRA} 2>&1 | tee /logs/vllm_moriep_${ROLE}.log + " +echo "[pp2tp8] $ROLE started. log: $LOGHOST/vllm_moriep_${ROLE}.log" +[ "$ROLE" = head ] && echo "[pp2tp8] health: curl http://${MASTER}:${PORT}/v1/models"