[zovonoir]

Summary

Qwen3.8-Flash-Next carries a four-stream hyper-connection (HC) residual, [T, 4 x 2560]. Every sub-layer is wrapped by an HC: a grouped norm, two low-rank GEMMs ([T,10240]x[10240,324] and [T,320]x[320,10240]) and a gated mean, and the sub-layer output is written back into all four streams. The HC weights are replicated under TP, so every rank repeats this per-token work on every token. On a TP4 prefill profile (4096 x 32, MTP on) the HC kernels took about 30% of GPU time.

This PR keeps the streams token-sharded across the TP ranks (sequence parallelism) for large eager forwards:

Before: plain TP, every rank holds and processes all T rows

flowchart TB
    S0["rank 0<br/>streams [T,4S]"]:::dup
    S1["rank 1<br/>streams [T,4S]"]:::dup
    S2["rank 2<br/>streams [T,4S]"]:::dup
    S3["rank 3<br/>streams [T,4S]"]:::dup

    M0["mix: T rows"]:::dup
    M1["mix: T rows"]:::dup
    M2["mix: T rows"]:::dup
    M3["mix: T rows"]:::dup

    X0["x [T,S]"]:::dup
    X1["x [T,S]"]:::dup
    X2["x [T,S]"]:::dup
    X3["x [T,S]"]:::dup

    L0["sub-layer<br/>heads / experts shard 0"]:::shard
    L1["sub-layer<br/>heads / experts shard 1"]:::shard
    L2["sub-layer<br/>heads / experts shard 2"]:::shard
    L3["sub-layer<br/>heads / experts shard 3"]:::shard

    P0["partial 0 [T,S]"]:::shard
    P1["partial 1 [T,S]"]:::shard
    P2["partial 2 [T,S]"]:::shard
    P3["partial 3 [T,S]"]:::shard

    AR{{"all-reduce [T,S]<br/>sum of all partials, full result to every rank"}}:::comm

    C0["combine: T rows"]:::dup
    C1["combine: T rows"]:::dup
    C2["combine: T rows"]:::dup
    C3["combine: T rows"]:::dup

    N0["streams [T,4S]"]:::dup
    N1["streams [T,4S]"]:::dup
    N2["streams [T,4S]"]:::dup
    N3["streams [T,4S]"]:::dup

    S0 --> M0 --> X0 --> L0 --> P0 --> AR
    S1 --> M1 --> X1 --> L1 --> P1 --> AR
    S2 --> M2 --> X2 --> L2 --> P2 --> AR
    S3 --> M3 --> X3 --> L3 --> P3 --> AR

    AR --> C0 --> N0
    AR --> C1 --> N1
    AR --> C2 --> N2
    AR --> C3 --> N3

    classDef dup fill:#f8d7da,stroke:#c0392b,color:#000
    classDef shard fill:#d4edda,stroke:#27ae60,color:#000
    classDef comm fill:#d6e4f5,stroke:#2e6da4,color:#000

Loading

After (this PR): HC on a T/tp row shard per rank, the sub-layer still sees all T rows

flowchart TB
    S0["rank 0<br/>streams shard 0 [T/4,4S]"]:::shard
    S1["rank 1<br/>streams shard 1 [T/4,4S]"]:::shard
    S2["rank 2<br/>streams shard 2 [T/4,4S]"]:::shard
    S3["rank 3<br/>streams shard 3 [T/4,4S]"]:::shard

    M0["mix: T/4 rows"]:::shard
    M1["mix: T/4 rows"]:::shard
    M2["mix: T/4 rows"]:::shard
    M3["mix: T/4 rows"]:::shard

    Y0["x shard 0 [T/4,S]"]:::shard
    Y1["x shard 1 [T/4,S]"]:::shard
    Y2["x shard 2 [T/4,S]"]:::shard
    Y3["x shard 3 [T/4,S]"]:::shard

    AG{{"all-gather [T,S]<br/>each rank contributes its T/4 rows"}}:::comm

    X0["x [T,S]"]:::full
    X1["x [T,S]"]:::full
    X2["x [T,S]"]:::full
    X3["x [T,S]"]:::full

    L0["sub-layer<br/>heads / experts shard 0"]:::shard
    L1["sub-layer<br/>heads / experts shard 1"]:::shard
    L2["sub-layer<br/>heads / experts shard 2"]:::shard
    L3["sub-layer<br/>heads / experts shard 3"]:::shard

    P0["partial 0 [T,S]"]:::shard
    P1["partial 1 [T,S]"]:::shard
    P2["partial 2 [T,S]"]:::shard
    P3["partial 3 [T,S]"]:::shard

    RS{{"reduce-scatter [T,S]<br/>sum of all partials, rank k keeps row shard k"}}:::comm

    R0["sum shard 0 [T/4,S]"]:::shard
    R1["sum shard 1 [T/4,S]"]:::shard
    R2["sum shard 2 [T/4,S]"]:::shard
    R3["sum shard 3 [T/4,S]"]:::shard

    C0["combine: T/4 rows"]:::shard
    C1["combine: T/4 rows"]:::shard
    C2["combine: T/4 rows"]:::shard
    C3["combine: T/4 rows"]:::shard

    N0["streams shard 0 [T/4,4S]"]:::shard
    N1["streams shard 1 [T/4,4S]"]:::shard
    N2["streams shard 2 [T/4,4S]"]:::shard
    N3["streams shard 3 [T/4,4S]"]:::shard

    S0 --> M0 --> Y0 --> AG
    S1 --> M1 --> Y1 --> AG
    S2 --> M2 --> Y2 --> AG
    S3 --> M3 --> Y3 --> AG

    AG --> X0 --> L0 --> P0 --> RS
    AG --> X1 --> L1 --> P1 --> RS
    AG --> X2 --> L2 --> P2 --> RS
    AG --> X3 --> L3 --> P3 --> RS

    RS --> R0 --> C0 --> N0
    RS --> R1 --> C1 --> N1
    RS --> R2 --> C2 --> N2
    RS --> R3 --> C3 --> N3

    classDef dup fill:#f8d7da,stroke:#c0392b,color:#000
    classDef shard fill:#d4edda,stroke:#27ae60,color:#000
    classDef comm fill:#d6e4f5,stroke:#2e6da4,color:#000
    classDef full fill:#ffffff,stroke:#888888,color:#000

Loading

S = hidden_size (2560), 4S = hc_count x hidden_size (10240), shown for TP = 4.

• Red: identical data / identical computation replicated on every rank (redundant).
• Green: each rank holds or computes a different part (heads, experts, or a T/tp row shard).
• Blue: collectives; every rank sends to and receives from them.
• White: the full x [T,S] the sub-layer needs on every rank, obtained by all-gather instead of recomputation.
• combine also reads the pre-mix streams and the inject logits of the same rows (local, omitted from the diagrams); in the code it is fused into the next mix (pending).

All-reduce = reduce-scatter + all-gather, so communication bytes are unchanged and each rank does 1/tp of the HC work.

File	Change
atom/models/qwen4_exp.py	_SequenceParallel (shard / gather / reduce-scatter, rows padded to a multiple of TP). Qwen4ExpDecoderLayer.forward(sp=) runs both HCs on the local shard, all-gathers the sub-layer input and reduce-scatters its partial sum. Qwen4ExpSparseMoeBlock.forward(reduce_results=False) returns the partial sum for the reduce-scatter. Qwen4ExpModel.forward shards after the last PLE layer and gathers the final mixer output, or the HC state when MTP is on.
atom/utils/envs.py	ATOM_QWEN4_EXP_SP_MIN_TOKENS (default 1024; 0 disables).

When it applies:

• Only when TP > 1, the forward has at least ATOM_QWEN4_EXP_SP_MIN_TOKENS tokens, and the stream is not capturing, so no captured graph is affected. With the follow-up in perf(qwen4_exp): Qwen3.8-Flash-Next TP prefill follow-ups: drafter FlyDSL QSA, row-sharded indexer selection, sequence-parallel drafter #2847 (239e24ab9) the forward must also have more than HC_QUANT_ROWS_MAX (1024) tokens, whatever the env value; see the next bullet but one. The gate checks the token count, not the execution phase, so any eager target forward that large is sharded. In practice that is prefill / extend.
• Decode and MTP target verification are sharded only when a batch runs eagerly (outside the CUDA graphs) and reaches the threshold. At the default 1024 that means 1024 or more decode requests, or 342 or more verify requests with EAGLE 2 steps (3 tokens per request). The deployment measured here (at most 64 requests, all captured) never reaches it. A batch that does reach it gets the same speed-up, and its numerics differ from the unsharded path only as described under Numerics below.
• Sharding starts after the PLE layer (layer 1), because its causal short convolution needs every token of the request.
• Padded rows are zero streams. They stay zero through the HC norm and are dropped on gather.
• Sharded layers do not use feat(qwen4-exp): online PTPC FP8 for the GDN in_proj, with fused activation quant #2650's mix_fused_quant. The FP8 attention-input projection quantizes the gathered BF16 rows itself, which the unsharded path already does above HC_QUANT_ROWS_MAX (1024) rows. Known gap in this PR, fixed in perf(qwen4_exp): Qwen3.8-Flash-Next TP prefill follow-ups: drafter FlyDSL QSA, row-sharded indexer selection, sequence-parallel drafter #2847 (239e24ab9): this PR's gate starts at 1024 tokens, while the unsharded path keeps the fused quant up to and including 1024. So an eager TP > 1 forward of exactly 1024 tokens (or any count between a lower env threshold and 1024) quantizes the projection input with aiter's per_token_quant_hip when sharded but with hc_gated_mean_quant when not. The FP32 scale can differ by about 1 ulp and an occasional FP8 element can round differently. Both paths are existing production paths (T > 1024 already uses the first). The fix shards only above HC_QUANT_ROWS_MAX, so sharded and unsharded forwards always quantize the same way. Since perf(qwen4_exp): sequence-parallel HC residual streams for Qwen3.8-Flash-Next TP prefill #2834 is merged directly before perf(qwen4_exp): Qwen3.8-Flash-Next TP prefill follow-ups: drafter FlyDSL QSA, row-sharded indexer selection, sequence-parallel drafter #2847, the fix lands in perf(qwen4_exp): Qwen3.8-Flash-Next TP prefill follow-ups: drafter FlyDSL QSA, row-sharded indexer selection, sequence-parallel drafter #2847 to keep this PR's reviewed code unchanged.
• The MTP drafter is not changed. It still receives the gathered [T, 4, 2560] HC state.

Numerics: mathematically identical, with rounding-level differences only. The sub-layer sum is a BF16 reduce-scatter instead of a BF16 all-reduce, so the summation order may differ. The HC low-rank GEMMs run on T/tp rows instead of T, so the GEMM backend may pick a different solution for the smaller M. At exactly 1024 tokens in this PR, the FP8 activation-quant path also differs (see above; fixed in #2847).

Relation to #2653

Independent of #2653: they share no code and the two can merge in either order. The only shared file is atom/utils/envs.py, and the new flag sits in the Qwen QSA block, away from ATOM_QSA_FLYDSL_ATTN. Checked locally: git merge-tree of this branch with #2653's head dc3a19bf5 is clean, and rebasing this commit onto dc3a19bf5 is clean.

Follow-up, not in this PR: split the replicated QSA indexer scoring / top-k across the same row shards, and shard the MTP drafter's prefill extend in the same way. Both build on _SequenceParallel, and the indexer part touches qsa_attention.py, which #2653 also changes. I will open them after this PR and #2653 land.

Results

4x MI308X (TP4 / EP4), Qwen3.8-Flash-Next-PTPC-FP8, MTP on (EAGLE 2 steps), radix cache on and flushed before each measured run; server-side #cached-token = 0 for every measured run.

Input throughput, before / after. The deployment image's tree: ATOM 6293e75b5 (main a0527c9db plus the Qwen3.8-Flash SGLang PRs that have since merged and an earlier revision of #2653, with ATOM_QSA_FLYDSL_ATTN=1), aiter 04a254638 with ROCm/aiter#6576, SGLang 94602c9c2b. Before = that tree. After = the same tree plus this change (4ac16ec96, which is this PR before the rebase onto main). Workload: sglang.bench_serving random, range 1.0, OSL 1, ~800K input tokens per case. Each value is the mean of runs r1 and r2 after a discarded warmup run; where the two runs differed by more than 3% (marked *), r3 was added and r1-r3 are averaged.

ISL x CON	before tok/s	after tok/s	gain	mean TTFT ms before / after
1024 x 64	22,953	27,450*	+19.6%	2,701 / 2,254
4096 x 32	23,169	28,316	+22.2%	5,160 / 4,198
16384 x 16	22,222	26,725*	+20.3%	10,223 / 8,491
30720 x 8	21,521	26,266	+22.0%	10,023 / 8,192
63488 x 4	20,818	24,876	+19.5%	11,032 / 9,222
129024 x 2	19,154	22,954	+19.8%	12,684 / 10,570
geomean	21,595	26,037	+20.6%

GSM8K 400 (run_eval, chat, thinking off): before 0.970, after 0.970.

SP accuracy A/B with every prefill sharded. The run_eval GSM8K above only partly covers SP. With radix cache on, the shared few-shot prefix is cached, so only the first wave of 64 questions is prefilled in batches above the threshold, about 55% of prefill tokens; the other 336 questions prefill around 180 new tokens each and stay replicated. So this A/B uses benchmark/gsm8k/bench_sglang.py --num-shots 10 (chat template, thinking off, temperature 0) with --disable-radix-cache. Every question then prefills about 2.1K new tokens, and 99.97% of prefill tokens ran in batches above 1024 tokens. This used the same TP4 / EP4 MTP image tree, with this PR plus #2847's sharded drafter and indexer, so SP on vs off toggles all three. Two runs per setting:

GSM8K 10-shot, all 1319 test questions, greedy	run 1	run 2	mean
SP on (default)	0.964	0.970	0.967
SP off (ATOM_QWEN4_EXP_SP_MIN_TOKENS=0)	0.967	0.969	0.968

Per-question McNemar tests across all six pairs: SP on vs SP off differ on 23-28 answers with p = 0.17-1.0. The two SP-on runs differ on 28 answers (p = 0.13) and the two SP-off runs on 22 (p = 0.80): the run-to-run variation from MTP and batch composition is as large as the SP on/off difference. No accuracy effect from SP is detectable.

TP2 / EP2, before / after. Same image tree, same MTP / radix launch and workload, with --tp 2 --ep-size 2. After = 4ac16ec96 plus this PR's follow-up fix 6c77e2a42. The two servers ran at the same time on separate GPU pairs: before on GPU 4-5, after on GPU 6-7. Each value is the mean of runs r1 and r2 after a discarded warmup run; runs differed by 1.7% at most.

ISL x CON	before tok/s	after tok/s	gain	mean TTFT ms before / after
1024 x 64	16,277	17,500	+7.5%	3,804 / 3,547
4096 x 32	16,345	17,638	+7.9%	7,259 / 6,712
16384 x 16	15,827	17,102	+8.1%	14,363 / 13,297
30720 x 8	15,623	16,863	+7.9%	13,510 / 13,048
63488 x 4	15,113	16,315	+8.0%	15,210 / 14,080
129024 x 2	14,278	15,373	+7.7%	17,032 / 15,817
geomean	15,561	16,780	+7.8%

GSM8K 400: before 0.975, after 0.9725. The gain is smaller than at TP4 because only half of the HC work is replicated at TP2 (each rank saves 1/2 of the HC work, against 3/4 at TP4).

This branch on current main. ATOM main ce46c7c8d plus this PR, aiter main b33beb00e. ATOM_QWEN4_EXP_SP_MIN_TOKENS=0 turns the change off. Each value is the mean of 2 runs after a warmup run; runs differed by 1.6% at most. Absolute numbers are lower than in the table above. That table also has #2653's FlyDSL attention on, and this run does not; the cause of the gap was not investigated.

ISL x CON	SP off tok/s	SP on tok/s	gain
4096 x 32	18,621	21,873	+17.5%
30720 x 8	17,803	20,754	+16.6%
129024 x 2	16,304	18,878	+15.8%

GSM8K 400: SP on 0.9775, SP off 0.9675. The difference is within sampling noise.

Native ATOM serving (atom.entrypoints.openai_server). This PR changes atom/models/qwen4_exp.py, which native ATOM uses too, not only the SGLang plugin. So the change was also checked on ATOM's own server. Setup:

• Code: this branch (58f995db4) merged with current main 606aceb11. The merge is clean.
• Image: rocm/atom-dev:sglang-latest with ATOM main d246797dc and aiter main 0dd271600; this tree is mounted over /app/ATOM.
• Hardware and model: 4x MI308X, Qwen3.8-Flash-Next-PTPC-FP8, -tp 4 --enable-expert-parallel, prefix caching off (the recipe's setting for MTP).
• SP off means ATOM_QWEN4_EXP_SP_MIN_TOKENS=0, which runs the same branches as main.
• Client: ATOM's own atom.benchmarks.benchmark_serving (random token ids, range 1.0, OSL 1). sglang.bench_serving's OpenAI client fails on the server's usage-only stream chunk.
• Each throughput value is the mean of r1 and r2 after a discarded warmup run r0; the two runs differed by 0.8% at most. GSM8K 400 is run_eval --api chat with thinking off.

# export ATOM_QWEN4_EXP_SP_MIN_TOKENS=0   # baseline
python -m atom.entrypoints.openai_server --model /model --trust-remote-code \
  -tp 4 --enable-expert-parallel --no-enable_prefix_caching --server-port 30480 \
  --method mtp --num-speculative-tokens 2

Configuration	Metric	SP off	SP on
MTP 2 steps, FULL cudagraph (default)	4096 x 32 input tok/s (mean TTFT ms)	18,923 (6,450)	22,326 (5,474), +18.0%
	30720 x 8 input tok/s (mean TTFT ms)	17,572 (12,527)	20,690 (10,655), +17.7%
	GSM8K 400	0.9675	0.9775
No MTP, FULL cudagraph	4096 x 32 input tok/s (mean TTFT ms)	20,559 (5,935)	25,026 (4,883), +21.7%
	GSM8K 400	0.9725	0.9725
MTP 2 steps, --enforce-eager, 400 concurrent requests	GSM8K 400	0.9725	0.9725

All requests succeeded, and no server log has a Traceback. A 4096-in / 512-out run with 16 concurrent requests completed 32/32, with MTP accept length 2.64 with SP on and 2.59 with SP off. The rows cover these paths:

• The no-MTP row runs the final sp.gather(mixed) path. The SGLang deployment always runs with MTP, so it never reaches this path.
• In the --enforce-eager row, decode and verify batches reach 1200 tokens (400 requests x 3), so SP also runs on eager decode / target verification. Temporary instrumentation (not committed) confirmed sharded decode forwards of up to 1200 tokens.
• Under the default FULL cudagraph, the bs = 512 decode bucket (1536 tokens) is sharded only in the eager warmup forward that the runner executes before capture (model_runner.py capture_cudagraph). Inside the capture torch.cuda.is_current_stream_capturing() keeps it unsharded, and replay does not re-enter the Python gate.
• --cudagraph-mode TR_HYBRID is rejected on main for this model ("Qwen3.8-Flash-Next has no TR_HYBRID graph support"), independent of this PR.
• PP, DP attention and TBO are unsupported for this model per recipes/Qwen3.8-Flash-Next.md.

Reproduction

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 SGLANG_AITER_KV_CACHE_LAYOUT=nhd AITER_MOE_PADDING_SIZE=128
# export ATOM_QWEN4_EXP_SP_MIN_TOKENS=0   # baseline
python -m sglang.launch_server --model-path /model --trust-remote-code --port 30480 \
  --tp 4 --ep-size 4 --attention-backend aiter --kv-cache-dtype bf16 --page-size 64 \
  --context-length 132096 --max-running-requests 64 --mem-fraction-static 0.88 \
  --cuda-graph-backend-decode full --cuda-graph-max-bs-decode 64 \
  --speculative-algorithm EAGLE --speculative-num-steps 2 --speculative-eagle-topk 1 \
  --sampling-defaults openai --mamba-radix-cache-strategy extra_buffer --max-mamba-cache-size 330

python -m sglang.bench_serving --backend sglang --port 30480 --model /model \
  --dataset-name random --random-input-len 4096 --random-output-len 1 --random-range-ratio 1.0 \
  --num-prompts 196 --max-concurrency 32 --request-rate inf --warmup-requests 1 --flush-cache

Not covered

• Decode / TPOT was not re-measured. The change does not run below 1024 tokens or under graph capture.
• Native ATOM serving at TP2 was not run (recipes/Qwen3.8-Flash-Next.md: pure TP2 is not validated); TP4 covers the row padding.

Test plan

• tests/test_qwen4_exp*.py: 104 passed, including the 13 new CPU cases in tests/test_qwen4_exp_sequence_parallel.py (ROCm, aiter main f0b96f49e)
• Follow-up fix 6c77e2a42 on current main, TP4: GSM8K 400 0.975; 4096 x 32 with SP on 21,902 tok/s (21,873 before the fix)
• black / ruff clean on the changed files
• Before / after on the deployment image tree, TP4 and TP2: six input-throughput cases and GSM8K 400, no prefix hits
• This branch on current main: GSM8K 400 and three input-throughput cases, SP on vs off
• Clean merge / rebase against perf(sglang): aiter FlyDSL QSA sparse attention for Qwen3.8-Flash-Next eager prefill #2653 head dc3a19bf5
• SP on vs off accuracy A/B with every prefill sharded (GSM8K 10-shot, radix off, 1319 questions, 2 runs each)
• Native ATOM server (atom.entrypoints.openai_server), TP4, merged with main 606aceb11: MTP / no MTP / --enforce-eager, SP on vs off, GSM8K 400 and input throughput; tests/test_qwen4_exp_sequence_parallel.py 13 passed on the merge
• 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 2834 --add-label <label>

[Copilot]

🟡 Changes recommended

Final-layer PLE leaves streams unsharded but still gathers them, producing extra output rows that break token alignment, while sequence-parallel correctness lacks test coverage.

2 open issues

• Avoid SP gather when PLE is in the final decoder layer · New
• Add CPU coverage for uneven _SequenceParallel token handling · New

What changed in this PR

This PR reduces duplicated hyper-connection work in Qwen3.8-Flash-Next by sharding residual streams across tensor-parallel ranks for large, non-captured forwards.

Changes:

• Adds token sharding, all-gather, and reduce-scatter around attention and MoE sublayers.
• Starts sharding after PLE and gathers final outputs, including the HC state used by MTP.
• Adds a configurable token threshold, with 0 disabling sequence parallelism.

File	Description
atom/​models/​qwen4_exp.py	Implements sequence-parallel residual streams and sublayer collective handling.
atom/​utils/​envs.py	Registers ATOM_QWEN4_EXP_SP_MIN_TOKENS, defaulting to 1024.

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The sharding path pads the entire replicated tensor before slicing, causing avoidable multi-gigabyte allocations per rank that remain retained by the shard views.

1 open issue

• Add CPU coverage for uneven _SequenceParallel token handling

1 previously missed issue

In code that hasn't changed since last review

Avoid full-tensor padding allocation before sharding

atom/​models/​qwen4_exp.py:683

For token counts not divisible by TP, shard() pads the entire replicated tensor before slicing, adding an avoidable full-size allocation and copy. With 131,071 tokens, TP4, and BF16 streams of width 10,240, each rank allocates 2.5 GiB to obtain a 0.625 GiB shard, whose view retains the full padded allocation. Slice the local rows first, then pad only that slice if it is shorter than self.chunk.

1 resolved since last review

• Avoid SP gather when PLE is in the final decoder layer

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The token-threshold gate can incorrectly enable sequence parallelism during MTP target verification, changing HC computation and reduction order despite the intended prefill-only scope.

1 open issue

• Add CPU coverage for uneven _SequenceParallel token handling

1 previously missed issue

In code that hasn't changed since last review

Restrict sequence parallelism to prefill, excluding target verification

atom/​models/​qwen4_exp.py:982

The token threshold does not enforce the stated prefill-only scope: with TP > 1, an eager target-verification batch of 256 requests × 5 tokens enables sequence parallelism at the default threshold. This changes verification's HC computation and reduction order despite the guarantee that MTP verification remains unchanged. Gate activation on the execution phase explicitly, excluding target verification, and add a regression case above the token threshold.

_{Review effort: Balanced}

[zovonoir]

Re Copilot's "Restrict sequence parallelism to prefill, excluding target verification" (atom/models/qwen4_exp.py:982): I kept the gate as it is and corrected the description, which overstated the scope.

• The gate is TP > 1, at least ATOM_QWEN4_EXP_SP_MIN_TOKENS tokens, and not capturing. Any eager target forward of that size is sharded, including a decode or target-verify batch that runs outside the CUDA graphs. At the default 1024 that needs 1024+ decode requests or 342+ verify requests (EAGLE 2 steps). Captured decode / verify graphs never shard. "When it applies" in the description now says this.
• This is intended. Sharding is mathematically identical: the only difference is that the BF16 sub-layer sum becomes a reduce-scatter instead of an all-reduce, so the summation order may change. That is the same kind of difference the existing all-reduce backends already show between each other. Excluding large eager verify batches would only give up the HC saving on those batches.
• Phase-gating would also need a different signal than context.is_prefill: the SGLang plugin sets it to True for target-verify batches whose attention state is a prefill (atom/plugin/sglang/runtime/forward_context.py:782), so it does not separate prefill from verification.

[ZLkanyo009]

LGTM, approving. Keeping the HC residual streams token-sharded after the last PLE layer is a nice way to split the replicated HC work TP ways without moving more bytes than before. I walked the collectives and they pair up the same way on every rank:

• Order within a sharded layer: gather, then attention/GDN, reduce-scatter, MLP mix on the shard, gather, MoE, reduce-scatter.
• No hidden reduces: o_proj, out_proj, the FusedMoE and the shared expert all run with reduce_results=False, so no sub-layer reduces behind the scheme's back.
• Rank-consistent gate: _sequence_parallel depends only on the TP size, the layer layout, the env value, the token count and the capture state, so ranks can't diverge and hang.
• Padding: shard/scatter use the same ceil(T/tp) split, and gather trims back to T rows, so attention, GDN, the QSA KV cache and the MoE only ever see real rows.
• Exits: both exit paths gather. The return_hc_state path gathers after apply_pending, so the MTP hand-off still gets the full [T, hc, hidden] state.

TP=1, CUDA-graph capture and forwards below the threshold run exactly the old code.

One non-blocking note: ATOM_QWEN4_EXP_SP_MIN_TOKENS defaults to 1024, so standalone ATOM and the RTP-LLM plugin at TP>1 also take this path, through the same Qwen4ExpModel.forward, while the numbers in the description are from SGLang. A quick accuracy sanity run on one of those at TP>1 (or a line in the description pointing to ATOM_QWEN4_EXP_SP_MIN_TOKENS=0 as the off switch) would round it out.

[zovonoir]

@valar-cop

[valar-cop]

@zovonoir checked both fixes: shard() now slices this rank's rows first and pads only that slice (atom/models/qwen4_exp.py:681), and _sequence_parallel returns None when sp_start_layer >= end_layer (atom/models/qwen4_exp.py:979), so the PLE-last-layer layout keeps the streams replicated and never gathers phantom rows. the new cpu tests cover both, and the padding case now asserts the shard owns no full-size storage.

one red ci: accuracy (kimi-k3 dspark dcp8 tp8), which these files can't affect, rerun it?

fine to land as long as #2847 follows right after, as the description says. no other concerns from my side.

[valar-cop]

a few small things..

ci Accuracy (Kimi-K3 DSpark DCP8 TP8) failed — not sure yet (low): The job log is empty (0 bytes), so the failure cannot be classified. None of the changed files is on this job's path: the PR touches atom/models/qwen4_exp.py, adds one inert atom/utils/envs.py entry, and adds a new test file, none of which Kimi-K3 DSpark DCP8 TP8 loads.

FAILED job Accuracy (Kimi-K3 DSpark DCP8 TP8) (ATOM vLLM Test), runner linux-atom-pitt-mi350x-8-6nfvf-runner-bjj47, log .valar/ci/job_124900318602.log
(.valar/ci/job_124900318602.log is empty; no log lines available)
next: Rerun the job and pull its log; treat as unrelated to this PR unless the log names a qwen4_exp/envs path. https://github.com/ROCm/ATOM/actions/runs/41324001342/job/124900318602

[zovonoir]

Added a Native ATOM serving section under Results. Because this PR changes atom/models/qwen4_exp.py, I also checked it on atom.entrypoints.openai_server, not only the SGLang plugin. Setup: TP4/EP4, PTPC-FP8, this branch merged with main 606aceb11.

The server starts and serves with SP on, and no server log has a Traceback.

Configuration	SP off	SP on
MTP, 4096 x 32 input throughput	18,923 tok/s	22,326 tok/s (+18.0%)
No MTP, 4096 x 32 input throughput	20,559 tok/s	25,026 tok/s (+21.7%)

GSM8K 400 in the same runs:

Configuration	SP off	SP on
MTP	0.9675	0.9775
No MTP	0.9725	0.9725
--enforce-eager, 400 concurrent requests (decode batches of 1200 tokens are sharded)	0.9725	0.9725

Under the default FULL cudagraph, captured graphs stay unsharded.

[valarLip]

Added a Native ATOM serving section under Results. Because this PR changes atom/models/qwen4_exp.py, I also checked it on atom.entrypoints.openai_server, not only the SGLang plugin. Setup: TP4/EP4, PTPC-FP8, this branch merged with main 606aceb11.

The server starts and serves with SP on, and no server log has a Traceback.

Configuration SP off SP on
MTP, 4096 x 32 input throughput 18,923 tok/s 22,326 tok/s (+18.0%)
No MTP, 4096 x 32 input throughput 20,559 tok/s 25,026 tok/s (+21.7%)
GSM8K 400 in the same runs:

Configuration SP off SP on
MTP 0.9675 0.9775
No MTP 0.9725 0.9725
--enforce-eager, 400 concurrent requests (decode batches of 1200 tokens are sharded) 0.9725 0.9725
Under the default FULL cudagraph, captured graphs stay unsharded.

mtp get lower throuput?

[zovonoir]

Added a Native ATOM serving section under Results. Because this PR changes atom/models/qwen4_exp.py, I also checked it on atom.entrypoints.openai_server, not only the SGLang plugin. Setup: TP4/EP4, PTPC-FP8, this branch merged with main 606aceb11.
The server starts and serves with SP on, and no server log has a Traceback.
Configuration SP off SP on
MTP, 4096 x 32 input throughput 18,923 tok/s 22,326 tok/s (+18.0%)
No MTP, 4096 x 32 input throughput 20,559 tok/s 25,026 tok/s (+21.7%)
GSM8K 400 in the same runs:
Configuration SP off SP on
MTP 0.9675 0.9775
No MTP 0.9725 0.9725
--enforce-eager, 400 concurrent requests (decode batches of 1200 tokens are sharded) 0.9725 0.9725
Under the default FULL cudagraph, captured graphs stay unsharded.

mtp get lower throuput?

Yes, this is a known issue, because I have enabled cuda graph for sglang plugin in : #2645 but not enabled for native atom server. At this situation, MTP is eager path and lead to performance drop compare with noMTP.

[Copilot]

🟢 Approval recommended

No active review issues are blocking approval of the token-sharded residual-stream changes.

No open issues 🎉

1 resolved since last review

• Add CPU coverage for uneven _SequenceParallel token handling

_{Review effort: Balanced}

[valar-cop]

LGTM

ci Accuracy (Kimi-K3 DSpark DCP8 TP8) failed — not sure yet (low): The job log is empty (0 bytes), so the failure cannot be classified. The job does not load any file this PR changes, and it failed the same way at the previous head on a different runner, which points away from the PR.

.valar/ci/job_126502556508.log is empty; no log lines available
next: Rerun the job with log capture (or pull the artifact); treat as unrelated to this PR unless the log names a qwen4_exp/envs path. https://github.com/ROCm/ATOM/actions/runs/41807150022/job/126502556508