diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index 7895df2bc..529fad125 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -1080,7 +1080,7 @@ jobs: fail-fast: false matrix: runner: [linux-mi35x-gpu-8] - part: [0, 1] + part: [0, 1, 2] runs-on: ${{matrix.runner}} steps: - name: Checkout code @@ -1101,7 +1101,7 @@ jobs: - name: Run test timeout-minutes: 60 run: | - bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} + bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} # =============================================== Disaggregation ==================================================== stage-b-test-large-8-gpu-mi35x-disaggregation-amd-rocm720: diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index b2162bdde..307aab6b0 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -1135,7 +1135,7 @@ jobs: fail-fast: false matrix: runner: [linux-mi35x-gpu-8] - part: [0, 1] + part: [0, 1, 2] runs-on: ${{matrix.runner}} steps: - name: Checkout code @@ -1157,7 +1157,7 @@ jobs: - name: Run test timeout-minutes: 60 run: | - bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} + bash scripts/ci/amd/amd_ci_exec.sh -w "/sglang-checkout/test" python3 run_suite.py --hw amd --suite stage-c-test-large-8-gpu-amd-mi35x --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 3600 ${{ needs.check-changes.outputs.continue_on_error == 'true' && '--continue-on-error' || '' }} # =============================================== Disaggregation ==================================================== stage-b-test-large-8-gpu-mi35x-disaggregation-amd: diff --git a/benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py b/benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py new file mode 100644 index 000000000..85cabe94c --- /dev/null +++ b/benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py @@ -0,0 +1,540 @@ +""" +Benchmark fused AllReduce + RMSNorm + per-group FP8 quant on AMD with +correctness checks. + +This script targets the three op paths used by SGLang on ROCm/aiter for +Qwen3.5-FP8 style models: + + 1. Split (3 kernels) - reference: + tensor_model_parallel_all_reduce -> RMSNorm -> aiter per-1x128 quant. + 2. Fused AR+RMSNorm + separate per-group quant (2 kernels): + tensor_model_parallel_fused_allreduce_rmsnorm -> aiter per-1x128 quant. + 3. Fully fused AR+RMSNorm+per-group-quant (1 kernel): + tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group. + +Default shape sets cover the Qwen3.5-397B-A17B-FP8 layout: + * hidden_size = 4096 + * TP = 8 (launched with torchrun --nproc_per_node=8) + * Prefill batch sizes up to a few thousand tokens. + * Decode batch sizes 1-512 covering typical steady-state running_req values. + +Usage: + torchrun --nproc_per_node=8 \ + benchmark/kernels/all_reduce/benchmark_fused_ar_rms_quant_amd.py \ + --dtype bf16 --group-size 128 +""" + +import argparse +import csv +import os +import statistics +from typing import Dict, List, Optional, Tuple + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from sglang.srt.distributed.communication_op import ( + tensor_model_parallel_all_reduce, + tensor_model_parallel_fused_allreduce_rmsnorm, + tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group, +) +from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + graph_capture, + init_distributed_environment, + initialize_model_parallel, + set_custom_all_reduce, +) + +Shape = Tuple[int, int] +FP8_DTYPE = torch.float8_e4m3fnuz + + +def parse_shapes(raw: str) -> List[Shape]: + shapes: List[Shape] = [] + for item in [x.strip() for x in raw.split(",") if x.strip()]: + if "x" not in item: + raise ValueError(f"Invalid shape '{item}', expected MxN format.") + m_str, n_str = item.split("x", 1) + m, n = int(m_str), int(n_str) + if m <= 0 or n <= 0: + raise ValueError(f"Invalid shape '{item}', both dims must be positive.") + shapes.append((m, n)) + if not shapes: + raise ValueError("Empty shape list is not allowed.") + return shapes + + +def dtype_from_name(name: str) -> torch.dtype: + mapping = { + "fp16": torch.float16, + "float16": torch.float16, + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + } + if name not in mapping: + raise ValueError(f"Unsupported dtype: {name}") + return mapping[name] + + +def _barrier(device: torch.device) -> None: + try: + dist.barrier(device_ids=[device.index]) + except TypeError: + dist.barrier() + + +def _mean_across_ranks(val: float, device: torch.device) -> float: + t = torch.tensor([val], dtype=torch.float64, device=device) + dist.all_reduce(t, op=dist.ReduceOp.SUM) + t /= dist.get_world_size() + return float(t.item()) + + +def _all_true_across_ranks(val: bool, device: torch.device) -> bool: + t = torch.tensor([1 if val else 0], dtype=torch.int32, device=device) + dist.all_reduce(t, op=dist.ReduceOp.MIN) + return bool(int(t.item())) + + +def _measure_us( + fn, warmup: int, iters: int, repeats: int, device: torch.device +) -> float: + for _ in range(max(1, warmup)): + fn() + torch.cuda.synchronize() + + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + samples: List[float] = [] + for _ in range(max(1, repeats)): + _barrier(device) + torch.cuda.synchronize() + start.record() + for _ in range(iters): + fn() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1000.0 / iters) + samples.sort() + return float(statistics.median(samples)) + + +def _make_inputs( + shape: Shape, dtype: torch.dtype, seed: int, rank: int, device: torch.device +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + m, n = shape + torch.manual_seed(seed + rank * 17) + # fp32 first then downcast so every rank has distinct values that still + # sum to a well-conditioned pre-norm tensor after all-reduce. + x = torch.randn((m, n), dtype=torch.float32, device=device).to(dtype) + residual = x.clone() + weight = torch.randn((n,), dtype=torch.float32, device=device).to(dtype) + return x, residual, weight + + +def _split_3_reference( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + group_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reference: plain all_reduce -> RMSNorm -> aiter per-1x128 quant.""" + import aiter + + ar_out = tensor_model_parallel_all_reduce(x.clone()) + residual_out = ar_out + residual + normed = F.rms_norm(residual_out, (residual_out.shape[-1],), weight, eps) + hip_quant = aiter.get_hip_quant(aiter.QuantType.per_1x128) + fp8_out, scale_out = hip_quant(normed, quant_dtype=aiter.dtypes.fp8) + return fp8_out, residual_out, scale_out + + +def _fused_ar_rms_then_quant( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + group_size: int, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """2-kernel: fused AR+RMSNorm (existing) + separate per-group quant.""" + import aiter + + result = tensor_model_parallel_fused_allreduce_rmsnorm( + x.clone(), residual.clone(), weight, eps + ) + if result is None: + return None + normed, residual_out = result + hip_quant = aiter.get_hip_quant(aiter.QuantType.per_1x128) + fp8_out, scale_out = hip_quant(normed, quant_dtype=aiter.dtypes.fp8) + return fp8_out, residual_out, scale_out + + +def _fully_fused( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + group_size: int, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """1-kernel: fused AR+RMSNorm+per-group-quant (fp8+scale only).""" + return tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group( + x.clone(), residual.clone(), weight, eps, group_size + ) + + +def _fully_fused_with_bf16( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + group_size: int, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]: + """1-kernel: fused AR+RMSNorm+per-group-quant with bf16 side-output + (GDN keep_bf16=True path — replaces fused_ar_rms + separate per-group + quant with a single kernel that writes BOTH fp8+scale and bf16). + """ + return tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group( + x.clone(), residual.clone(), weight, eps, group_size, emit_bf16=True + ) + + +def _check_quant_close( + fp8_a: torch.Tensor, + scale_a: torch.Tensor, + fp8_b: torch.Tensor, + scale_b: torch.Tensor, + group_size: int, +) -> Tuple[bool, str]: + """Compare two (fp8, scale) per-group quantized outputs by dequantizing.""" + dq_a = fp8_a.float() * scale_a.repeat_interleave(group_size, dim=-1) + dq_b = fp8_b.float() * scale_b.repeat_interleave(group_size, dim=-1) + max_diff = (dq_a - dq_b).abs().max().item() + denom = dq_a.abs().max().item() + 1e-6 + rel_err = max_diff / denom + ok = rel_err < 0.15 + return ok, f"max_diff={max_diff:.4f},rel={rel_err:.4f}" + + +def bench_shape( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float, + group_size: int, + warmup: int, + iters: int, + repeats: int, + mode: str, +) -> Dict[str, object]: + device = x.device + + # --- Split 3-kernel baseline --- + split_fn = lambda: _split_3_reference(x, residual, weight, eps, group_size) + if mode == "graph": + with graph_capture() as gc: + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=gc.stream): + _split_3_reference(x, residual, weight, eps, group_size) + split_fn = g.replay + split_us = _measure_us(split_fn, warmup, iters, repeats, device) + + # --- Fused AR+RMSNorm + separate quant (2 kernels) --- + probe2 = _fused_ar_rms_then_quant(x, residual, weight, eps, group_size) + fused2_available = probe2 is not None + fused2_us: Optional[float] = None + if fused2_available: + fused2_fn = lambda: _fused_ar_rms_then_quant( + x, residual, weight, eps, group_size + ) + if mode == "graph": + with graph_capture() as gc: + g2 = torch.cuda.CUDAGraph() + with torch.cuda.graph(g2, stream=gc.stream): + _fused_ar_rms_then_quant(x, residual, weight, eps, group_size) + fused2_fn = g2.replay + fused2_us = _measure_us(fused2_fn, warmup, iters, repeats, device) + + # --- Fully fused, fp8-only (1 kernel) — std-attention path --- + probe1 = _fully_fused(x, residual, weight, eps, group_size) + fused1_available = probe1 is not None + fused1_us: Optional[float] = None + if fused1_available: + fused1_fn = lambda: _fully_fused(x, residual, weight, eps, group_size) + if mode == "graph": + with graph_capture() as gc: + g1 = torch.cuda.CUDAGraph() + with torch.cuda.graph(g1, stream=gc.stream): + _fully_fused(x, residual, weight, eps, group_size) + fused1_fn = g1.replay + fused1_us = _measure_us(fused1_fn, warmup, iters, repeats, device) + + # --- Fully fused, fp8+bf16 (1 kernel) — GDN keep_bf16=True path --- + probe1b = _fully_fused_with_bf16(x, residual, weight, eps, group_size) + # The bf16 side-output is only emitted when the call actually returned + # a 4-tuple; a 3-tuple means the aiter build doesn't support it. + fused1bf16_available = ( + probe1b is not None and isinstance(probe1b, tuple) and len(probe1b) == 4 + ) + fused1bf16_us: Optional[float] = None + if fused1bf16_available: + fused1bf16_fn = lambda: _fully_fused_with_bf16( + x, residual, weight, eps, group_size + ) + if mode == "graph": + with graph_capture() as gc: + g1b = torch.cuda.CUDAGraph() + with torch.cuda.graph(g1b, stream=gc.stream): + _fully_fused_with_bf16(x, residual, weight, eps, group_size) + fused1bf16_fn = g1b.replay + fused1bf16_us = _measure_us(fused1bf16_fn, warmup, iters, repeats, device) + + # --- Correctness --- + # (a) fused1 fp8+scale vs fused2 fp8+scale (both emit fp8 pair) + # (b) fused1_bf16 bf16 side-output vs fused2 bf16 (both describe the + # same normed value; fused2 writes bf16 explicitly, fused1_bf16 + # writes bf16 AND fp8; comparing bf16 against bf16 is the tightest + # check of the bf16 hook). + correctness = "N/A" + if fused1_available and fused2_available: + res1 = _fully_fused(x, residual, weight, eps, group_size) + res2 = _fused_ar_rms_then_quant(x, residual, weight, eps, group_size) + ok, detail = _check_quant_close(res2[0], res2[2], res1[0], res1[2], group_size) + correctness = "PASS" if ok else f"FAIL({detail})" + + correctness_bf16 = "N/A" + if fused1bf16_available and fused2_available: + res1b = _fully_fused_with_bf16(x, residual, weight, eps, group_size) + res2 = _fused_ar_rms_then_quant(x, residual, weight, eps, group_size) + # res1b = (fp8, res_out, scale, bf16); res2 = (fp8, res_out, scale) + # Cross-check: dequant(res1b.fp8) ≈ dequant(res2.fp8) and + # res1b.bf16 ≈ dequant(res1b.fp8) within ~one FP8 step. + ok_fp8, detail_fp8 = _check_quant_close( + res2[0], res2[2], res1b[0], res1b[2], group_size + ) + bf16_vs_fp8 = ( + ( + res1b[3].float() + - (res1b[0].float() * res1b[2].repeat_interleave(group_size, dim=-1)) + ) + .abs() + .max() + .item() + ) + if not ok_fp8: + correctness_bf16 = f"FAIL_fp8({detail_fp8})" + elif bf16_vs_fp8 > 1.0: + correctness_bf16 = f"FAIL_bf16(diff={bf16_vs_fp8:.4f})" + else: + correctness_bf16 = f"PASS(bf16_diff={bf16_vs_fp8:.3f})" + + return { + "split_us": split_us, + "fused2_available": fused2_available, + "fused2_us": fused2_us, + "fused1_available": fused1_available, + "fused1_us": fused1_us, + "fused1bf16_available": fused1bf16_available, + "fused1bf16_us": fused1bf16_us, + "correctness": correctness, + "correctness_bf16": correctness_bf16, + } + + +# Qwen3.5-397B-A17B-FP8 has hidden_size=4096 (both GDN and standard attention +# layers go through input_layernorm at full hidden dim before the TP projections +# are applied, so the fused op sees [M, 4096] inputs on every rank). +_DEFAULT_PREFILL_SHAPES = ( + "64x4096,128x4096,256x4096,512x4096,1024x4096,2048x4096,4096x4096,8192x4096" +) +_DEFAULT_DECODE_SHAPES = ( + "1x4096,2x4096,4x4096,8x4096,16x4096,32x4096,64x4096,128x4096,256x4096,512x4096" +) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Benchmark fused AR+RMSNorm+per-group-quant for Qwen3.5-FP8 shapes." + ) + ) + parser.add_argument( + "--dtype", + type=str, + default="bf16", + choices=["fp16", "bf16", "float16", "bfloat16"], + ) + parser.add_argument("--eps", type=float, default=1e-6) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--group-size", type=int, default=128) + parser.add_argument("--prefill-shapes", type=str, default=_DEFAULT_PREFILL_SHAPES) + parser.add_argument("--decode-shapes", type=str, default=_DEFAULT_DECODE_SHAPES) + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--iters", type=int, default=30) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument( + "--mode", + type=str, + default="both", + choices=["eager", "graph", "both"], + ) + parser.add_argument("--csv-out", type=str, default=None) + args = parser.parse_args() + + dtype = dtype_from_name(args.dtype) + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + torch.cuda.set_device(local_rank % torch.cuda.device_count()) + device = torch.device(f"cuda:{local_rank % torch.cuda.device_count()}") + + set_custom_all_reduce(True) + init_distributed_environment( + world_size=world_size, + rank=rank, + local_rank=local_rank, + distributed_init_method="env://", + backend="nccl", + ) + initialize_model_parallel(tensor_model_parallel_size=world_size) + + if rank == 0: + print( + f"Config: world_size={world_size}, dtype={dtype}, " + f"group_size={args.group_size}" + ) + print( + f" 1-stage boundary: total_bytes <= 128KB " + f"(M <= {128 * 1024 // (4096 * 2)} for hidden=4096 bf16)" + ) + print( + f" 2-stage boundary: total_bytes <= 512KB " + f"(M <= {512 * 1024 // (4096 * 2)} for hidden=4096 bf16)" + ) + print( + " fallback: fused_ar_rms + per_group_quant (2 kernels) " + "when single-kernel path is unavailable" + ) + + run_modes = ("eager", "graph") if args.mode == "both" else (args.mode,) + csv_rows: List[Dict[str, object]] = [] + + for mode in run_modes: + shapes = parse_shapes( + args.prefill_shapes if mode == "eager" else args.decode_shapes + ) + if rank == 0: + phase = "prefill(eager)" if mode == "eager" else "decode(graph)" + print(f"\n{'=' * 145}") + print(f"Mode: {phase}") + print( + "| Shape | Bytes/rank | Split(3k) us | Fused2(2k) us | " + "Fused1(1k) us | Fused1+bf16(1k) us | Speedup(2k) | " + "Speedup(1k) | Speedup(1k+bf16) | Corr fp8 | Corr bf16 |" + ) + print( + "|:------|----------:|-----------:|------------:|-----------:|" + "-----------:|-----------:|-----------:|-----------:|" + ":---------|:----------|" + ) + + for shape in shapes: + x, residual, weight = _make_inputs(shape, dtype, args.seed, rank, device) + m = bench_shape( + x, + residual, + weight, + args.eps, + args.group_size, + args.warmup, + args.iters, + args.repeats, + mode, + ) + + split_us = _mean_across_ranks(m["split_us"], device) + fused2_avail = _all_true_across_ranks(m["fused2_available"], device) + fused1_avail = _all_true_across_ranks(m["fused1_available"], device) + fused1bf16_avail = _all_true_across_ranks(m["fused1bf16_available"], device) + fused2_us = ( + _mean_across_ranks(m["fused2_us"], device) + if m["fused2_us"] is not None + else None + ) + fused1_us = ( + _mean_across_ranks(m["fused1_us"], device) + if m["fused1_us"] is not None + else None + ) + fused1bf16_us = ( + _mean_across_ranks(m["fused1bf16_us"], device) + if m["fused1bf16_us"] is not None + else None + ) + + if rank == 0: + M, N = shape + nbytes = M * N * 2 + f2_str = f"{fused2_us:.1f}" if fused2_us else "N/A" + f1_str = f"{fused1_us:.1f}" if fused1_us else "N/A" + f1b_str = f"{fused1bf16_us:.1f}" if fused1bf16_us else "N/A" + s2 = ( + f"{split_us / fused2_us:.2f}x" + if fused2_us and fused2_us > 0 + else "N/A" + ) + s1 = ( + f"{split_us / fused1_us:.2f}x" + if fused1_us and fused1_us > 0 + else "N/A" + ) + s1b = ( + f"{fused2_us / fused1bf16_us:.2f}x" + if fused1bf16_us and fused2_us and fused1bf16_us > 0 + else "N/A" + ) + print( + f"| {M}x{N} | {nbytes} | {split_us:.1f} | {f2_str} | " + f"{f1_str} | {f1b_str} | {s2} | {s1} | {s1b} | " + f"{m['correctness']} | {m['correctness_bf16']} |" + ) + csv_rows.append( + { + "mode": mode, + "shape": f"{M}x{N}", + "m": M, + "n": N, + "bytes_per_rank": nbytes, + "split_us": split_us, + "fused2_us": fused2_us if fused2_us is not None else "", + "fused1_us": fused1_us if fused1_us is not None else "", + "fused1bf16_us": ( + fused1bf16_us if fused1bf16_us is not None else "" + ), + "fused1_available": fused1_avail, + "fused2_available": fused2_avail, + "fused1bf16_available": fused1bf16_avail, + "correctness": m["correctness"], + "correctness_bf16": m["correctness_bf16"], + } + ) + + if rank == 0 and args.csv_out and csv_rows: + os.makedirs(os.path.dirname(args.csv_out) or ".", exist_ok=True) + with open(args.csv_out, "w", newline="") as f: + w = csv.DictWriter(f, fieldnames=list(csv_rows[0].keys())) + w.writeheader() + w.writerows(csv_rows) + print(f"\nSaved CSV: {args.csv_out}") + + _barrier(device) + destroy_model_parallel() + destroy_distributed_environment() + + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/distributed/communication_op.py b/python/sglang/srt/distributed/communication_op.py index 89f9986e4..e8f91bd53 100644 --- a/python/sglang/srt/distributed/communication_op.py +++ b/python/sglang/srt/distributed/communication_op.py @@ -40,6 +40,29 @@ def tensor_model_parallel_fused_allreduce_rmsnorm( return get_tp_group().fused_allreduce_rmsnorm(input_, residual_inp_, weight_, eps) +def tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group( + input_: torch.Tensor, + residual_inp_: torch.Tensor, + weight_: torch.Tensor, + eps: float, + group_size: int = 128, + emit_bf16: bool = False, +) -> Optional[Tuple[torch.Tensor, ...]]: + """Fused TP all-reduce + RMSNorm + per-group FP8 quant (ROCm/aiter). + + Returns ``(fp8_output, residual_out, per_group_scale)`` by default, or + ``(fp8_output, residual_out, per_group_scale, bf16_output)`` when + ``emit_bf16=True`` (kernel writes both fp8 and the pre-quantization bf16 + normed output — no extra kernel). ``None`` when the backend cannot + service the request (non-AMD, custom AR disabled, shape unsupported). + Callers MUST handle ``None`` by falling back to the separate + fused-AR-RMSNorm + per-group-quant path. + """ + return get_tp_group().fused_allreduce_rmsnorm_quant_per_group( + input_, residual_inp_, weight_, eps, group_size, emit_bf16=emit_bf16 + ) + + def tensor_model_parallel_all_gather( input_: torch.Tensor, dim: int = -1 ) -> torch.Tensor: diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 241be9026..01ca678c3 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -60,6 +60,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_cuda_alike, + is_gfx95_supported, is_hip, is_musa, is_npu, @@ -784,6 +785,77 @@ class GroupCoordinator: ) return fused_outputs + def fused_allreduce_rmsnorm_quant_per_group( + self, + input_: torch.Tensor, + residual_inp_: torch.Tensor, + weight_: torch.Tensor, + eps: float, + group_size: int = 128, + emit_bf16: bool = False, + ) -> Optional[Tuple[torch.Tensor, ...]]: + """Attempt fused all-reduce + RMSNorm + per-group FP8 quant. + + ROCm/aiter/gfx95-only entry point. Returns ``None`` on any other + platform or when the aiter custom-all-reduce communicator cannot + service the request, letting the caller fall back to the existing + ``fused_allreduce_rmsnorm`` + separate per-group quant path. + + When ``emit_bf16=True`` the fused kernel also writes the + pre-quantization bf16/fp16 normed output and returns + ``(fp8, residual_out, scale, bf16)`` — used by GDN-style layers that + need both an FP8 projection and a bf16 gating projection without + launching a separate per-group quant kernel. + """ + if not (is_hip() and is_gfx95_supported()): + return None + + ca_comm = self.ca_comm + if ca_comm is None or getattr(ca_comm, "disabled", True): + return None + if not hasattr(ca_comm, "custom_fused_ar_rms_per_group_quant"): + return None + + # Shape / size eligibility mirrors aiter's internal gate so we fail + # fast without entering the HIP kernel dispatch. + K = input_.shape[-1] + if K % group_size != 0 or K > 16384: + return None + total_bytes = input_.numel() * input_.element_size() + if total_bytes == 0 or total_bytes > 8 * 1024 * 8192: + return None + if self.world_size == 6: + return None + + if envs.SGLANG_USE_1STAGE_ALLREDUCE.is_set(): + use_1stage_ar = envs.SGLANG_USE_1STAGE_ALLREDUCE.get() + else: + token_num = input_.numel() // K + use_1stage_ar = total_bytes <= 128 * 1024 + if ( + # Keep the default 128 KiB cutoff except for the measured TP=8 + # K=7168 graph-replay crossover. K=4096 remains on the default + # rule because token_num=8/16 still favored 1-stage there. + self.world_size == 8 + and 4096 < K <= 7168 + and token_num >= 8 + and use_1stage_ar + ): + use_1stage_ar = False + + try: + return ca_comm.custom_fused_ar_rms_per_group_quant( + input_, + residual_inp_, + weight_, + eps, + group_size, + use_1stage_ar, + emit_bf16=emit_bf16, + ) + except Exception: + return None + def _all_reduce_out_place( self, input_: torch.Tensor, outplace_all_reduce_method: str ) -> torch.Tensor: diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 9faa78a28..ddc914c96 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -457,6 +457,8 @@ class LayerCommunicator: is_last_layer: bool = False, qkv_latent_func: Optional[Callable] = None, force_layernorm_before_dp_gather: bool = False, + enable_fused_ar_quant: bool = False, + fused_ar_quant_keep_bf16: bool = False, ): self.layer_scatter_modes = layer_scatter_modes self.input_layernorm = input_layernorm @@ -465,6 +467,8 @@ class LayerCommunicator: self.is_last_layer = is_last_layer self.qkv_latent_func = qkv_latent_func self.force_layernorm_before_dp_gather = force_layernorm_before_dp_gather + self.enable_fused_ar_quant = enable_fused_ar_quant + self.fused_ar_quant_keep_bf16 = fused_ar_quant_keep_bf16 self._context = CommunicateContext.init_new() self._context.force_layernorm_before_dp_gather = ( @@ -579,11 +583,32 @@ class LayerCommunicator: apply_aiter_all_reduce_fusion(hidden_states) or apply_flashinfer_allreduce_fusion(hidden_states.shape[0]) ) and hasattr(self.input_layernorm, "forward_with_allreduce_fusion"): - hidden_states, residual = ( - self.input_layernorm.forward_with_allreduce_fusion( - hidden_states, residual, use_attn_tp_group=False + quant_result = None + if ( + self.enable_fused_ar_quant + and _use_aiter + and hasattr( + self.input_layernorm, + "forward_with_allreduce_fusion_quant_per_group", + ) + ): + # Try fused AR+RMSNorm+per-group-quant. Internally + # falls back to AR+RMSNorm + separate quant when the + # fully-fused kernel cannot service the shape. + quant_result = self.input_layernorm.forward_with_allreduce_fusion_quant_per_group( + hidden_states, + residual, + use_attn_tp_group=False, + keep_bf16=self.fused_ar_quant_keep_bf16, + ) + if quant_result is not None: + hidden_states, residual = quant_result + else: + hidden_states, residual = ( + self.input_layernorm.forward_with_allreduce_fusion( + hidden_states, residual, use_attn_tp_group=False + ) ) - ) else: hidden_states = moe_tensor_model_parallel_all_reduce(hidden_states) hidden_states, residual = self.input_layernorm( diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index f8b02ed5a..8f7582a93 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -14,6 +14,7 @@ """Fused operators for normalization layers.""" import logging +from functools import lru_cache from typing import Optional, Tuple, Union import torch @@ -95,6 +96,7 @@ _has_aiter_layer_norm = False _has_vllm_rms_norm = False _has_rocm_triton_gemma_rms_norm = False if _use_aiter: + import aiter as _aiter from aiter import layernorm2d_fwd as layer_norm from aiter import rmsnorm2d_fwd as rms_norm from aiter import rmsnorm2d_fwd_with_add as fused_add_rms_norm @@ -156,6 +158,17 @@ if _is_npu: from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm +@lru_cache(maxsize=1) +def _get_aiter_per_group_quant(): + """Resolve aiter's per-1x128 HIP quant functor + FP8 dtype on first use. + + Memoized locally (rather than cached as module-level globals) so this + aiter-specific state stays out of layernorm.py's shared namespace and is + handed to callers as explicit values instead of being read implicitly. + """ + return _aiter.get_hip_quant(_aiter.QuantType.per_1x128), _aiter.dtypes.fp8 + + def _forward_with_allreduce_fusion( norm_module, x: torch.Tensor, @@ -213,6 +226,130 @@ def _forward_with_allreduce_fusion( return norm_module.forward(x, residual, post_residual_addition) +def _forward_with_allreduce_fusion_quant_per_group( + norm_module, + x: torch.Tensor, + residual: Optional[torch.Tensor], + weight: torch.Tensor, + group_size: int = 128, + use_attn_tp_group: bool = True, + keep_bf16: bool = False, +): + """Fused AR + RMSNorm + per-group FP8 quant with graceful staged fallback. + + The single-kernel quantized backend dispatch is ROCm + aiter + gfx95-only. + Other HIP/aiter runs can still use the 2-kernel fallback below to preserve + the existing tuple handoff behavior. + + The helper returns one of: + + * ``((fp8, scale), residual)`` when keep_bf16=False + * ``((bf16, fp8, scale), residual)`` when keep_bf16=True + * ``None`` when no fusion is possible + + Fallback chain (best → worst): + + 1. Fully-fused AR+RMSNorm+per-group-quant (aiter single kernel). + 2. Fused AR+RMSNorm followed by a separate per-1x128 quant + (two kernels, still saves the 3-kernel unfused baseline path). + 3. ``None`` so the caller can run the generic unfused path. + + ``keep_bf16`` is required for GDN-style layers that have one FP8 projection + (``in_proj_qkvz``) **and** one bf16 projection (``in_proj_ba``) on the + same normed output; without the bf16 we would have to dequantize which is + lossy. Standard attention layers (single FP8 ``qkv_proj``) use + ``keep_bf16=False``. + """ + if residual is None or not _use_aiter: + return None + + from sglang.srt.distributed import ( + tensor_model_parallel_fused_allreduce_rmsnorm, + tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group, + ) + from sglang.srt.layers.quantization.fp8_utils import ( + _use_aiter_bpreshuffle_gfx95 as use_bpreshuffle, + ) + from sglang.srt.layers.quantization.fp8_utils import ( + materialize_bpreshuffle_fp8_scale, + ) + + if use_attn_tp_group: + world_size = get_parallel().attn_tp_size + else: + if get_parallel().moe_ep_size > 1: + world_size = get_parallel().moe_ep_size + else: + world_size = get_parallel().moe_tp_size + if world_size <= 1: + return None + + # TODO: When ROCm/aiter#3652 is available in our bundled aiter, plumb + # transpose_scale=use_bpreshuffle into the fused AR+RMSNorm+quant kernel + # and drop this explicit post-kernel scale materialization. + if not keep_bf16: + result = tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group( + x, residual, weight, norm_module.variance_epsilon, group_size + ) + if result is not None: + fp8_out, residual_out, scale_out = result + if use_bpreshuffle: + scale_out = materialize_bpreshuffle_fp8_scale(scale_out) + return (fp8_out, scale_out), residual_out + + # Fallback: fused AR+RMSNorm then separate per-group quant. + fused_result = tensor_model_parallel_fused_allreduce_rmsnorm( + x, residual, weight, norm_module.variance_epsilon + ) + if fused_result is None: + return None + bf16_out, residual_out = fused_result + per_1x128_quant, fp8_dtype = _get_aiter_per_group_quant() + fp8_out, scale_out = per_1x128_quant( + bf16_out, + quant_dtype=fp8_dtype, + transpose_scale=False, + ) + if use_bpreshuffle: + scale_out = materialize_bpreshuffle_fp8_scale(scale_out) + return (fp8_out, scale_out), residual_out + + # keep_bf16=True: GDN path — need both an unquantized bf16 normed output + # (for in_proj_ba) AND (fp8, scale) (for in_proj_qkvz). Preferred path: + # use the fully-fused AR+RMSNorm+per-group-quant kernel with the optional + # bf16 side-output, so we avoid the separate per-group quant launch + # entirely. Fallback: fused AR+RMSNorm + separate per-group quant. + result = tensor_model_parallel_fused_allreduce_rmsnorm_quant_per_group( + x, + residual, + weight, + norm_module.variance_epsilon, + group_size, + emit_bf16=True, + ) + if result is not None and len(result) == 4: + fp8_out, residual_out, scale_out, bf16_out = result + if use_bpreshuffle: + scale_out = materialize_bpreshuffle_fp8_scale(scale_out) + return (bf16_out, fp8_out, scale_out), residual_out + + fused_result = tensor_model_parallel_fused_allreduce_rmsnorm( + x, residual, weight, norm_module.variance_epsilon + ) + if fused_result is None: + return None + bf16_out, residual_out = fused_result + per_1x128_quant, fp8_dtype = _get_aiter_per_group_quant() + fp8_out, scale_out = per_1x128_quant( + bf16_out, + quant_dtype=fp8_dtype, + transpose_scale=False, + ) + if use_bpreshuffle: + scale_out = materialize_bpreshuffle_fp8_scale(scale_out) + return (bf16_out, fp8_out, scale_out), residual_out + + class RMSNorm(MultiPlatformOp): def __init__( self, @@ -609,6 +746,25 @@ class RMSNorm(MultiPlatformOp): self, x, residual, post_residual_addition, self.weight, use_attn_tp_group ) + def forward_with_allreduce_fusion_quant_per_group( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + group_size: int = 128, + use_attn_tp_group: bool = True, + keep_bf16: bool = False, + ): + """Fused AR + RMSNorm + per-group FP8 quant (ROCm/aiter path). + + Returns ``((fp8, scale), residual)`` when ``keep_bf16=False``; + ``((bf16, fp8, scale), residual)`` when ``keep_bf16=True``; + or ``None`` when no fused path is available (caller must fall back to + the standard fused AR+RMSNorm + separate quant path). + """ + return _forward_with_allreduce_fusion_quant_per_group( + self, x, residual, self.weight, group_size, use_attn_tp_group, keep_bf16 + ) + class LayerNorm(MultiPlatformOp): def __init__( @@ -871,6 +1027,25 @@ class GemmaRMSNorm(MultiPlatformOp): use_attn_tp_group=True, ) + def forward_with_allreduce_fusion_quant_per_group( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + group_size: int = 128, + use_attn_tp_group: bool = True, + keep_bf16: bool = False, + ): + """Fused AR + RMSNorm + per-group FP8 quant (Gemma-style: weight + 1).""" + return _forward_with_allreduce_fusion_quant_per_group( + self, + x, + residual, + self.gemma_weight, + group_size, + use_attn_tp_group, + keep_bf16, + ) + class Gemma3RMSNorm(MultiPlatformOp): def __init__(self, dim: int, eps: float = 1e-6): diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index 19a6b6a95..dc6800a96 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -154,6 +154,56 @@ if _is_cpu: torch.ops.sgl_kernel.fused_qkvzba_split_reshape_cat_contiguous_cpu ) + +@lru_cache(maxsize=1) +def _enable_qwen35_fused_ar_quant() -> bool: + """Gate the fused AR+RMSNorm+per-group-FP8-quant path for Qwen3.5. + + The single-kernel backend is ROCm/aiter/gfx95-only. The model gate stays + tied to ROCm/aiter so non-gfx95 HIP can keep the existing 2-kernel fallback + behavior for tuple handoff when this branch is used. It replaces the + existing ``--enable-aiter-allreduce-fusion`` 3-kernel path + (AR → RMSNorm → per-group quant) with either a single fused kernel (when + the fully-fused variant is eligible) or a 2-kernel path + (fused AR+RMSNorm + separate per-group quant) that still saves one + kernel launch vs. baseline. The LayerCommunicator gracefully falls back + to ``forward_with_allreduce_fusion`` (plain AR+RMSNorm) when the fused + quant helper returns ``None``, so turning this on never regresses the + AR+RMSNorm fusion itself. + + Opt-out: set ``SGLANG_DISABLE_FUSED_AR_QUANT=1`` to fall back to the + unmodified AR+RMSNorm fusion path. + """ + if not _use_aiter: + return False + if get_bool_env_var("SGLANG_DISABLE_FUSED_AR_QUANT", default="false"): + return False + return bool(get_server_args().enable_aiter_allreduce_fusion) + + +def _linear_accepts_fp8_tuple(linear: nn.Module) -> bool: + quant_method = getattr(linear, "quant_method", None) + return quant_method.__class__.__name__ == "Fp8LinearMethod" and ( + getattr(quant_method, "block_quant", False) + or getattr(quant_method, "use_mxfp8", False) + ) + + +def _select_fused_ar_input_for_linear(hidden_states, linear: nn.Module): + if not isinstance(hidden_states, tuple): + return hidden_states + if len(hidden_states) == 3: + hs_bf16, hs_fp8, hs_scale = hidden_states + if _linear_accepts_fp8_tuple(linear): + return (hs_fp8, hs_scale) + return hs_bf16 + if len(hidden_states) == 2 and _linear_accepts_fp8_tuple(linear): + return hidden_states + raise TypeError( + f"{linear.__class__.__name__} cannot consume fused AR quant tuple input" + ) + + if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import ( split_qkvgate_gemma_rmsnorm_rope, @@ -487,6 +537,14 @@ class Qwen3_5GatedDeltaNet(nn.Module): return query, key, value, z, b, a def _forward_input_proj(self, hidden_states: torch.Tensor): + # AMD/aiter fused AR+RMSNorm+per-group-quant path ships a + # ``(bf16, fp8, scale)`` 3-tuple so the FP8 ``in_proj_qkvz`` can + # consume ``(fp8, scale)`` (skipping its internal quant) while the + # bf16 ``in_proj_ba`` consumes the unquantized bf16. Non-aiter runs skip + # the tuple branch and keep the original control flow below unchanged. + if _use_aiter and isinstance(hidden_states, tuple): + return self._forward_input_proj_fused_quant_amd(hidden_states) + if ( _is_cpu or _is_npu @@ -523,6 +581,39 @@ class Qwen3_5GatedDeltaNet(nn.Module): projected_states_ba, _ = self.in_proj_ba(hidden_states) return projected_states_qkvz, projected_states_ba + def _forward_input_proj_fused_quant_amd(self, hidden_states): + """AMD-only variant for the fused AR+RMSNorm+per-group-quant path. + + ``hidden_states`` is a ``(bf16, fp8, scale)`` 3-tuple produced by the + upstream fused kernel. FP8 ``in_proj_qkvz`` takes ``(fp8, scale)`` + directly; unquantized variants take the bf16 side-output. + """ + hs_bf16 = hidden_states[0] + hs_qkvz = _select_fused_ar_input_for_linear(hidden_states, self.in_proj_qkvz) + seq_len = hs_bf16.shape[0] + + if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE): + DUAL_STREAM_TOKEN_THRESHOLD = 0 + else: + DUAL_STREAM_TOKEN_THRESHOLD = 1024 + + if ( + self.alt_stream is not None + and get_is_capture_mode() + and seq_len < DUAL_STREAM_TOKEN_THRESHOLD + and _gdn_use_alt_stream + ): + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + projected_states_qkvz, _ = self.in_proj_qkvz(hs_qkvz) + with torch.cuda.stream(self.alt_stream): + projected_states_ba, _ = self.in_proj_ba(hs_bf16) + current_stream.wait_stream(self.alt_stream) + else: + projected_states_qkvz, _ = self.in_proj_qkvz(hs_qkvz) + projected_states_ba, _ = self.in_proj_ba(hs_bf16) + return projected_states_qkvz, projected_states_ba + def forward( self, hidden_states: torch.Tensor, @@ -661,12 +752,21 @@ class Qwen3_5LinearDecoderLayer(nn.Module): self.post_attention_layernorm = GemmaRMSNorm( config.hidden_size, eps=config.rms_norm_eps ) + # GDN layers need both bf16 (for the small in_proj_ba gating + # projection) and a quantized tuple only when in_proj_qkvz can consume + # it. Otherwise, stay on the plain AR+RMSNorm path. + enable_fused_ar_quant = ( + _enable_qwen35_fused_ar_quant() + and _linear_accepts_fp8_tuple(self.linear_attn.in_proj_qkvz) + ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=(layer_id == config.num_hidden_layers - 1), + enable_fused_ar_quant=enable_fused_ar_quant, + fused_ar_quant_keep_bf16=enable_fused_ar_quant, ) def forward( @@ -870,12 +970,19 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + # Standard attention layers benefit from a fused quant epilogue only + # when qkv_proj can consume the returned quantized tuple. + enable_fused_ar_quant = ( + _enable_qwen35_fused_ar_quant() and _linear_accepts_fp8_tuple(self.qkv_proj) + ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=(layer_id == config.num_hidden_layers - 1), + enable_fused_ar_quant=enable_fused_ar_quant, + fused_ar_quant_keep_bf16=False, ) self.alt_stream = alt_stream @@ -945,6 +1052,10 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): return q, k, v, gate def forward_prepare_native(self, positions, hidden_states): + if _use_aiter and isinstance(hidden_states, tuple): + hidden_states = _select_fused_ar_input_for_linear( + hidden_states, self.qkv_proj + ) qkv, _ = self.qkv_proj(hidden_states) if self.attn_output_gate: q_gate, k, v = qkv.split( @@ -964,6 +1075,10 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): return q, k, v, gate def forward_prepare_fused_gate(self, positions, hidden_states): + if _use_aiter and isinstance(hidden_states, tuple): + hidden_states = _select_fused_ar_input_for_linear( + hidden_states, self.qkv_proj + ) qkv, _ = self.qkv_proj(hidden_states) if self.attn_output_gate: q_gate, k, v = qkv.split( diff --git a/test/registered/amd/perf/mi35x/test_qwen35_fp8_ar_fusion_mi35x.py b/test/registered/amd/perf/mi35x/test_qwen35_fp8_ar_fusion_mi35x.py new file mode 100644 index 000000000..a60373316 --- /dev/null +++ b/test/registered/amd/perf/mi35x/test_qwen35_fp8_ar_fusion_mi35x.py @@ -0,0 +1,234 @@ +"""MI35x PR-CI accuracy coverage for Qwen3.5-FP8 aiter AR-fusion. + +PR-specific fused AR+RMSNorm+per-group-quant accuracy check. The file runs in +the 8-GPU MI35x stage-c suite and launches two TP4 servers in parallel: + +* GPUs 0-3: fused AR+RMSNorm+per-group FP8 quant enabled. +* GPUs 4-7: same launch with SGLANG_DISABLE_FUSED_AR_QUANT=1 fallback. + +Each server runs GSM8K and reports accuracy, invalid rate, latency, and output +throughput. This is the PR-CI accuracy signal; the nightly throughput/latency +perf benchmark lives separately in test_qwen35_fp8_perf_mi35x.py. +""" + +import os +import re +import subprocess +import unittest +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field +from pathlib import Path +from typing import Dict, List + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) +from sglang.utils import download_and_cache_file + +register_amd_ci(est_time=4800, suite="stage-c-test-large-8-gpu-amd-mi35x") + +QWEN35_FP8_MODEL_PATH = os.environ.get( + "QWEN35_FP8_MODEL_PATH", + "Qwen/Qwen3.5-397B-A17B-FP8", +) +SERVER_LAUNCH_TIMEOUT = 4800 +GSM8K_NUM_QUESTIONS = int(os.environ.get("GSM8K_NUM_QUESTIONS", "1319")) +ACCURACY_THRESHOLD = 0.94 + +# bench_sglang.py lives at the repo root (this file is 5 levels below), not under +# test/. Resolve it absolutely so it works regardless of the CI working directory. +REPO_ROOT = Path(__file__).resolve().parents[5] +GSM8K_BENCH_SCRIPT = REPO_ROOT / "benchmark" / "gsm8k" / "bench_sglang.py" +GSM8K_DATA_URL = ( + "https://raw.githubusercontent.com/openai/grade-school-math/" + "master/grade_school_math/data/test.jsonl" +) + + +@dataclass +class FusionVariant: + """A Qwen3.5-FP8 AR-fusion configuration to validate.""" + + variant: str + hip_visible_devices: str + port_offset: int + env_vars: Dict[str, str] = field(default_factory=dict) + + +COMMON_ARGS: List[str] = [ + "--tensor-parallel-size", + "4", + "--trust-remote-code", + "--attention-backend", + "aiter", + "--kv-cache-dtype", + "fp8_e4m3", + "--page-size", + "16", + "--chunked-prefill-size", + "8192", + "--mem-fraction-static", + "0.8", + "--disable-radix-cache", + "--enable-aiter-allreduce-fusion", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + "--watchdog-timeout", + "1200", +] + + +def _base_url_with_port_offset(offset: int) -> str: + host, port = DEFAULT_URL_FOR_TEST.rsplit(":", 1) + return f"{host}:{int(port) + offset}" + + +def get_fusion_variants() -> List[FusionVariant]: + return [ + FusionVariant( + variant="fused-ar-rms-per-group-quant", + hip_visible_devices="0,1,2,3", + port_offset=0, + env_vars={ + "SGLANG_USE_AITER": "1", + "SGLANG_USE_AITER_UNIFIED_ATTN": "1", + }, + ), + FusionVariant( + variant="disable-fused-ar-quant-opt-out", + hip_visible_devices="4,5,6,7", + port_offset=1, + env_vars={ + "SGLANG_USE_AITER": "1", + "SGLANG_USE_AITER_UNIFIED_ATTN": "1", + "SGLANG_DISABLE_FUSED_AR_QUANT": "1", + }, + ), + ] + + +def _parse_gsm8k_metrics(stdout: str) -> Dict[str, float]: + metrics = {} + for key, pattern in { + "accuracy": r"Accuracy:\s*([0-9.]+)", + "invalid": r"Invalid:\s*([0-9.]+)", + "latency": r"Latency:\s*([0-9.]+)\s*s", + "output_throughput": r"Output throughput:\s*([0-9.]+)\s*token/s", + }.items(): + match = re.search(pattern, stdout) + if match is None: + raise AssertionError(f"Could not parse {key} from GSM8K output:\n{stdout}") + metrics[key] = float(match.group(1)) + return metrics + + +class TestQwen35Fp8ArFusionMI35x(CustomTestCase): + """Validate Qwen3.5-FP8 AR-fusion accuracy and throughput on MI35x.""" + + @classmethod + def setUpClass(cls): + cls.model = QWEN35_FP8_MODEL_PATH + cls.variants = get_fusion_variants() + # Pre-fetch the dataset once (single-threaded) so the two parallel + # benchmark subprocesses don't race writing the shared /tmp cache file. + cls.gsm8k_data_path = download_and_cache_file(GSM8K_DATA_URL) + + def _run_gsm8k(self, base_url: str) -> Dict[str, float]: + port = int(base_url.rsplit(":", 1)[-1]) + command = [ + "python3", + str(GSM8K_BENCH_SCRIPT), + "--num-questions", + str(GSM8K_NUM_QUESTIONS), + "--parallel", + str(GSM8K_NUM_QUESTIONS), + "--num-shots", + "5", + "--data-path", + str(self.gsm8k_data_path), + "--port", + str(port), + ] + result = subprocess.run(command, capture_output=True, text=True) + if result.returncode != 0: + raise AssertionError( + "GSM8K benchmark failed:\n" + f"stdout:\n{result.stdout}\n" + f"stderr:\n{result.stderr}" + ) + print(result.stdout) + return _parse_gsm8k_metrics(result.stdout) + + def _run_variant(self, variant: FusionVariant) -> Dict[str, float]: + env = os.environ.copy() + env["HIP_VISIBLE_DEVICES"] = variant.hip_visible_devices + env.update(variant.env_vars) + base_url = _base_url_with_port_offset(variant.port_offset) + + process = popen_launch_server( + self.model, + base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=list(COMMON_ARGS), + env=env, + ) + try: + requests.get(base_url + "/flush_cache", timeout=10) + metrics = self._run_gsm8k(base_url) + print(f"[{variant.variant}] {metrics=}") + return metrics + finally: + kill_process_tree(process.pid) + + def test_qwen35_fp8_ar_fusion_accuracy_and_perf(self): + summary = "### Qwen3.5-FP8 aiter AR-fusion (MI35x, parallel TP4)\n\n" + summary += ( + "| Variant | GPUs | Accuracy | Invalid | Latency (s) | Output tok/s | " + "Threshold | Status |\n" + ) + summary += "| ------- | ---- | -------- | ------- | ----------- | ------------ | --------- | ------ |\n" + + failures = [] + with ThreadPoolExecutor(max_workers=len(self.variants)) as executor: + future_to_variant = { + executor.submit(self._run_variant, variant): variant + for variant in self.variants + } + for future in as_completed(future_to_variant): + variant = future_to_variant[future] + with self.subTest(variant=variant.variant): + metrics = future.result() + accuracy = metrics["accuracy"] + passed = accuracy >= ACCURACY_THRESHOLD + status = "PASS" if passed else "FAIL" + summary += ( + f"| {variant.variant} | {variant.hip_visible_devices} | " + f"{accuracy:.3f} | {metrics['invalid']:.3f} | " + f"{metrics['latency']:.2f} | " + f"{metrics['output_throughput']:.2f} | " + f"{ACCURACY_THRESHOLD} | {status} |\n" + ) + if not passed: + failures.append((variant.variant, accuracy)) + + if is_in_ci(): + write_github_step_summary(summary) + print(summary) + + self.assertEqual( + failures, + [], + f"Qwen3.5-FP8 AR-fusion accuracy below {ACCURACY_THRESHOLD}: {failures}", + ) + + +if __name__ == "__main__": + unittest.main()