[DO NOT MERGE][AMD][DI][CI] Kimi-K2.6 disagg diagnostic recipes for non-MTP GSM8K drop

[michaelzhang-ai]

Motivation

On the Nightly Test (AMD MI355X 2N 1P1D Disagg), Kimi-K2.6 non-MTP GSM8K only reaches ~0.88 (gate is 0.92), while:

• the MTP variant passes at ~0.95, and
• single-node non-MTP passes at 0.944.

Already ruled out by prior investigation (see #30336 side-notes):

• page_size: 0.887 @256 vs 0.880 @1 — no meaningful difference.
• all disagg runtime args (kv_cache_dtype, disable_radix_cache, chunked_prefill, mem_fraction) are identical between the passing MTP and failing non-MTP legs.
• swapping the decode backend did not help.

Why the backend is likely NOT the cause: dispatch_attn_forward_method routes non-MTP decode → decode backend and MTP verify → prefill backend (because speculative_attention_mode defaults to prefill), but for their respective forward modes both dispatch to AttnForwardMethod.MLA (absorbed) — see models/deepseek_common/attention_backend_handler.py (handle_attention_aiter / handle_attention_triton). So the remaining axis is the decode FORWARD-MODE path reading the MORI-transferred prefix, not the attention backend/kernel.

SGLANG_FP8_PAGED_MQA_LOGITS_TORCH is intentionally NOT tested here: it only affects layers/attention/dsv4/ (the DSA indexer), which Kimi's aiter/triton MLA path does not use.

Modifications

Three diagnostic recipes under scripts/ci/slurm/recipes/mi355x-fp8/kimik26/1k1k/, registered in nightly-configs.yaml. Each isolates ONE variable vs the base 1p1d.yaml / 1p1d-mtp.yaml and trims the perf sweep to conc=1 (only the GSM8K gate matters). They are not part of the default nightly signal — trigger explicitly via the workflow_dispatch configs input:

config name	change vs base	what it tests
kimik26-fp8-1k1k-1p1d-aiterdecode	non-MTP, decode triton → aiter	control: is it the backend, or the decode forward-mode path? (expect still ~0.88)
kimik26-fp8-1k1k-1p1d-nocudagraph	non-MTP + --disable-cuda-graph	is it HIP graph capture/replay over the transferred prefix? (expect ~0.94 if yes)
kimik26-fp8-1k1k-1p1d-mtp-specdecode	MTP + --speculative-attention-mode decode	does routing verify through the decode-mode path also drop MTP to ~0.88? (confirms the axis)

Interpretation:

• If aiterdecode ≈ 0.88 and mtp-specdecode ≈ 0.88 → confirmed: the regression is the decode forward-mode path over MORI-transferred KV (metadata/graph), independent of backend and of MTP.
• If nocudagraph recovers to ~0.94 → the bug is in decode HIP graph capture/replay under disaggregation.

How to run

gh workflow run "Nightly Test (AMD MI355X 2N 1P1D Disagg)" \
  --ref amd/kimik26-disagg-decode-repro \
  -f configs=kimik26-fp8-1k1k-1p1d-aiterdecode,kimik26-fp8-1k1k-1p1d-nocudagraph,kimik26-fp8-1k1k-1p1d-mtp-specdecode

Base — must be stacked on #30386, not plain main. The MI355X disagg CI reinstalls the checked-out sglang via the runtime-checkout mechanism, which currently only exists on #30386 (amd-disagg-checkout). This branch is therefore based on 0aac976 (the #30386 head, and the exact commit the reference scheduled nightly ran on — it reproduces the clean 0.878 baseline). Running these recipe-only changes on plain main does NOT work: without the runtime-checkout mechanism the baked image is used and Kimi-K2.6 fails to load with ValueError: 'KimiK25ForConditionalGeneration' is not a registered model, and any later main regressions would also leak in. In short: rebase onto #30386 or the run eats the upstream error.

Checklist

• Diagnostic-only; no production code path changed.
• generate_matrix.py --runner mi355x produces the three new config names and --filter accepts them.
• All three recipe YAMLs parse via the launcher's recipe parser.

Validation runs

Baseline (bug is live, not a serving/MR failure): the scheduled nightly's base kimik26-fp8-1k1k-1p1d leg serves correctly (PD probe OK, coherent output) but the GSM8K gate reports 0.878 < 0.92 and aborts before the sweep — run 28873810961, job:

[probe] ok:  Paris. Paris is located in the north-central part of the country...
Accuracy: 0.878
[gsm8k] accuracy=0.878 threshold=0.92
[gsm8k] accuracy below threshold -- failing before sweep

Diagnostic run (this PR's 3 variants): run 28907609511 on the corrected base (#30386 / 0aac976), dispatched with configs=kimik26-fp8-1k1k-1p1d-aiterdecode,kimik26-fp8-1k1k-1p1d-nocudagraph,kimik26-fp8-1k1k-1p1d-mtp-specdecode.

An earlier run 28903343346 was mistakenly based on plain main and failed all three legs at model load ('KimiK25ForConditionalGeneration' is not a registered model) — concrete proof the #30386 base is required.

Results filled in once the legs finish:

config	Accuracy	reading
baseline 1p1d (non-MTP)	0.878	bug reproduced
reference 1p1d-mtp	~0.95	MTP verify (extend-mode / prefill backend) over the SAME transferred KV is correct
1p1d-aiterdecode	0.879	≈ baseline ⇒ NOT the attention backend (triton→aiter changes nothing)
1p1d-nocudagraph	0.883	≈ baseline ⇒ NOT HIP graph capture/replay (eager decode changes nothing)
1p1d-mtp-specdecode	crash at init (SIGABRT -6)	inconclusive: --speculative-attention-mode decode + EAGLE3 external draft fails to start; the decode-mode axis needs a code-level probe instead

(Diagnostic run: 28907609511.)

Conclusion so far

Two mechanistic suspects are eliminated: the attention backend (aiter decode ≡ triton decode ≈ 0.88) and HIP graph capture/replay (eager ≈ 0.88). Meanwhile MTP verify — which reads the same MORI-transferred KV but in extend-mode via the prefill backend — stays correct at ~0.95, and single-node non-MTP decode is 0.944.

So the drop is intrinsic to the single-token, absorbed-MLA decode forward-mode path reading disagg-transferred KV, independent of backend and graph. The transferred KV data itself is very likely fine (extend-mode reads it correctly). Prime remaining suspect: the decode forward metadata for the transferred prefix (seq_len / kv page indices / kv_indptr) on the decode worker, or the absorbed-MLA decode read of it.

Decode-side KV dump result (run 28909881181)

Added an env-gated one-shot probe (SGLANG_DEBUG_DISAGG_DECODE_DUMP) in ModelRunner.forward that logs, per decode batch on the decode worker: seq_lens, the req_to_token KV-slot mapping, and layer-0 KV-pool content at those slots. Captured the warmup probe ("The capital of France is", bs=1, decode seq_len 6→9):

[DISAGG_DECODE_DUMP #2] mode=DECODE bs=1 seq_lens=[6] | req0 seq_len=6 kv_slots=6 neg_slots=0 | KV L0 zero_rows=5 first3_norm=[0.0,0.0,0.0] last3_norm=[0.0,20.198,0.0]
[DISAGG_DECODE_DUMP #3] ... seq_lens=[7] ... zero_rows=5 first3_norm=[0.0,0.0,0.0] last3_norm=[20.198,21.346,0.0]
[DISAGG_DECODE_DUMP #5] ... seq_lens=[9] ... zero_rows=5 first3_norm=[0.0,0.0,0.0] last3_norm=[24.171,24.142,0.0]

Reading:

• seq_len == kv_slots, neg_slots=0 ⇒ the decode seq_len / req_to_token slot mapping is correct.
• The leading prefix KV rows read exactly 0.0 (first3_norm=[0,0,0], zero_rows constant), while only the locally-generated decode tokens are populated (contiguous non-zero norms ~20–24, growing by one per step; the trailing 0.0 is the current token, written later in the forward).

⇒ In the standard KV buffer (layer 0), the MORI-transferred prefill KV is absent (unwritten) at the decode-side attended slots. This points at KV transfer / decode-side KV placement, not the attention compute — consistent with backend & HIP-graph being ruled out.

Caveats to close next: only the short warmup probe was captured (5-dump budget spent before GSM8K); layer 0 only; and confirm this isn't a MORI staging-layout artifact (i.e. that the decode attention kernel really reads these zero slots). Follow-up: target a GSM8K request, dump several layers + the PREBUILT handoff, and compare the transferred length the prefill reported vs the decode prefix_indices.

Correction — KV transfer is fine on real (GSM8K) requests (run 28917883474)

Re-ran the probe skipping the short warmup (only seq_len >= 64), across layers 0/30/60, page-bucketing the zero rows. On a real GSM8K request (prefix = 1298 tokens):

[DISAGG_DECODE_DUMP] seq_len=1299 prefix=1298 page_size=256 neg_slots=0 |
  L0:zero=0/1298 mean=2.378e-01  L30:zero=0/1298  L60:zero=0/1298 | zero_pages=0/6 []

zero=0 on every layer, zero_pages=0/6, neg_slots=0. So the earlier "prefix KV is all zero" was an artifact of the tiny warmup probe (<1 page); for a real long prefix the MORI-transferred KV is fully present and correctly placed, and MTP-verify reading the same KV at ~0.95 confirms the values are correct too.

Revised conclusion. Eliminated: attention backend, HIP graph, KV-transfer/placement, and decode seq_len/slot metadata. Given identical correct KV + metadata, the ~0.88 vs single-node 0.944 (absorbed decode) vs ~0.95 (extend/verify) gap must live in the decode-mode compute itself under disaggregation — leading suspect is the query position / RoPE the decode worker reconstructs for the transferred prefix (an off-by-one there corrupts attention scores while leaving KV norms untouched), or another absorbed-MLA-decode numeric detail that single-node decode gets right.

Position / RoPE check (run 28974822944) — also clean

Extended the probe to log the decode current-token position vs seq_len-1. On GSM8K decodes:

[DISAGG_DECODE_DUMP] seq_len=1299 prefix=1298 ... pos0=1298 exp=1298 pos_vs_(seq-1)_mismatch=0
[DISAGG_DECODE_DUMP] seq_len=1300 prefix=1299 ... pos0=1299 exp=1299 pos_vs_(seq-1)_mismatch=0

pos0 == seq_len-1, mismatch=0 on every request/rank ⇒ decode positions/RoPE are correct.

Elimination summary

hypothesis	status	evidence
attention backend	ruled out	aiterdecode 0.879 == triton 0.878
HIP graph capture/replay	ruled out	nocudagraph 0.883
KV transfer / placement	ruled out	GSM8K prefix zero=0/1298 all layers, pages complete
decode metadata (seq_len/slots)	ruled out	seq_len==slots, neg_slots=0
KV values	correct	MTP-verify (MHA over same KV) = ~0.95
decode position / RoPE	ruled out	pos == seq_len-1, mismatch=0

Every ForwardBatch-level input to the decode attention is correct, yet absorbed-MLA decode over the transferred KV = ~0.88 while single-node absorbed-MLA decode = 0.944 and MHA/extend over the same KV = ~0.95. _dispatch_mla_subtype shows a pure decode step is always absorbed MLA — there is no server-arg to force decode→MHA, so the decode(absorbed)-vs-extend(MHA) axis can't be flipped by a recipe flag (the -mtp-specdecode attempt to flip it crashed at init).

Historical context (from #29855 + CI triage R241)

• [AMD][DI][CI] 3/N Add Kimi K2.6 FP8 MI355X 1P1D nightly recipes #29855 dev validation had disagg base Kimi at GSM8K 0.941 (Lzy17, 2-node g09/g29), and MTP at 0.948. So the disagg path did pass ~0.94 during development — the current ~0.88 is likely a regression, not an inherent disagg limitation. This is consistent with every ForwardBatch-level attention input being correct in my probes.
• Timeline: [AMD][DI][CI] 3/N Add Kimi K2.6 FP8 MI355X 1P1D nightly recipes #29855 dev 0.941 -> bingxche 07-03 (PR29986-hotpatch image) 0.892 -> current nightly 0.887.
• CI triage (R241) names in-flight model-loading suspects: Fix Kimi K2.6 compressed-tensors MoE w2 scale loading #28905 (K2.6 MoE w2 scale loading) and [AMD] Add MI355X Kimi-K2.6 tuning artifacts #23381 (K2.6 tuning artifacts). A weight/quant-loading bug would hit single-node and disagg equally — which the single-node control below directly tests.

Single-node control result (run 28984448331) — decisive

Same passing single-node MI35x Kimi eval, default vs --page-size 256, back-to-back on one image:

default        accuracy=0.947  PASS
page_size=256  accuracy=0.945  PASS

• page_size ruled out: single-node@256 = 0.945 (== default). (Matches disagg page 1 vs 256 both ~0.88 — page size is not the variable.)
• weight/quant loading ruled out (Fix Kimi K2.6 compressed-tensors MoE w2 scale loading #28905 w2 scale etc.): single-node = 0.947; a load bug would degrade single-node too.

Verdict: the drop is disagg-specific, and a regression (vs #29855 dev disagg 0.941). Single-node absorbed-MLA decode = 0.947 with correct KV; disagg absorbed-MLA decode = ~0.88 with also-correct KV/metadata/positions. Every non-disagg explanation is eliminated.

Convergence with the parallel amd/kimik26-disagg-decodemeta branch (yctseng0211)

A parallel investigation reached the same conclusion and is now probing the decode compute directly:

• decode-metadata + MORI transfer-fidelity probes (extend-mode + fixed-absolute-KV-position sampling): disagg runs still 0.876 / 0.879 / 0.885 — probes only observe; the transferred KV is confirmed present and the drop persists. Independently corroborates: KV correct, decode compute is the issue.
• New live hypothesis: the triton MLA decode split-KV online-softmax reduction (num_kv_splits). It is decode-only (extend/verify use no split-KV → explains why they stay ~0.95) and reduces across KV splits over the long transferred prefix — a plausible precision-loss site. Tested via --triton-attention-num-kv-splits 1 + SGLANG_DEBUG_MLA_DECODE_NO_SPLIT_CAP (recipe 1p1d-nsplit1).
• nsplit=1 result is currently unavailable: run 28983593561 finished as failure with no artifact and its log has been purged (BlobNotFound) — needs a re-run to get a clean number.

Update — nsplit=1 result (run 29042111257, ported to this branch): accuracy 0.875 → split-KV reduction is RULED OUT. Forcing a single decode KV split (--triton-attention-num-kv-splits 1 + no-split-cap + static fill) leaves accuracy at baseline ~0.88. So the online-softmax cross-split reduction is not the cause.

This is the most promising lead and supersedes the python/image bisect for now (the stuck bisect run was cancelled to free the single MI355X slot).

Regression bisect — step 1: python is NOT the cause (run 29058233535)

Reverted python/sglang/srt to 67361ff (#29855 kimi-merge, the 0.941 era) and ran disagg base Kimi via checkout-runtime (artifact confirms sha=629912085, 0.5.15.dev520). Result: 0.879 — still ~0.88.

Key confound found: the disagg workflow's setup job resolves the latest MI35x image dynamically and overrides the recipe's pinned tag. So:

• step 1 = old python (67361ff) + latest image = 0.879
• Lzy17's 0.941 = old python + v0.5.13.post1-rocm720-mi35x-20260623 (Jun-23 image)

⇒ reverting python doesn't help; the differing variable is the image (aiter / ROCm / sgl-kernel). Also relevant: DI CI and Lzy17's 0.941 use the same physical amd-sglang g09/g29 nodes (the DI runner only sallocs them), so pass/fail is a code/image version difference, not a machine difference.

Step 2 (in flight): full known-good repro — bisect branch (old python) + Jun-23 image pinned (-f image=...20260623). If ~0.94, the regression is the image and known-good is reproduced; then bisect the image. If still ~0.88, the regression is neither current code nor image (environmental: NFS weights / aiter build).

BISECT RESULT — the regression is the DOCKER IMAGE, not sglang code (run 29116830660)

Reproduced Lzy17's known-good cleanly: full checkout @ 67361ff + June-23 image (v0.5.13.post1-rocm720-mi35x-20260623), no checkout-runtime (that commit predates it, so it runs the image's coherent baked sglang) → GSM8K 0.941 ✅ (== Lzy17's 0.941).

run	python	image	GSM8K
scheduled baseline	main	latest	0.887
bisect step 1	srt@67361ff	latest	0.879
known-good repro	full 67361ff	Jun-23	0.941

• known-good (old code + old image) = 0.941 ⇒ environmental causes (g09/g29 nodes, NFS weights) are ruled out; known-good is reproducible.
• step 1 (old srt + latest image) = 0.879 ⇒ swapping in the newer image breaks it even with old python.
• ⇒ the regression is in the docker image (aiter / ROCm / sgl-kernel), not sglang python.

(Note: bisect step 2 = mixed python on old image crashed at init on an ABI mismatch — discarded, uninformative.)

Immediate mitigation available

Pin the Kimi-K2.6 disagg recipes to the Jun-23 image (or set the workflow image input) to restore ~0.94 while the image regression is bisected. The scheduled nightly overrides the recipe's pinned image with the dynamically-resolved latest tag, which is what regressed.

Image bisect result — first bad build is the v0.5.14 bump (Jun-26)

image tag	version	GSM8K
...mi35x-20260623	v0.5.13.post1	0.941 ✅
...mi35x-20260626	v0.5.14 (first)	0.888 ❌ (run 29277682836)

Jun-24 and Jun-25 are still v0.5.13.post1 (same line as the good Jun-23), so the regression is the v0.5.13.post1 → v0.5.14 image bump landing on Jun-26 (aiter / ROCm / sgl-kernel changes baked into the v0.5.14 image). Confirmed airtight: Jun-25 (v0.5.13.post1) = 0.939 (run 29282623403) — last good is Jun-25, first bad is Jun-26 (v0.5.14). Mitigation (pin known-good image + baked build) verified green at 0.939 (PR #31063).

Next: bisect the image

Sweep lmsysorg/sglang-rocm:v*-rocm720-mi35x-YYYYMMDD tags between Jun-23 (good) and the current latest (bad) via the workflow image input on the 67361ff known-good branch, to find the first bad image build.

Open next steps (branching) — now that it's confirmed disagg-specific + a regression

1. KV latent dtype/scale check (cheapest) — dump the compressed-KV dtype + raw nope/k_pe value stats on the decode side; if disagg stores/transfers the latent in a different precision than single-node, absorbed-MQA (direct dot-products) would degrade where MHA-decompress (large matmul) tolerates it.
2. Bisect the regression — [AMD][DI][CI] 3/N Add Kimi K2.6 FP8 MI355X 1P1D nightly recipes #29855 dev disagg was 0.941; run disagg Kimi via checkout-runtime at the [AMD][DI][CI] 3/N Add Kimi K2.6 FP8 MI355X 1P1D nightly recipes #29855 validation commit, confirm ~0.94, then bisect forward to the breaking change (python and/or image/aiter kernel).
3. In-process attention diff — hook the MLA layer to compare absorbed-decode vs MHA-decompress output on the same transferred KV (decisive: is the absorbed-decode kernel itself wrong on transferred KV?).

Next step — this branch already runs via the runtime-checkout mechanism on #30386): a code-level probe on the decode worker for a fixed prompt — compare (a) per-layer KV-pool bytes and (b) the decode forward metadata vs a single-node baseline. A data match + metadata/read mismatch localizes the bug.

---

CI States

Latest PR Test (Base): ❌ Run #29042115255
Latest PR Test (Extra): ❌ Run #29042114884

[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!