Fix: add missing end_sync barrier in fused allreduce+rmsnorm kernel

[yuzho-amd]

Motivation

Issue:
in sglang with fused allreduce+rmsnorm, it would have acc issue in lower concurrency(16), and the root cause is missing end_sync in all reduce fusion kernel.

How to reproduce:
image: lmsysorg/sglang-rocm:v0.5.15-rocm720-mi35x-20260713
model: amd/Qwen3.5-397B-A17B-MXFP4
start sglang server with:
CUDA_VISIBLE_DEVICES=0,1 \ SGLANG_USE_AITER=1 \ SGLANG_USE_AITER_UNIFIED_ATTN=1 \ AITER_FLYDSL_FORCE=1 \ SGLANG_MAMBA_SSM_DTYPE=bfloat16 \ SGLANG_USE_AITER_FP8_PER_TOKEN=1 \ python3 -m sglang.launch_server \ --model-path /data/amd/Qwen3.5-397B-A17B-MXFP4 \ --trust-remote-code \ --host 0.0.0.0 \ --port 6666 \ --tensor-parallel-size 2 \ --attention-backend aiter \ --mem-fraction-static 0.8 \ --model-loader-extra-config '{"enable_multithread_load": true}' \ --watchdog-timeout 7200 \ --disable-radix-cache \ --enable-aiter-allreduce-fusion \ --max-running-requests 16 \ --page-size 16
send prompt with:
'python3 /home/yuzho/inferencex_perf/run_refill_sequence_check.py
--base-url http://127.0.0.1:6666
--model /data/amd/Qwen3.5-397B-A17B-MXFP4
--concurrency 16
--num-requests 160
--timeout 7200
--output /tmp/bf16_conc16_refill.json'
the content of 'run_refill_sequence_check.py' is:
`#!/usr/bin/env python3
import argparse
import asyncio
import json
import re
import time
from collections import Counter
from pathlib import Path

import aiohttp

EXPECTED = list(range(1, 101))
REPEATED_SYMBOL = re.compile(r"([^\w\s，,。；;：:\-])\1{5,}")

def classify(text, finish_reason):
numbers = [int(value) for value in re.findall(r"\d+", text)]
errors = []
if not text.strip():
errors.append("empty")
if "\ufffd" in text:
errors.append("replacement_character")
if REPEATED_SYMBOL.search(text):
errors.append("repeated_symbols")
if numbers[:100] != EXPECTED:
errors.append(f"sequence_mismatch_{len(numbers)}")
if finish_reason not in ("stop", "length"):
errors.append(f"finish_reason_{finish_reason}")
return errors, numbers

async def main():
parser = argparse.ArgumentParser()
parser.add_argument("--base-url", default="http://127.0.0.1:6666")
parser.add_argument("--model", required=True)
parser.add_argument("--concurrency", type=int, required=True)
parser.add_argument("--num-requests", type=int, required=True)
parser.add_argument("--timeout", type=float, default=7200)
parser.add_argument("--output", required=True)
args = parser.parse_args()

output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
timeout = aiohttp.ClientTimeout(total=args.timeout)
connector = aiohttp.TCPConnector(limit=args.concurrency)
queue = asyncio.Queue()
for index in range(args.num_requests):
    queue.put_nowait(index)

results = []
started = time.monotonic()
async with aiohttp.ClientSession(timeout=timeout, connector=connector) as session:
    async def worker():
        while True:
            try:
                index = queue.get_nowait()
            except asyncio.QueueEmpty:
                return
            request_started = time.monotonic()
            prompt = (
                f"请求编号{index}。请从1到100按顺序输出整数，使用中文逗号分隔，"
                "只输出数字序列，不要解释，不要省略。"
            )
            record = {"index": index}
            try:
                async with session.post(
                    f"{args.base_url}/v1/chat/completions",
                    json={
                        "model": args.model,
                        "messages": [{"role": "user", "content": prompt}],
                        "temperature": 0,
                        "max_tokens": 512,
                        "chat_template_kwargs": {"enable_thinking": False},
                    },
                ) as response:
                    body = await response.text()
                    record["http_status"] = response.status
                    if response.status != 200:
                        raise RuntimeError(body[:1000])
                    data = json.loads(body)
                    choice = data["choices"][0]
                    text = choice["message"]["content"] or ""
                    errors, numbers = classify(text, choice.get("finish_reason"))
                    record.update(
                        text=text,
                        errors=errors,
                        ok=not errors,
                        finish_reason=choice.get("finish_reason"),
                        number_count=len(numbers),
                    )
            except Exception as exc:
                record.update(
                    text="",
                    errors=[f"request_error:{type(exc).__name__}:{exc}"],
                    ok=False,
                    finish_reason=None,
                    number_count=0,
                )
            record["elapsed_seconds"] = time.monotonic() - request_started
            results.append(record)
            queue.task_done()

    await asyncio.gather(*(worker() for _ in range(args.concurrency)))

results.sort(key=lambda item: item["index"])
failures = [item for item in results if not item["ok"]]
error_counts = Counter(error for item in failures for error in item["errors"])
summary = {
    "model": args.model,
    "concurrency": args.concurrency,
    "num_requests": args.num_requests,
    "passed": len(results) - len(failures),
    "failed": len(failures),
    "elapsed_seconds": time.monotonic() - started,
    "error_counts": dict(sorted(error_counts.items())),
    "first_failure": failures[0] if failures else None,
    "results": results,
}
output_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2))
print(json.dumps({key: value for key, value in summary.items() if key != "results"},
                 ensure_ascii=False))
raise SystemExit(0 if not failures else 2)

if name == "main":
asyncio.run(main())
`

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4346 --add-label <label>

[Copilot]

Pull request overview

Adds the missing final cross-rank synchronization (end_sync) to the 1-stage fused allreduce+RMSNorm kernels so a rank cannot start the next fused collective (and reuse the registered buffers/signal slots) before all peers have finished the prior invocation, preventing silent cross-request corruption under certain concurrency patterns.

Changes:

• Add end_sync<ngpus, true>(...) at kernel exit for the 1-stage fused allreduce+rmsnorm kernel.
• Add end_sync<ngpus, true>(...) at kernel exit for the per-group quant 1-stage variant.
• Add end_sync<ngpus, true>(...) at kernel exit for the MXFP4 1-stage variant (and document the rationale in-code).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[TennyWang1223]

LGTM