[merceod]

mstar::allreduce_add_rmsnorm is still the op models call. What runs behind it is now an ARFusion strategy, picked per platform with PlatformDispatch.

• FlashInferARFusion is today's CUDA code moved into a class, including the backend fallback from distributed: fall back to FlashInfer's trtllm all-reduce fusion, agreed across TP ranks #563 ("auto", then "trtllm", agreed across TP ranks). test_ar_fusion_fallback.py now calls FlashInferARFusion.create.
• AiterARFusion (ROCm) uses AITER's custom all-reduce with its one-stage fused RMSNorm, only where vLLM found that kernel faster than RCCL + norm. That's up to 80 tokens, at most 1024 16-byte packs per row, and under 256 KiB at TP4 or 128 KiB at TP8. Everything else stays on RCCL + norm, and the two-stage kernel is never used. The IPC buffer is sized to that token bound instead of AITER's default.
• A startup check runs the kernel against a host reference on every rank (1 and max tokens, gemma on and off, outputs prefilled with NaN), then the ranks agree over gloo. Every rank runs every case even after one fails, so no rank is left spinning in the kernel. It refuses AITER older than 0.1.19.post1 (the one-stage kernel could return before its peers finished reading), any numeric mismatch, and builds that round bf16 toward zero. A refusal on any rank means RCCL + norm on all of them.
• AITER's eager input pool is built with AITER_CUSTOM_AR_RAW_INPUT_POOL=1, set only while our CustomAllreduce is constructed. AITER shares that pool's IPC handle with offset 0, so when torch's caching allocator puts the pool partway into a segment (the usual case after model load) the peers read the wrong memory. That caused NaN logits on Qwen3.5-27B at TP4, and the per-rank startup check failures and hangs we saw. Filed as CustomAllreduce: fused_ar_rms reads the wrong peer memory when the eager input pool doesn't start its allocation (IPC offset is always 0) ROCm/aiter#6918.
• Both graph runners wrap each capture in tp_group.graph_capture(). On ROCm it registers the recorded inputs with AITER once the ranks agree on the count, and otherwise drops that graph on every rank.

On rounding, AITER's release wheels truncate fp32 to bf16 on gfx942. Over 128 steps against an fp64 stream the wheel's kernel drifted 17.5% and shrank the residual by 15.5%, against 1.37% / -0.02% for a round-to-nearest build (OPUS_FP32_to_BF16_DEFAULT=0) and 1.4% / -0.02% for RCCL + add (MI325X, test/rocm/test_ar_fusion_tp.py). So the check refuses the stock wheel and a stock install serves TP on RCCL + norm. Rebuilding AITER's module_custom_all_reduce with OPUS_FP32_to_BF16_DEFAULT=0 into AITER_JIT_DIR turns the fused path on, and the numbers below use that build.

NVIDIA: the FlashInfer path is the same code, and graph_capture() is a no-op scope on CUDA and at TP1.

scripts/ci/select_cpu_tests.py lists test/rocm with the modules CPU Core doesn't run (they need an AMD GPU), since main's build check wants every test module accounted for.

Tests

What	Result
ruff check .	clean
test_ar_fusion.py (new), test_ar_fusion_fallback.py, test_accelerator_graph_capture.py, test_parallel_groups.py, test_piecewise_shared_pool.py, test_packed_capture_multiplier.py, test_platform_dispatch.py, test_preplan_error_paths.py	pass in the run of the whole integration branch on MI300X below
test/rocm/test_ar_fusion_tp.py (multi-rank harness) on 8x MI325X, eager	0 failures at TP2/4/8 x hidden 2048/4096/5120/8192 x 1-4096 tokens x gemma on/off (240 cases per TP size), bit-repeatable. At TP2 hidden 8192 AITER's kernel leaves the output untouched, and the startup check refuses it
same, captured 16-call chains, copy-in and registered	333 cases, 0 failures
Kokoro on one MI300X, through both changed capture loops at TP1	110 graphs captured, serves
test/modular, test/distributed and test/sampling_test on an MI300X node with GPUs hidden, on the integration branch that carries this whole stack, rebased on main	only the 7 environment failures main has there too (rsync, ftfy)
test/rust there	359 passed
test/rocm on MI300X, same branch	674 passed, plus the fused all-reduce harness at 2, 4 and 8 GPUs
CI's CPU Core jobs (since #581 all of test/modular, test/distributed and the model test dirs, plus the TP lockstep sim; torch 2.9.1 CPU, merged with main)	green on every cut so far (last run 2026-10-21: 5956 passed, 749 skipped), TP sim passes
Rebase onto main 54f29da (2026-10-26, the whole stack)	CPU on a devel node, GPUs hidden: test/modular only fails the two test_dump_stacks environment cases main fails too (7174 passed before the last three test fixes, targeted rerun on the tip 398 passed), other CPU suites 248 passed, TP lockstep sim 7 passed, test/rust 368 passed. One MI300X: test/rocm 675 passed, fused-MoE and sampler GPU tests 228 passed, served smoke rows for Qwen3.5 0.8B, Kokoro, Whisper large-v3 and Cosmos3 DROID match the earlier runs

Greedy parity against TP1, 16 prompts x 64 tokens at c=32, 8x MI300X. Batch composition alone flips 1-4 prompts, so +-3 is noise.

	fused	RCCL + norm
Qwen3.5-27B TP4	13/16	13/16
Qwen3.5-9B TP2	15/16	15/16

No degenerate outputs in either arm. For the input pool fix: the older gating from users/kirill/amd-integration (no startup check) on the allocator's pool gave NaN logits on 32 of 32 requests at 27B TP4, and 0 of 32 with the raw pool.

Numbers

Per call inside a graph, MI325X, test/rocm/test_ar_fusion_tp.py, TP4 hidden 5120 (us):

tokens	fused	RCCL + norm
1	6.5	21.1
8	8.0	18.4
25	13.8	31.6

End to end, tok/s, benchmark.runner closed loop under benchmark.worker_phases, 256 output tokens, two repeats, 8x MI300X node using GPUs 0-3 (other jobs ran on the other GPUs):

	c=1	c=4	c=16	c=64
9B TP2 fused	180-185	649-666	2157-2220	4774-4952
9B TP2 RCCL + norm	147	547	1891-1898	4555-4605
27B TP4 fused	112	380	1356-1380	2830-2872
27B TP4 RCCL + norm	91	322	1108	2823-2832

That's +15-23% at c<=16. At c=64 most batches are over the fused bound (25 tokens at TP4 hidden 5120) and take RCCL + norm anyway. At batch 1, worker 0's event_sync goes from 5.2 to 4.1 ms (9B) and from 9.3 to 7.3 ms (27B), so the gain is GPU time. Registered capture is a little faster than copy-in at c<=16 (9B c=1 180-185 vs 175.6 tok/s) and the same at c=64. MI325X end-to-end reruns follow.

Not in this PR

• installation.rst doesn't describe the round-to-nearest AITER build yet. That's a docs follow-up.
• AITER's fused kernel doesn't run at TP2 hidden 8192. The check refuses it, so those configs use RCCL + norm.
• Whisper and Qwen3-Omni don't call the fused op.

Since the first cut

• ARFusion.new(group, hidden, dtype, max_tokens) is the factory, a classmethod that takes the CommGroup like main's AllReduce.new (GLM-5.2/GLM-5.3 PR 1: engine: MLA attention, fp8 fused MoE, symmetric all-reduce #386), and the FlashInfer every-rank vote uses _all_ranks_agree from all_reduce.py instead of a copy.

Since r8

Since the r8 cut: rebuilt on main's #881. The strategies take the CommGroup, FlashInfer's all-reduce is built with ARFusion.new (main's workspace layout, non-CUDA refused), and both strategies vote through CommGroup.agree, so the module-level gloo helper is gone.

Checks on this rebase

The stack now sits on main f7cb5ca (2026-10-29), every branch cut from one integration tree.

Check	Result
CPU on a devel node with the GPUs hidden: ruff, test/ci, the TP lockstep sim, test/rust, CPU Core	ruff clean, sim 14 passed, test/rust 436 passed, CPU Core fails nothing main doesn't (main's own 34: test_dump_stacks x2, loop_output_ownership x32)
test/rocm on one MI325X (a VM slice, torch 2.12.1+rocm7.2)	715 passed, 9 skipped, 0 failed
MoE and sampler tests there	MoE EP 53 passed, sampler 145 passed. The 75 test_*moe* failures are main's GLM-5.2 kernel tests, the same set on the previous head
Served rows there through the correctness runner	all PASS: Qwen3.5 0.8B and 9B text (no divergence away from near ties, 4/4 needles), Whisper large-v3 (WER 1.64%, 99/100 transcripts equal to the reference), Kokoro (WER delta 0), Qwen3-Omni colocated (50/50, spoken WER 0), pi0.5 and Cosmos3 DROID (two instances bit-same)
Qwen3.5 9B TP2 on two GPUs there	PASS (21/32 exact, 11 near ties, 0 unexplained), the fused all-reduce on both ranks

Nothing fails on this tree that the previous head doesn't: the failing-set diff against it is empty for test/rocm, MoE, the sampler and the rows.

For this PR: The TP2 run: both ranks log the AITER CustomAllreduce workspace built, so every CommGroup.agree vote agreed, and test/rocm/test_ar_fusion_tp.py passes at world 2 (world 4 and 8 need a node).