[AMD] Add fused all-reduce RMSNorm per-group quant for Qwen3.5 FP8 (#24651)
Co-authored-by: jacky.cheng <yichiche@amd.com> Co-authored-by: yctseng0211 <yctseng@amd.com> Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Hubert Lu
co-authored by
jacky.cheng
yctseng0211
HAI
jacky.cheng
yctseng0211
HAI
parent
b855efd9e6
commit
e8e765b9d6
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user