[AMD] WIP - Amd/kimik26 disagg decodemeta

[yctseng0211]

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #28983578683
Latest PR Test (Extra): ❌ Run #28983578532
Latest PR Test (AMD ROCm 10): ➖ No AMD PR run found for this commit.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[yctseng0211]

Triggered: Nightly Test (AMD MI355X 2N 1P1D Disagg) run https://github.com/sgl-project/sglang/actions/runs/28957394001 on this branch, config kimik26-fp8-1k1k-1p1d-metadump.

What this probes

Prior analysis (see #30433) narrowed the non-MTP GSM8K drop (~0.88 vs single-node 0.944) to the decode-mode read of the MORI-transferred prefix, independent of attention backend (aiterdecode ≈ triton), HIP graph (nocudagraph ≈ baseline), and paging (page_size=1 is also ~0.88).

Reading the two decode kernels confirms triton decode_attention_fwd and aiter _mla_decode_fwd_with_head_pad consume the same scheduler metadata — kv_indptr / kv_indices / seq_lens / num_kv_splits (built from req_to_token) — as their only shared input. The verify/extend path uses a different set (qo_indptr / custom_mask / mask_indptr) and reads the current token's k/v from direct args, which is why MTP verify over the same transferred KV stays correct at ~0.95. So the remaining suspect is the decode-mode metadata (or the transferred KV values) for the prefix.

This branch adds an env-gated (SGLANG_DEBUG_DISAGG_DECODE_META), read-only, capture-safe probe in {triton,aiter}.forward_decode that dumps, per real decode step: kv_indptr / kv_indices (head/tail) / seq_lens / num_kv_splits, an idx==rtt sanity check, and the KV-pool row norms at the attended slots. 1p1d-metadump.yaml wires it with --disable-cuda-graph (eager decode so forward_decode is entered in Python; with HIP graph, decode is replayed and the probe never fires) and MINLEN=128 (skips warmup + the short PD probe so the dump budget lands on real GSM8K decode over the long transferred prefix — fixes the warmup-only limitation of the earlier KV dump).

How to read the log

Decode server log in the run artifact, grep -A6 "\[DDM". Best diffed against a single-node eager run of the same prompt:

• seq_len per step must match single-node — a mismatch is an off-by-one in the disagg decode seq_len.
• KV L0 zero_rows must be 0 — >0 means the transferred prefix KV is absent at the attended slots (on a real request this time, unlike the warmup-only dump).
• KV-norm head/tail sequence: same prompt ⇒ same norms regardless of physical slots — a mismatch means the transferred KV values are wrong.
• idx==rtt is a sanity check (kv_indices is a copy of req_to_token[req, :seq_len]).

Caveat

This branch merged main on top of the #30386 (0aac976) checkout base. If the run fails at model load (KimiK25ForConditionalGeneration is not a registered model / cutlass import), reset to the 0aac976-only base (drop the merge) and re-trigger — same reason #30433 stayed pinned there.

[yctseng0211]

:https://github.com/sgl-project/sglang/actions/runs/28962093651

[yctseng0211]

Probe results so far

Completed runs 28957394001 (acc 0.876) and 28962093651 (acc 0.879); new run 28963785464 in progress. The decode-metadata probe fires on real GSM8K decode (not warmup), bs=1, 8 TP ranks, seq_len ~1263–1299.

Refuted: the "transferred KV is absent" reading was a warmup-only artifact

On real decode steps:

• zero_rows = 0/seq_len at every attended slot → the transferred prefix KV is present and non-zero.
• idx==rtt: True → kv_indices exactly equals req_to_token[req, :seq_len].
• seq_len monotonic (+1/step) → no off-by-one.

So the decode metadata (indices / length) and KV presence are all correct.

Transfer-fidelity check (prefill-stored vs decode-received, within ONE disagg run)

Added a prefill-side (mode=EXTEND) dump so a single run yields both what the prefill worker stored and what the decode worker received, over the shared 8-shot prefix. The first-page latent norms match exactly:

prefill EXTEND (rp=5, seq_len=1232):  KV L0 norm head = [19.617, 22.029, 17.145, 23.779, 17.698, 19.234]
decode  DECODE (rp=6, seq_len=1299):  KV L0 norm head = [19.617, 22.029, 17.145, 23.779, 17.698, 19.234]

(The ~1e-3 spread across prefill requests is fp-nondeterminism, same magnitude on both sides — not transfer corruption.)

This points away from the transfer / metadata and toward the decode compute path (or state the decode worker lacks because it never ran a local prefill/extend for the request).

Caveat + what run 28963785464 adds

The head sample only covers page 0 (page_size=256). Run 28963785464 adds fixed absolute-position samples samples@abs={0,128,256,512,768,1024} spanning pages 0–4 — all within the shared 8-shot prefix, so directly comparable prefill vs decode:

• all positions match → whole-prefix transfer is faithful → the bug is in decode compute / missing local-extend state;
• a ≥256 position mismatches → a per-page transfer bug in disaggregation/mori/conn.py::send_kvcache (grouping/offset).

[yctseng0211]

https://github.com/sgl-project/sglang/actions/runs/28963785464