Add Kimi-K3 MI300X (gfx942) serving recipes (3 colocated + 2P/2D EP16 disagg, NIAH to 900K)

[raviguptaamd]

Summary

Adds scripts/vllm/kimik3_mi300x/ — self-contained vLLM serving recipes for
Kimi-K3 (MXFP4) on AMD Instinct MI300X / gfx942, complementing MAD's existing
single-node gfx950/MI355X K3 recipe (benchmark/kimi_k3/).

Why MI300X needs its own recipe: MI300X has 192 GB/GPU, so the ~1.5 TB K3 checkpoint
does not fit a single 8-GPU node under TP8. The colocated recipes 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).

This PR contains three colocated recipes plus a fourth prefill/decode
disaggregated (2P/2D) recipe — validated: single-needle NIAH passes deterministically
across the full native context range 10K–900K.

Colocated recipes — each: run.sh + README.md + niah_probe.py

Folder	Parallelism	Expert all2all	MoE path
pp2xtp8/	PP2×TP8, no EP	—	a16w4
wideep_int4_allgather/	PP2×TP8, EP8/node	allgather_reducescatter (generic)	a8w4 (AITER_SITUV2_A8W4=1)
wideep_int4_moriep/	PP2×TP8, EP8/node	mori_low_latency (MoRI-EP)	a8w4 (AITER_SITUV2_A8W4=1)

All colocated (single instance; no prefill/decode disaggregation). These use the public
vendor image amdsiloai/vllm:kimi-k3-mi325x-release-v2, overridable via -e IMAGE=.

EP scope note: the 896 experts split 8-way across each node's 8 GPUs (112/GPU, [EP Rank x/8]), replicated on each of the 2 PP stages. Expert all2all runs intra-node; cross-node traffic is the PP activation hand-off over NCCL. ("16" = GPU count, not EP width.)

wideep_disagg_2p2d/ — 2P/2D EP16 MoRIIO disaggregation (✅ validated)

A prefill/decode disaggregated Kimi-K3 recipe across 4 MI300X nodes: 2 prefill +
2 decode, TP2×DP8 → EP16 per pool via MoRI-EP all2all (no pipeline parallelism),
joined by the MoRIIO connector for prefill→decode transfer of both the MLA attention
KV and the Kimi-Delta-Attention (KDA/mamba) recurrent+conv state.

Three root-cause fixes (all folded into vLLM source on the branch below):

1. 4-KV-cache-group block routing. K3's hybrid attention allocates 4 KV-cache
   groups (3 KDA/mamba + 1 MLA); the connector hardcoded 2-group indices and sent MLA KV to
   mamba block-ids → decode read empty blocks (fluent but context-free). Fix routes each
   layer by its own group index.
2. Multi-chunk prefill transfer. The final-chunk gate used block count, which fires
   after chunk 1 when a prompt fits in ≤1 padded block → only max_num_batched_tokens of KV
   crossed (a razor cliff at 2048). Fix gates on compute progress from scheduler_output.
3. KDA gather sync-free (>500K unblock). gather_initial_states ran a diagnostic
   bool((indices>=n).any()) device→CPU sync per KDA layer per prefill chunk (~25k full
   stream drains at 750K), stalling long-context prefills into an apparent hang (py-spy'd:
   the DP rank with real work stuck in the sync while other DP ranks waited at the batch
   all_reduce). The clamp above already made the address safe; the diagnostic is now gated
   behind K3_KDA_GATHER_LOG=1 (default OFF), correctness unchanged.

Results — single-needle NIAH (HELIOTROPE-7492, greedy, depths 0.1/0.5/0.9)

ctx	result	eval time/req
10K–200K	3/3 PASS	5–88s
300K	3/3 PASS	~150s
500K	3/3 PASS	~301s
750K	3/3 PASS	~534s
900K	3/3 PASS	~717s

Deterministic across the full range; sub-quadratic scaling. Throughput (the DP disagg
payoff, 20K ctx): 5.7× at concurrency 8, 7.3× at 16 vs single-stream. Single-stream
latency is ~4× a colocated PP2×TP8 serve by design (1 request on 2 GPUs vs all 16) — use
disagg for high-QPS/batch, colocated for low-latency interactive. Tables in RESULTS.md;
tuning in OPTIMIZATION.md (KV_CACHE_MEMORY_BYTES=40e9 → 2.84M-token KV cache, required
for single requests > ~600K; MAX_MODEL_LEN=1000000 for full native ctx).

Known residual (does not block single-needle NIAH): the stricter 10-needle stress dips
to ~9/10 at ≥20K (an RDMA write-visibility race); single-needle is unaffected. Documented in
STATUS.md. Note: MAX_NUM_BATCHED_TOKENS=8192 still corrupts generation on this stack —
keep the default 2048.

Self-contained image: this folder ships Dockerfile.kimik3_disagg, which builds the
whole stack from source on a public base (rocm/vllm-dev:ci_base-...): MoRI v1.2.2 +
AITER 0.1.19 + vLLM (from the fork branch below, all fixes baked in) + the
DP-rank/KV-notify vllm-router. Validated from a scratch build (NIAH re-passed with the
runtime patchers reporting "already applied" — idempotent no-ops on a folded image).

Companion vLLM branch: the connector + KDA fixes are folded into vLLM source at
raviguptaamd/vllm branch kimi-k3-wideep-disagg-fullsource-v3; the Dockerfile builds it
directly (VLLM_REF=kimi-k3-wideep-disagg-fullsource-v3). The runtime patchers under
wideep_disagg_2p2d/patchers/ no-op idempotently on a v3 image (so the recipe also works on
an unfolded image).

Notes

• Fabric env (NCCL_IB_HCA/GID, MORI_*, SOCKET_IFNAME) is fully overridable; defaults
  are the validated Broadcom Thor2 (bnxt) values with a THOR2_BNXT_FIX toggle. Override
  for other fabrics (e.g. Mellanox mlx5) per the README.
• A fresh serve pays a one-time aiter MLA-kernel JIT (15 min) on the first ≥200K request;
  cached thereafter.

Test plan

• bash -n on all run.sh / launch scripts — pass.
• python3 -m py_compile on all .py (probes + patchers) — pass.
• Disagg image built from scratch off the vLLM fork branch — succeeds (incl. the
  MAX_JOBS empty-env build fix).
• Disagg NIAH validated: single-needle PASS across 10K–900K, all depths,
  deterministic; fixes served from the baked-in image.
• No local-cluster leakage (paths/IPs/job IDs/tokens) and no bytecode committed.

AI assistance (Claude) was used to prepare this PR; changes were human-reviewed.

[raviguptaamd]

Pushed a 4th recipe subfolder: scripts/vllm/kimik3_mi300x/wideep_disagg_2p2d/ — a 2P/2D EP16 MoRIIO disaggregated Kimi-K3 recipe (2 prefill + 2 decode nodes, TP2×DP8→EP16 per pool, no PP).

Status: WIP / known-issue — not a working deployment. The disagg transport is byte-perfect (attention KV + KDA/mamba state arrive byte-identical on all 8 decode ranks) and all connector-level bugs are fixed, but exact multi-token NIAH recall is still wrong and non-deterministic at greedy temp=0 (a forward-pass bug below the sampler, downstream of the proven-correct KV). Full repro + ruled-out list in wideep_disagg_2p2d/STATUS.md.

Unlike the colocated recipes, this folder ships a self-contained Dockerfile.kimik3_disagg that builds the whole stack from source on a public base (rocm/vllm-dev:ci_base-...): MoRI 1.2.2 + AITER 0.1.19 + vLLM from raviguptaamd/vllm@kimi-k3-wideep-disagg-fullsource (connector fixes folded into source, one commit per fix) + the DP-rank/KV-notify vllm-router.

Rebased cleanly on top of the recent docs commits already on this branch. Intend to split the disagg work into its own PR later; grouping here for review context.