minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Add fused all-reduce RMSNorm per-group quant for Qwen3.5 FP8 (#24651)

Browse Source 
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
2026-07-22 07:33:03 -07:00
committed by GitHub
co-authored by jacky.cheng yctseng0211 HAI
parent b855efd9e6
commit e8e765b9d6
9 changed files with 1192 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/communication_op.py
+23
View File
@@ -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:
python/sglang/srt/distributed/parallel_state.py
+72
View File
@@ -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:
python/sglang/srt/layers/communicator.py
+29 -4
View File
@@ -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(
python/sglang/srt/layers/layernorm.py
+175
View File
@@ -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):
python/sglang/srt/models/qwen3_5.py
+115
View File
@@ -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(