[zovonoir]

#2637 has merged; this branch is rebased onto main.

Summary

Qwen3.8-Flash-Next MTP on SGLang-ATOM runs every draft step eagerly. #2385 skips draft graph capture for Qwen4ExpForCausalLMNextN, citing the [3, N] mRoPE positions that native ATOM's DraftGraph cannot stage. SGLang's EAGLEDraftCudaGraphRunner / EAGLEDraftExtendCudaGraphRunner already carry mrope_positions buffers, so that reason does not hold on the SGLang path. As #2385 itself measured, the eager draft step is host-launch bound: one eager draft decode step costs about 8 ms, while the same step graphed costs about 0.7 ms. With MTP that overhead more than cancels the accepted tokens, so after #2418 MTP was slower than no MTP (TP2/EP1).

Both draft graphs are now captured whenever the Flash drafter runs (no flags; --disable-cuda-graph turns them off with every other graph):

Graph	Effect
Draft decode	Capture the Flash draft decode graph. The draft runner's out-of-graph hook no longer fills the target's _DECODE_GRAPH QSA buffers (TARGET_VERIFY replays from them); inside the captured draft forward the QSA metadata is rebuilt into the private _DRAFT_QSA buffers, so each draft step's positions / out_cache_loc are recomputed on replay.
Draft extend	Also graph the draft extend after every verify. SGLang's HIP gate only admits aiter multi-step / DSV4 / DSA backends, so the Flash drafter takes the plugin-built EAGLEDraftExtendCudaGraphRunner path DSV4 / GLM-5.2 already use (without their long pad-row seq_len fill).

Also fixes the plugin's draft-extend replay wrapper, which took only (self, forward_batch) while SGLang 0.5.20 calls execute(forward_batch, select_index); any draft-extend graph on this SGLang hit TypeError before.

Plugin-only: atom/plugin/register.py, atom/plugin/sglang/qwen4_exp_bridge.py. No SGLang change, no new environment variables.

Results

2x MI308X, Qwen3.8-Flash-Next-PTPC-FP8, sglang.bench_serving random (the #2418 harness: range 1.0, OSL 2048, num-prompts = 2 x concurrency, 1 warmup).

Not a recovered regression. Current main + #2637 + this PR on rocm/atom-dev:sglang-latest e021a3935c1c. Baseline = #2418 recipe (TP2/EP1, no MTP); MTP = TP2/EP2, EAGLE 4 steps, both graphs, ctx 132096, max-running 64, decode graph bs 64. The baseline matches #2418's published numbers (output throughput geomean -0.2%, every case within 3.4%), and MTP is measured on top of it:

ISL x CON	TPOT #2418 / baseline / MTP (ms)	out tok/s #2418 / baseline / MTP	MTP vs baseline	accept len
1024 x 1	6.43 / 6.50 / 4.24	153.4 / 151.9 / 230.2	+51.6%	3.74
1024 x 8	11.59 / 11.50 / 8.44	662.4 / 678.9 / 769.3	+13.3%	3.23
1024 x 32	20.48 / 19.95 / 17.91	1,476.6 / 1,507.5 / 1,462.5	-3.0%	3.00
4096 x 1	6.75 / 6.61 / 4.23	145.1 / 148.1 / 227.4	+53.5%	3.70
4096 x 16	16.14 / 16.06 / 12.13	910.2 / 894.7 / 1,051.2	+17.5%	3.01
6144 x 4	10.70 / 10.60 / 7.21	336.7 / 344.2 / 468.4	+36.1%	3.02
30720 x 1	6.63 / 6.78 / 4.49	130.5 / 127.9 / 177.3	+38.6%	3.61
30720 x 8	14.91 / 15.07 / 13.48	394.4 / 381.0 / 437.0	+14.7%	3.28
63488 x 4	14.13 / 14.40 / 11.81	201.6 / 197.1 / 233.6	+18.5%	3.17
63488 x 16	31.64 / 31.71 / 37.06	318.2 / 317.5 / 316.8	-0.2%	3.46
129024 x 1	7.38 / 7.39 / 4.82	83.7 / 83.5 / 101.3	+21.3%	3.57

MTP vs baseline: output throughput geomean +22.5%, TPOT -23.5%, ahead in 9/11 cases.

Draft steps, release-note launch (TP2/EP2, ctx 8192, max-running 8), ISL 1024 / OSL 2048, concurrency 1, on sglang-v0.5.20-nightly_20261012 (TPOT; no MTP TP2/EP1 on the same image: 6.4 ms):

steps	eager draft	decode graph	decode + extend graph	GSM8K 400, both graphs
2	12.37 ms	9.16 ms	5.90 ms	97.00%
3	13.57 ms	8.19 ms	6.01 ms	97.00% (accept 3.44, same as eager extend)
4	-	8.69 ms	5.17 ms	97.25%
5	-	-	5.02 ms	97.75%

Full #2418 sweep (29 cases, ISL 1024-129024, concurrency 1-32, 668/668 requests) on the nightly image, both graphs, output throughput geomean vs no MTP TP2/EP1: +14.5% (2 steps), +15.4% (4 steps), ahead in 25/29. Behind: 63488 x 16/32 and 129024 x 8/16 (-0.3% to -9.2%), where long-context, high-concurrency decode dominates and EP2 costs the target about 2.3 ms per step.

Accuracy / multimodal. GSM8K 400 (run_eval, chat, thinking off) 97.0-97.75% across configurations, eager 97.75%. 1080p image + ~3.8K-token text, concurrency 1-16: every answer correct.

Review follow-up (b5809b6df, 94cef79ce)

Rebased onto main (f30f8b0cc, which includes #2637); results below are on this tree.

mRoPE in the draft graphs. Confirmed with a probe on the draft runners: the decode graph's mrope_positions buffer was never written, so every draft step ran at RoPE position 1, and the draft-extend graph read the buffer it shares by name and shape with the target-verify runner. The same staleness existed in eager drafting: EAGLE draft_forward advances only the 1-D positions, so draft steps 1+ reused step 0's mRoPE.

• Before each draft decode / draft-extend replay, stage SGLang's live mrope_positions (multimodal offsets kept; the 1-D positions on all three rows when SGLang computed none; rows past the batch zeroed): stage_draft_graph_mrope.
• After every Flash draft decode forward, advance the caller's mRoPE together with the 1-D positions: advance_draft_decode_mrope, called from a draft_forward wrapper, because the eager and graph runners hand the model a copy of the batch. During capture the add is recorded in the graph. This also applies to eager drafting.
• Verified in serving: after a decode replay the buffer holds the staged row + 3 (4 steps), and an in-graph recorder shows the draft RoPE at staged+1, +2, +3 on an image request (offset of -992 kept).

Flag parsing → no flags. ATOM_SGLANG_QWEN4_DRAFT_CG / ATOM_SGLANG_QWEN4_DRAFT_EXTEND_CG are removed (94cef79ce): with the mRoPE fix the graphs draft as well as eager drafting, and eager drafting is host-launch bound (slower than no MTP), so both graphs are always captured for the Flash drafter. The out-of-graph QSA hook now keys on the draft batch alone, so there is no second gate to disagree with the first.

Tests. tests/plugin/test_sglang_qwen4_exp_mtp.py covers staging (text broadcast, multimodal offsets, padding), the per-step advance, and that a draft decode batch never fills the target's graph metadata.

Results. 2x MI308X, rocm/atom-dev:sglang-latest e021a3935c1c, TP2/EP2, EAGLE 4 steps, concurrency 1.

Release-note launch (ctx 8192); text is bench_serving random 1024/2048:

draft	GSM8K 400, before / after	text accept, before / after	text TPOT, before / after
eager	97.75% / 97.00%	3.97 / 4.00	12.23 / 12.56 ms
both graphs	97.75% / 97.25%	3.81 / 3.96	5.37 / 5.05 ms

Long-context launch (ctx 132096, the launch below), after the fix:

workload	draft	accept	TPOT
bench_serving image dataset (720p image + 256 tokens, OSL 512), two runs	both graphs	2.73 / 2.76	6.70 / 6.47 ms
	eager	2.74 / 2.73	15.92 / 15.75 ms
Real multi-turn multimodal agent request (65 messages, 18 tools, 30 images, 65,666 prompt tokens; greedy, about 3-4k output tokens; 3 runs)	both graphs	3.42	5.59 ms
	eager	3.42	12.51 ms
	no MTP, TP2/EP1 (#2418 recipe)	-	7.67 ms
	no MTP, TP2/EP2	-	10.06 ms

On the agent request, MTP with both graphs cuts TPOT by 27% against the best no-MTP launch (TP2/EP1); TTFT goes up by about 0.55 s (5.09 s to 5.63 s). Image requests accept the same with and without the draft graphs.

Accept length on few-prompt runs is noisy: two identical greedy runs diverge within the first tokens, and outputs that fall into repetition accept far more, so 4-8 prompt runs swing by about ±0.5. Removing the flags was checked with both builds served side by side (GPUs 2-3 and 4-5, same host load), 16 prompts, OSL 512, concurrency 1, release-note launch, two rounds:

build	text 1024 accept / TPOT	image 720p + 256 accept / TPOT
no flags (94cef79ce)	3.13, 3.10 / 4.74, 4.68 ms	3.07, 3.04 / 6.55, 6.91 ms
both flags set (399206b7e)	3.34, 3.10 / 4.74, 4.61 ms	2.97, 2.97 / 6.80, 7.00 ms

Reproduction

These are the commands behind the "Not a recovered regression" table (image rocm/atom-dev:sglang-latest e021a3935c1c, SGLang 0.5.20 94602c9c2b). The step-sweep and 29-case numbers used the same launches on rocm/atom-dev:sglang-v0.5.20-nightly_20261012, with this branch's two commits on ATOM a0527c9db (no #2637 needed there).

Tree and container

git clone https://github.com/zovonoir/ATOM.git ATOM-mtp && cd ATOM-mtp
git checkout qwen38-flash-mtp-base1019        # main + #2637 + this PR
cd ..

MODEL_DIR=/models/Qwen3.8-Flash-Next-PTPC-FP8
IMAGE=rocm/atom-dev:sglang-latest              # measured: e021a3935c1c

docker run -d --name qwen38-mtp \
  --device=/dev/kfd --device=/dev/dri --group-add video \
  --ipc=host --shm-size=32g --cap-add=SYS_PTRACE \
  --security-opt seccomp=unconfined --network=host \
  -e HIP_VISIBLE_DEVICES=0,1 \
  -v ${MODEL_DIR}:/model:ro \
  -v $PWD/ATOM-mtp:/atom \
  ${IMAGE} sleep infinity

Every launch below puts the branch ahead of the image's ATOM with PYTHONPATH=/app/sglang/python:/atom. Confirm with python -c "import atom; print(atom.__file__)" -> /atom/atom/__init__.py. Run one server at a time; first start takes about 3.5 minutes.

Baseline: no MTP, TP2/EP1 (#2418 recipe)

docker exec -d qwen38-mtp bash -lc '
  export PYTHONPATH=/app/sglang/python:/atom
  export SGLANG_PLUGINS=atom_sglang
  export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
  python -m sglang.launch_server \
    --model-path /model --host 0.0.0.0 --port 30080 \
    --tp-size 2 --ep-size 1 \
    --kv-cache-dtype bf16 --page-size 64 \
    --context-length 132096 --max-running-requests 64 \
    --mem-fraction-static 0.88 \
    --disable-radix-cache --trust-remote-code > /tmp/server_base.log 2>&1
'

Expected log: KV Cache is allocated. dtype: torch.bfloat16, #tokens: 3578944.

MTP with draft CUDA graphs: TP2/EP2, EAGLE 4 steps, long context

docker exec -d qwen38-mtp bash -lc '
  export PYTHONPATH=/app/sglang/python:/atom
  export SGLANG_PLUGINS=atom_sglang
  export SGLANG_EXTERNAL_MODEL_PACKAGE=atom.plugin.sglang.models
  export SGLANG_EXTERNAL_MM_PROCESSOR_PACKAGE=atom.plugin.sglang.models
  export SGLANG_USE_AITER=1
  export SGLANG_AITER_KV_CACHE_LAYOUT=nhd
  export AITER_MOE_PADDING_SIZE=128
  python -m sglang.launch_server \
    --model-path /model --trust-remote-code \
    --host 0.0.0.0 --port 30080 \
    --tp 2 --ep-size 2 --attention-backend aiter \
    --kv-cache-dtype bf16 --page-size 64 \
    --context-length 132096 --max-running-requests 64 \
    --mem-fraction-static 0.88 --disable-radix-cache \
    --cuda-graph-backend-decode full --cuda-graph-max-bs-decode 64 \
    --speculative-algorithm EAGLE \
    --speculative-num-steps 4 --speculative-eagle-topk 1 \
    --sampling-defaults openai > /tmp/server_mtp.log 2>&1
'

Expected log: Capture target verify CUDA graph end, Capture draft decode CUDA graph end, and no Failed to enable draft-extend cuda graph in ATOM plugin. The plugin-built draft-extend runner is not counted in SGLang's startup timing line, so draft_extend=0.00 there is expected. KV pool: 1,698,944 tokens at 4 steps (2,229,568 at 2 steps).

• Eager drafting (the behaviour before this PR) has no switch any more; --disable-cuda-graph disables every graph, the target's included.
• Draft-step sweep: same command with --context-length 8192 --max-running-requests 8 --mem-fraction-static 0.70 --cuda-graph-max-bs-decode 8 (the release-note launch) and --speculative-num-steps {2..6}.
• MTP requires --ep-size 2: with EP1 the BF16 drafter MoE (intermediate 320 per rank) has no CK kernel and the first request fails.

Benchmark (the #2418 harness)

for c in "1024 1" "1024 8" "1024 32" "4096 1" "4096 16" "6144 4" \
         "30720 1" "30720 8" "63488 4" "63488 16" "129024 1"; do
  set -- $c
  docker exec qwen38-mtp python -m sglang.bench_serving \
    --backend sglang --host 127.0.0.1 --port 30080 \
    --model /model --tokenizer /model --dataset-name random \
    --random-input-len $1 --random-output-len 2048 --random-range-ratio 1.0 \
    --num-prompts $(( $2 * 2 )) --max-concurrency $2 --request-rate inf \
    --warmup-requests 1 --disable-tqdm
done

MTP acceptance is the mean of the server log's accept len: values printed during each case.

Accuracy

docker exec qwen38-mtp bash -lc "cd /tmp && python -m sglang.test.run_eval \
  --host 127.0.0.1 --port 30080 --eval-name gsm8k --num-examples 400 --num-threads 64 \
  --api chat --chat-template-kwargs '{\"enable_thinking\": false}'"

Do not use sglang.test.few_shot_gsm8k for this model: it sends prompts without the chat template, and the model answers <|im_end|> immediately.

Not covered

• MTP still requires EP2 (the BF16 drafter MoE at TP2/EP1 has no CK kernel for intermediate 320); EP1 MTP could save another ~2.3 ms per target step.
• Decode graph bs > 64, other parallel sizes, gfx950.
• Unit tests for the draft graph metadata path.

Test plan

• Startup, smoke, GSM8K 400 and multimodal with each graph combination (before the flags were removed) and with the default
• Draft-step sweep 2-6, release-note launch
• 29-case perf(qwen4_exp): MI308X optimizations for Qwen3.8-Flash-Next on native ATOM and SGLang-ATOM #2418 sweep (2 and 4 steps), nightly image
• 11-case baseline vs MTP on current main + fix(sglang-atom): fill the per-step GDN prefill indices for Qwen3.8-Flash-Next #2637, latest image
• black / ruff clean on changed files
• Unit tests for draft mRoPE staging / advance and the draft hook isolation
• Unit tests for the draft graph QSA metadata path
• SGLang CI

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2645 --add-label <label>

[Copilot]

🟡 Changes recommended

Draft graph mRoPE inputs are not refreshed correctly, and inconsistent flag parsing bypasses target metadata isolation.

2 open issues

• Stage and advance mRoPE positions during draft replay · New
• Align draft graph flag parsing with _env_flag values · New

What changed in this PR

Adds opt-in draft CUDA graphs for Qwen3.8-Flash-Next MTP in the SGLang plugin to reduce host-launch overhead.

Changes:

• Enables draft decode and extend graphs with separate draft QSA metadata.
• Forwards additional SGLang draft-extend replay arguments.
• Includes #2637’s GDN prefill metadata fix and regression test.

File	Description
tests/​plugin/​test_sglang_gdn_forward_context.py	Tests required per-step prefill indices.
atom/​plugin/​sglang/​qwen4_exp_bridge.py	Separates draft and target graph metadata.
atom/​plugin/​sglang/​attention_backend/​attention_gdn.py	Populates GDN prefill state indices.
atom/​plugin/​register.py	Gates draft graphs and adapts replay arguments.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Model validation, draft buffer capacity, and shared-pool synchronization issues remain unresolved.

3 open issues

• Publish the scheduler fence only after draft graph shared reads · New
• Validate Qwen4 configuration before returning QSA marker · New
• Size draft QSA buffers for the CUDA graph capture batch · New

2 resolved since last review

• Stage and advance mRoPE positions during draft replay
• Align draft graph flag parsing with _env_flag values

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Legacy SGLang patch ordering overwrites the mRoPE wrapper, leaving later draft steps with stale positions.

3 open issues

• Publish the scheduler fence only after draft graph shared reads
• Size draft QSA buffers for the CUDA graph capture batch
• Validate Qwen4 configuration before returning QSA marker

1 previously missed issue

In code that hasn't changed since last review

mRoPE wrapper is lost when argmax patch replaces draft_forward

atom/​plugin/​register.py:725

[verified] On the legacy SGLang runner API without execute, register_ops_to_sglang() installs this wrapper and then calls _patch_sglang_eagle_v2_draft_argmax(), which replaces EagleDraftWorker.draft_forward outright. The mRoPE advancement is lost, while _atom_qwen4_draft_mrope_patched prevents reinstalling it. Flash draft steps therefore keep using step-0 mRoPE in eager execution and graph capture. Install the argmax replacement before capturing original_draft_forward so this wrapper composes with it.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Legacy SGLang patch ordering overwrites the mRoPE wrapper, leaving later draft steps with stale positions.

4 open issues

• Preserve argmax patching when wrapping legacy draft_forward · New
• Publish the scheduler fence only after draft graph shared reads
• Size draft QSA buffers for the CUDA graph capture batch
• Validate Qwen4 configuration before returning QSA marker

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Shared-buffer GPU replay needs final human validation, with draft attention graph-path tests and SGLang integration CI still pending.

No open issues 🎉

4 resolved since last review

• Preserve argmax patching when wrapping legacy draft_forward
• Publish the scheduler fence only after draft graph shared reads
• Size draft QSA buffers for the CUDA graph capture batch
• Validate Qwen4 configuration before returning QSA marker

_{Review effort: Balanced}

[zovonoir]

@valarLip — friendly ping for a review. This has been open since Oct 20; the last automated review round reported no open issues and only flagged the shared-buffer GPU replay for human validation.

State

• Both draft CUDA graphs (decode + extend) are now captured by default for the Flash drafter. No flags and no SGLang change: atom/plugin/register.py and atom/plugin/sglang/qwen4_exp_bridge.py only.
• Fixes the plugin's draft-extend replay wrapper, which took (self, forward_batch) while SGLang 0.5.20 calls execute(forward_batch, select_index).

Verification (2x MI308X, Qwen3.8-Flash-Next-PTPC-FP8)

• MTP vs the perf(qwen4_exp): MI308X optimizations for Qwen3.8-Flash-Next on native ATOM and SGLang-ATOM #2418 no-MTP baseline on current main: output throughput geomean +22.5%, TPOT -23.5%, ahead in 9/11 cases.
• Draft step cost drops from ~8 ms eager to ~0.7 ms graphed, which is what makes MTP a net win.
• GSM8K 400 97.0-97.75% across configurations; multimodal (1080p + ~3.8K text) correct at concurrency 1-16.

Deliberately left out of scope (documented under "Not covered" in the PR body)

• ROCm + forced SGLANG_ENABLE_WAR_BARRIER=1 + the Flash draft-extend graph. The WAR barrier is off by default on ROCm (is_cuda() or SGLANG_ENABLE_WAR_BARRIER), so this combination needs the event published after the draft-extend replay; the fix is a follow-up.
• MTP still requires --ep-size 2 (the BF16 drafter MoE at TP2/EP1 has no CK kernel for intermediate 320).

Ready for review.

[ZLkanyo009]

LGTM, approving. Keeping this plugin-only makes sense, since SGLang's draft runners already carry mrope_positions. The stale-mRoPE issue is now covered in both modes: the replay wrappers stage the live rows (multimodal offsets included), and the draft_forward wrapper advances them per step for eager drafting. _DRAFT_QSA keeps the target-verify graph metadata untouched, and the *args, **kwargs passthrough also fixes the execute(forward_batch, select_index) TypeError for the other draft-extend users. Tracking the WAR-barrier caveat as a follow-up is fine, since ROCm doesn't enable that barrier by default.

One optional thought, not blocking. ATOM_SGLANG_V4_DISABLE_DRAFT_EXTEND_CG still turns off the Flash draft-extend graph, but the Flash draft decode graph can now only be turned off with --disable-cuda-graph, which drops the target graphs too. If you ever need to bisect a draft-graph issue in production, letting ATOM_SGLANG_V4_DISABLE_DRAFT_CG also cover the Flash drafter would give you a cheap escape hatch.