[ZLkanyo009]

Motivation

Enable Qwen3.8-Flash-Next MTP in SGLang-ATOM.

The target and the same-checkpoint 1-layer QSA draft both run Native ATOM
Qwen4Exp / Qwen4ExpMTP. The SGLang plugin only adapts scheduling, QSA page
tables, PLE state, and hyper-connection hidden storage. MTP is scheduled as
SGLang EAGLE with --speculative-num-steps 2 and --speculative-eagle-topk 1
(3 draft tokens). Target verify uses CUDA graphs. Draft stays eager because
Flash mRoPE positions are [3, N], matching Native ATOM.

Command

MI308X, TP2 + EP2, BF16 KV, page size 64, no prefix / radix cache.
Checkpoint: /models/Qwen3.8-Flash-Next-PTPC-FP8.

SGLang-ATOM with MTP:

export ATOM_PATH=/workspace/qwen/ATOM_main_pr3
export SGLANG_PATH=/workspace/qwen/sglang-v517
export PYTHONPATH="${SGLANG_PATH}/python:${ATOM_PATH}"
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_LOG_LEVEL=WARNING
export AITER_MOE_PADDING_SIZE=128
export CUDA_VISIBLE_DEVICES=0,1
export HIP_VISIBLE_DEVICES=0,1

python3 -m sglang.launch_server \
  --model-path /models/Qwen3.8-Flash-Next-PTPC-FP8 \
  --host 127.0.0.1 --port 8000 --trust-remote-code \
  --tp 2 --ep-size 2 --attention-backend aiter \
  --kv-cache-dtype bf16 --page-size 64 \
  --mem-fraction-static 0.70 \
  --context-length 65536 --max-running-requests 32 \
  --disable-radix-cache \
  --cuda-graph-backend-decode full --cuda-graph-max-bs-decode 32 \
  --chunked-prefill-size 16384 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 2 --speculative-eagle-topk 1 \
  --reasoning-parser qwen3 \
  --sampling-defaults openai

Omit the three --speculative-* flags to serve without MTP.

Native ATOM with MTP, same shape:

export ATOM_PATH=/workspace/qwen/ATOM_main_pr3
export PYTHONPATH="${ATOM_PATH}"
export AITER_LOG_LEVEL=WARNING
export CUDA_VISIBLE_DEVICES=0,1
export HIP_VISIBLE_DEVICES=0,1

python3 -m atom.entrypoints.openai_server \
  --model /models/Qwen3.8-Flash-Next-PTPC-FP8 \
  --host 127.0.0.1 --server-port 8000 --trust-remote-code \
  -tp 2 --enable-expert-parallel \
  --kv-cache-dtype bf16 --block-size 64 \
  --max-model-len 65536 --max-num-seqs 32 \
  --max-num-batched-tokens 16384 \
  --gpu-memory-utilization 0.70 \
  --cudagraph-mode FULL \
  --cudagraph-capture-sizes "[1,2,4,8,16,32]" \
  --no-enable_prefix_caching \
  --method mtp --num-speculative-tokens 2

Large-output benchmark (ISL 1024, OSL 4096, concurrency 4, 8 prompts, ignore EOS):

PYTHONPATH=/workspace/qwen/ATOM_main_pr3 python3 -m atom.benchmarks.benchmark_serving \
  --backend sglang --host 127.0.0.1 --port 8000 \
  --model /models/Qwen3.8-Flash-Next-PTPC-FP8 \
  --dataset-name random \
  --random-input-len 1024 --random-output-len 4096 --random-range-ratio 1.0 \
  --num-prompts 8 --max-concurrency 4 \
  --request-rate inf --ignore-eos --seed 0 --num-warmups 0

GSM8K, 5-shot, 1319 items, max_gen_toks=4096, stop on <|im_end|>:

lm_eval --model local-chat-completions \
  --model_args model=/models/Qwen3.8-Flash-Next-PTPC-FP8,base_url=http://127.0.0.1:8000/v1/chat/completions,num_concurrent=8,tokenized_requests=False,trust_remote_code=True \
  --tasks gsm8k --num_fewshot 5 \
  --gen_kwargs max_gen_toks=4096,until='<|im_end|>' \
  --apply_chat_template --confirm_run_unsafe_code \
  --output_path ./gsm8k

RealWorldQA, 765 items, rule score on the answer text (the Qwen3.5-27B judge was not run):

python realworldqa_eval.py infer \
  --model /models/Qwen3.8-Flash-Next-PTPC-FP8 \
  --base-url http://127.0.0.1:8000/v1 \
  --data /workspace/qwen/sglang-v517/eval/data/RealWorldQA.tsv \
  --concurrency 32 \
  --output predictions.jsonl

Performance

Same server shape as above. Random ISL 1024 / OSL 4096, concurrency 4, 8 prompts, ignore EOS. MTP off uses the same benchmark with the speculative flags omitted. Output-token throughput and TPOT are the decode comparison.

	Output tok/s	Mean TPOT	Median TPOT	Mean ITL
Native ATOM, MTP off	249.48	15.84 ms	15.84 ms	15.84 ms
Native ATOM, MTP on	376.79	8.88 ms	7.76 ms	20.05 ms
SGLang-ATOM, MTP off	287.18	13.74 ms	13.74 ms	13.74 ms
SGLang-ATOM, MTP on	328.12	10.70 ms	10.11 ms	25.03 ms

MTP helps on this long-output shape. With MTP off, one decode step emits one token, so ITL equals TPOT. With MTP on, a step emits more than one accepted token, so ITL goes up while time per output token goes down.

• Native ATOM: TPOT 15.84 ms to 8.88 ms, output throughput 249.48 to 376.79 tok/s (about 1.5x).
• SGLang-ATOM: TPOT 13.74 ms to 10.70 ms, output throughput 287.18 to 328.12 tok/s (about 1.14x).

SGLang-ATOM with MTP is still behind Native ATOM with MTP (10.70 ms vs 8.88 ms, about 13% less output throughput). Verify GPU time matches (about 18.2 ms). The remaining gap is the eager 1-token draft step: its kernels match Native in length, but they are launched on an idle GPU, so host launch latency shows up between kernels. The 3-token draft extend is already packed. Draft CUDA graphs stay off.

ACC

GSM8K flexible-extract / strict-match, 5-shot, 1319 items, max_gen_toks=4096, stop <|im_end|>, temperature 0:

	flexible-extract	strict-match
Native ATOM, MTP off	0.9788	0.9757
Native ATOM, MTP on	0.9780	0.9750
SGLang-ATOM, MTP off	0.9780	0.9735
SGLang-ATOM, MTP on	0.9757	0.9727

RealWorldQA rule score, 765 items. This counts extracted answers only. The Qwen3.5-27B judge pass was not run.

	correct	accuracy
Native ATOM, MTP off	591 / 765	77.25%
Native ATOM, MTP on	594 / 765	77.65%
SGLang-ATOM, MTP off	599 / 765	78.30%
SGLang-ATOM, MTP on	599 / 765	78.30%

MTP acceptance

The MTP-on row above (328.12 tok/s, TPOT 10.70 ms) is before the draft QSA seq_lens fix. That build kept GSM8K answer accuracy, but GSM8K draft acceptance dropped from about 79% to 68%.

QSA requires logical_pos < seq_lens[req]. Target decode sits at seq_lens - 1, so no lift is needed. The second draft step is also 1 token per request (ForwardMode.DECODE), but its position is seq_lens. The packed draft path in _fill_qsa_like_native only lifted when tokens_per_req > 1, so that step left the current token outside the indexer context. The first draft extend (width 3) was lifted. The eager path used by the 79% tree already lifts on any draft forward. ReplaySSM stayed enabled; it was not the change that recovered acceptance.

Same GSM8K protocol as the table above, plus temperature=0, top_k=1, top_p=1. Acceptance is accepted draft tokens / drafted tokens (the bonus token is not in the denominator). accept len includes the bonus.

	flexible-extract	strict-match	accept rate	accept len
SGLang-ATOM before the lift fix, 1319 items	0.9773	0.9757	0.680	2.36
SGLang-ATOM without the packed-draft path, 1319 items	0.9773	0.9719	0.792	2.58
SGLang-ATOM after the lift fix, first 32 items	0.9688	0.9688	0.812	2.63

The 32-item run is only an acceptance check. It is not a replacement for the 1319-item score.

Remeasured the long-output benchmark after the fix, same shape as the MTP-on row (ISL 1024 / OSL 4096, concurrency 4, 8 prompts, ignore EOS, context 65536, decode graph batch 32):

	Output tok/s	Mean TPOT	Median TPOT	Mean ITL	accept rate	accept len
SGLang-ATOM, MTP on, before	328.12	10.70 ms	10.11 ms	25.03 ms	—	—
SGLang-ATOM, MTP on, after	399.69	9.31 ms	9.08 ms	24.89 ms	0.858	2.71

Mean ITL is unchanged, so the extra lift kernel did not lengthen the step. TPOT improved because each step accepted more draft tokens.

[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 2385 --add-label <label>