Two changes to the --enable-aiter-allreduce-fusion path:

1. communicator.py (commit 449f7c293): Fix an off-by-one (< → <=) in apply_aiter_all_reduce_fusion so the fused kernel activates at the exact boundary size, matching AITER's internal should_custom_ar (inp_size <= max_size/2).

2. parallel_state.py (this PR): Remove the hardcoded hidden_dim ∈ {512, 1024, 2048, 4096} allowlist for 1-stage vs 2-stage selection. Keep the 128 KB byte-size threshold (the 1-stage kernel is capped at 80 tokens via kMaxBlocks) but let AITER's C++ dispatch decide which hidden_dims are supported.

Together these eliminate the per-model maintenance burden (no more manual hidden_dim additions for new models) and ensure AITER's own heuristics are respected.

Depends on:

• ROCm/aiter#2586 — fixes numerical
  accuracy in the 1-stage fused kernel (bf16 round-trip before residual addition).
• ROCm/aiter#2453 — refactors the fused allreduce dispatch so the 1-stage kernel accepts any hidden_dim divisible by pack_size (no more hardcoded template set), and the 2-stage kernel uses n_packs-based dispatch (no more n_bytes % 1024 alignment requirement).

Subsumes: #21947 — that PR adds 2880 to the Python-side allowlist. With this change, the allowlist is removed entirely, so all AITER-supported hidden_dims work automatically.

Motivation

Problem 1: Off-by-one in outer gate

The activation gate in communicator.py used < for the byte-size threshold while AITER uses <=. For hidden_size=4096 with bf16 at 8192 tokens, total_bytes exactly equals the threshold (67,108,864), so < rejected it.

Problem 2: Redundant hidden_dim allowlist

The 1-stage vs 2-stage decision in parallel_state.py gated on hidden_dim ∈ {512, 1024, 2048, 4096}. This was redundant at three levels:

AITER's C++ dispatch already enforces which hidden_dims have 1-stage support. Passing use_1stage=True for an unsupported dim is safe — AITER overrides it to false and falls through to 2-stage. The Python-side allowlist only created a maintenance burden (cf. PR #21947).

The 128 KB threshold is kept (but the allowlist is not)

The 1-stage fused kernel launches one CTA per token and is hard-capped at
kMaxBlocks = 80 tokens. For hidden_dim=4096 with bf16:
80 × 4096 × 2 = 655,360 bytes. The 128 KB threshold safely ensures only
small decode-like batches take the 1-stage path; larger prefill batches
automatically fall through to the 2-stage kernel.

Modifications

File 1: python/sglang/srt/layers/communicator.py

 def apply_aiter_all_reduce_fusion(input_tensor: torch.Tensor):
     n = input_tensor.shape[-1]
     total_bytes = input_tensor.numel() * input_tensor.element_size()
+    # Aiter's should_custom_ar uses <= max_size/2 (64 MB); match that boundary.
     return (
         _use_aiter
         and total_bytes > 0
         and n <= 16384
-        and total_bytes < 8 * 1024 * 8192
+        and total_bytes <= 8 * 1024 * 8192
         and get_tensor_model_parallel_world_size() != 6
         and not is_dp_attention_enabled()
         and get_global_server_args().enable_aiter_allreduce_fusion
     )

File 2: python/sglang/srt/distributed/parallel_state.py

-        # 1-stage policy for fused AR+RMSNorm:
-        # 1) Explicit env override wins.
-        # 2) Deterministic inference forces 1-stage for reproducibility.
-        # 3) Otherwise follow AITER's heuristic (small payloads only).
+        # 1-stage vs 2-stage selection for fused AR+RMSNorm:
+        # The 1-stage kernel launches one block per token and is capped at
+        # 80 tokens (kMaxBlocks).  Guard with a byte threshold so large
+        # prefill batches fall through to the 2-stage kernel instead of
+        # hitting a runtime error.  AITER's C++ dispatch already gates
+        # which hidden_dims have valid 1-stage support.
         if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set():
             use_1stage_ar = envs.SGLANG_USE_1STAGE_ALLREDUCE.get()
-        elif envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get():
-            use_1stage_ar = True
         else:
             total_bytes = input_.numel() * input_.element_size()
-            hidden_dim = input_.shape[-1]
-            use_1stage_ar = total_bytes <= 128 * 1024 and hidden_dim in {
-                512,
-                1024,
-                2048,
-                4096,
-            }
+            use_1stage_ar = total_bytes <= 128 * 1024

Accuracy Tests

Setup

• Model: Qwen3.5-397B-A17B-FP8 (60 layers, hidden=4096, 512 experts, MoE)
• Hardware: 4× AMD MI355X / gfx950 (TP=4, XGMI fully connected)
• Backend: SGLANG_USE_AITER=1, --attention-backend aiter
• AITER version: main branch (a530b92bf), includes both
  ROCm/aiter#2586 (1-stage accuracy fix) and
  ROCm/aiter#2453 (allreduce refactor).
  Full rebuild with GPU_ARCHS=gfx950 MAX_JOBS=192 PREBUILD_KERNELS=1.

GSM8K (1319 questions, 5-shot, parallel=1319)

Three rounds per configuration, with full PID/GPU cleanup between cases.

No accuracy regression. Fusion ON is +2.3pp above baseline, confirming the AITER accuracy fix (#2586) is correctly integrated and the fused path produces numerically sound results.

Commands

# Server (Case A — fusion ON)
SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
  --model-path /data2/Qwen/Qwen3.5-397B-A17B-FP8/ \
  --tp 4 --attention-backend aiter --trust-remote-code \
  --watchdog-timeout 1200 --mem-fraction-static 0.9 \
  --host 0.0.0.0 --port 9000 \
  --enable-aiter-allreduce-fusion

# Server (Case B — fusion OFF)
SGLANG_USE_AITER=1 python3 -m sglang.launch_server \
  --model-path /data2/Qwen/Qwen3.5-397B-A17B-FP8/ \
  --tp 4 --attention-backend aiter --trust-remote-code \
  --watchdog-timeout 1200 --mem-fraction-static 0.9 \
  --host 0.0.0.0 --port 9000

# GSM8K benchmark (run 3 rounds per case)
python3 benchmark/gsm8k/bench_sglang.py \
  --num-questions 1319 --parallel 1319 --num-shots 5 --port 9000

# Serving benchmark
python3 -m sglang.bench_serving \
  --host localhost --port 9000 \
  --model /data2/Qwen/Qwen3.5-397B-A17B-FP8/ \
  --dataset-name random \
  --random-input 8192 --random-output 1024 \
  --random-range-ratio 1.0 \
  --max-concurrency 1 --num-prompt 8

Benchmarking

Serving Benchmark (random, 8 × 8192-in / 1024-out, concurrency=1)

Kernel Microbenchmark (TP=4, bf16)

The fused kernel is faster across all tested shapes and hidden_dims:

torchrun --nproc_per_node=4 \
  benchmark/kernels/all_reduce/benchmark_fused_ar_rms_amd.py \
  --dtype bf16 \
  --decode-shapes 1x4096,4x4096,16x4096,32x4096,1x5120,16x5120,1x7168,16x7168,1x8192,1x14336,4x16384 \
  --prefill-shapes 128x4096,128x5120,128x7168,128x16384,8192x4096 \
  --warmup 10 --iters 30 --repeats 5

Decode (graph-captured)

With AITER #2453 applied, all hidden_dims
work — including GPT-OSS's hidden_dim=2880:

Prefill (eager)

All hidden_dims benefit — confirming that the allowlist was unnecessarily restrictive.
GPT-OSS (hidden_dim=2880) gets up to 2.10× speedup on decode.

Trace Verification

After the threshold fix, the EXTEND (prefill) trace on TP-0 confirms fused
kernel activation:

119 out of 121 allreduce sites now use the fused path. The remaining 2 are
smaller-hidden-dim layers (64-wide) that fall below AITER's internal threshold.

Summary

CC: @kkHuang-amd @HaiShaw