[AMD][CI] Move the MI35x DeepSeek-V4-Flash nightly to the official 0731 checkpoint (blocked: triton FlashMLA accuracy regression)

[michaelzhang-ai]

Blocked. One line — the MI35x DeepSeek-V4-Flash FP4 nightly's default checkpoint — but it fails the triton FlashMLA accuracy leg, and the fix is not in this PR. Details below; the useful output here is the regression, not the diff.

The other two changes this branch once carried have both merged: the MI35x V3.2 four-into-two job merge as #36736, and the V3-0324 / V3.1 MI30x drop as #36396.

Motivation

The MI35x Flash nightly serves the preview checkpoint on its FP4 leg. deepseek-ai/DeepSeek-V4-Flash-0731 "is the official release of DeepSeek-V4-Flash, superseding the preview version" per its model card, has ~2x the downloads, and is far stronger on agentic benchmarks (Terminal Bench 2.1 61.8 -> 82.7, DeepSWE 7.3 -> 54.4). The cookbook also ships MI355X flash-official (0731) FP4 cells marked verified: false that nightly coverage would back.

Modifications

test/registered/amd/test_deepseek_v4_flash_fp4.py: DEEPSEEK_V4_FP4_MODEL_PATH default becomes deepseek-ai/DeepSeek-V4-Flash-0731, plus three docstring lines on why target-only serving is safe for a checkpoint carrying a draft head. The env override is unchanged, so the preview can still be pinned.

Structurally the swap is inert: the two config.json files are identical for the target (43 layers, head_dim 512, index_topk 512, same quantization_config, expert_dtype: fp4), 0731 adds only four dspark_* keys, and DeepseekV4ForCausalLM.load_weights skips mtp.* outside a draft worker so the bundled DSpark head never reaches VRAM.

FP8 legs are untouched and stay on sgl-project/DeepSeek-V4-Flash-FP8 — FP8 needs a repackaged checkpoint and no 0731 repack exists.

Accuracy Tests

The Flash job runs the suite twice, once per SGLANG_HACK_FLASHMLA_BACKEND. GSM8K on the FP4 leg, from run 32906052690 (rocm720, continue_on_error=false):

FlashMLA backend	preview (main)	0731 (this PR)	delta
unified_kv_triton	0.9280	0.9227	-0.005
triton	0.9249	0.9007, 0.8908 on retry	-0.024 to -0.034

unified_kv_triton passed cleanly — 0731 FP4 at 0.9227, untouched FP8 legs at 0.9287 and 0.9325 (TBO). triton missed the 0.91 floor on both attempts and stopped there, so the FP8 legs under triton were not reached.

Not flake, not pre-existing: the triton step passed in all four Flash job runs across the 8/24 (32757901075) and 8/26 (32997660973) nightlies, both image flavors, with the checkpoint as the only variable here.

What the failure is, and is not

Mechanically: test_a_gsm8k asserts accuracy > 0.91, CustomTestCase allows one retry in CI, both attempts were under, so the retry helper raised and the step exited 1.

Two candidate causes ruled out:

• Not truncation or malformed output. invalid is 0.0 on both failing runs, so the model emits well-formed answers that are simply wrong. This kills the theory that 0731's max reasoning effort produces longer chains of thought that hit the eval's max_new_tokens=512.
• Not a compress_ratios index shift. 0731 has 46 entries against the preview's 44, but the first 44 are element-wise identical, the extras are trailing zeros for the DSpark head layers, and the derived num_c4_layers is 21 either way.

One finding that complicates the picture: the triton cell is nondeterministic. Both 0731 evals ran against a single server (launched 12:28, ready 12:31; evals at 12:33 and 12:34) at temperature 0, and scored 0.9007 then 0.8908 — roughly 13 of 1319 questions flipping between identical runs. So that cell carries >=1 point of noise on its own, and the 0.91 floor only ever had 1.5 points of headroom there against the preview's 0.9249.

The 2.9-point mean gap (0.8958 vs 0.9249) still exceeds that noise, which favours a real checkpoint x backend interaction rather than luck. But preview-on-triton is a single sample, so its variance in that cell is unknown.

The two backends are separate implementations, not a tuning flag: triton uses nsa_triton_decode with the separate packed KV layout, unified_kv_triton uses dsv4/unified_kv_kernels with the unified layout. Both run --kv-cache-dtype fp8_e4m3, so FP8 KV quant/dequant and the sparse indexer (index_topk=512) are the natural places to look.

Options

Lowering the 0.91 floor or dropping the triton leg would turn this green while discarding the result. That sweep exists (#28290) to catch backend divergence on this model, and it worked. So, for the DSV4 / AMD owners:

0. First, characterise the variance. Given the nondeterminism above, 3-5 repeats each of preview-on-triton and 0731-on-triton would separate "real regression" from "thin threshold meeting a noisy cell" before anyone goes into the kernel. Cheap relative to a kernel hunt, and it decides between the options below.
1. Fix the triton FlashMLA path against 0731, then land this unchanged. unified_kv_triton is what the MI355X cookbook recipes use, so this would be a non-production backend regressing on the official checkpoint — worth understanding independently of this PR. The temperature-0 nondeterminism is arguably its own bug regardless of the checkpoint.
2. Take 0731 and narrow the FP4 sweep to unified_kv_triton, with a tracking issue. Cheapest; trades away the signal.
3. Stay on the preview until 1 is resolved, accepting that the nightly keeps validating a superseded checkpoint.

For the record, my pre-run assessment called this low-risk from config equality plus the loader's mtp.* skip. That was structurally correct and insufficient — it says nothing about numerics, and numerics are what moved.

Speed Tests and Profiling

No cost change; this is a currency fix. On the unified_kv_triton leg that passed, 0731 measured 57.6 s / 2238 tok/s, in line with the preview. test_b_perf_8k_1k asserts no threshold, so perf shifts surface in the reported table rather than as failures.

If this lands: stage-c-dsv4-flash-fp4-fp8-amd-mi35x-rocm720 gates every main-package PR on the Flash suite with a 60-minute step timeout. 0731 is 155 GiB against the preview's 149, cached per host under HF_HOME=/sgl-data/hf-cache, so the first run on each MI35x 8-GPU host pays a cold download and could exceed that timeout. Pre-warm /home/runner/sglang-data/hf-cache, or pin the preview on the PR gate via DEEPSEEK_V4_FP4_MODEL_PATH.

Checklist

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

---

CI States

Latest PR Test (Base): ❌ Run #34829118668
Latest PR Test (Extra): ❌ Run #34829118366
Latest PR Test (AMD ROCm 10): ❌ Run #34829118509