vllm.model_executor.layers.fused_allreduce_gemma_rms_norm

Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm.

Under tensor parallelism a RowParallelLinear (e.g. attention o_proj) produces a per-rank partial sum that is all-reduced, and the result is then fed into a GemmaRMSNorm that adds the residual and normalizes. flashinfer ships a kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this helper drives it directly (no torch.compile pass) for models that run eager.

Scope: attention output only, no quantization. When the flashinfer fast path is not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an oversize batch) it falls back to all_reduce + GemmaRMSNorm, which is numerically identical to the unfused model path.

Functions:

• allreduce_gemma_rms_norm_mxfp8 –

  All-reduce + residual add + GemmaRMSNorm with MXFP8-quantized output.

• fused_allreduce_gemma_rms_norm –

  All-reduce hidden_states + add residual + GemmaRMSNorm, fused.

_can_use_flashinfer(hidden_states, tp_size)

Whether the flashinfer fused path applies; returns (ok, max_token_num).

Source code in vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
    """Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
    if (
        flashinfer_trtllm_fused_allreduce_norm is None
        or get_fi_ar_workspace is None
        or _AR_RESIDUAL_RMS_NORM is None
    ):
        return False, 0
    if (
        not hidden_states.is_cuda
        or hidden_states.dim() != 2
        or not hidden_states.is_contiguous()
        or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
    ):
        return False, 0

    num_tokens, hidden_size = hidden_states.shape
    max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
    if max_token_num is None or num_tokens > max_token_num:
        return False, 0

    # Lazily create / fetch the (globally cached) workspace; returns None on
    # GPUs without NVSwitch, in which case we fall back gracefully.
    workspace = get_fi_ar_workspace(
        world_size=tp_size,
        rank=get_tensor_model_parallel_rank(),
        max_token_num=max_token_num,
        hidden_dim=hidden_size,
        dtype=hidden_states.dtype,
        group=get_tp_group().cpu_group,
    )
    if workspace is None:
        return False, 0
    # The token-count bound above uses the whole workspace budget, but a backend
    # may use only a fraction of it per call (mnnvl rotates through three Lamport
    # buffers). Ask the workspace so we don't admit tensors that the kernel will
    # reject with "The buffer size in the given workspace is insufficient".
    if not workspace.is_buffer_size_sufficient(
        tp_size=tp_size,
        num_tokens=num_tokens,
        hidden_dim=hidden_size,
        dtype=hidden_states.dtype,
    ):
        return False, 0
    return True, max_token_num

_max_token_num(tp_size, hidden_size, dtype)

Workspace token budget for flashinfer fused all-reduce, or None if the current world size / device is unsupported. Mirrors FlashInferAllReduce.

Source code in vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
    """Workspace token budget for flashinfer fused all-reduce, or None if the
    current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
    from vllm.config.compilation import PassConfig

    max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
    if not max_size_mb:
        return None
    element_size = torch.tensor([], dtype=dtype).element_size()
    return int(max_size_mb * MiB) // (hidden_size * element_size)

allreduce_gemma_rms_norm_mxfp8(hidden_states, residual, norm, return_normed=False)

All-reduce + residual add + GemmaRMSNorm with MXFP8-quantized output.

fused_allreduce_gemma_rms_norm followed by an MXFP8 activation quant, with the quant as an epilogue of the Triton norm instead of its own launch. The fused all-reduce + norm kernels have no quantized output, so this only covers the plain collective + Triton norm; when one of them would run it returns None and the caller keeps the unfused path.

Returns (x_q, x_scale, normed, residual_out): float8_e4m3fn values [M, N], row-major E8M0 scales [M, N // 32], the bf16 normed rows (only if return_normed, else None) and the new residual.

Source code in vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

def allreduce_gemma_rms_norm_mxfp8(
    hidden_states: torch.Tensor,
    residual: torch.Tensor,
    norm: GemmaRMSNorm,
    return_normed: bool = False,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor] | None:
    """All-reduce + residual add + GemmaRMSNorm with MXFP8-quantized output.

    ``fused_allreduce_gemma_rms_norm`` followed by an MXFP8 activation quant,
    with the quant as an epilogue of the Triton norm instead of its own launch.
    The fused all-reduce + norm kernels have no quantized output, so this only
    covers the plain collective + Triton norm; when one of them would run it
    returns ``None`` and the caller keeps the unfused path.

    Returns ``(x_q, x_scale, normed, residual_out)``: float8_e4m3fn values
    ``[M, N]``, row-major E8M0 scales ``[M, N // 32]``, the bf16 normed rows
    (only if ``return_normed``, else ``None``) and the new residual.
    """
    from vllm.models.minimax_m3.amd.ops import (
        gemma_fused_add_rmsnorm_mxfp8,
        gemma_fused_add_rmsnorm_mxfp8_and_bf16,
    )

    tp_size = get_tensor_model_parallel_world_size()
    if tp_size > 1:
        if _can_use_flashinfer(hidden_states, tp_size)[0]:
            return None
        if _can_use_aiter_fused_ar_rms(hidden_states):
            return None
        hidden_states = tensor_model_parallel_all_reduce(hidden_states)
    if return_normed:
        return gemma_fused_add_rmsnorm_mxfp8_and_bf16(
            hidden_states, residual, norm.weight, norm.variance_epsilon
        )
    x_q, x_scale, res_out = gemma_fused_add_rmsnorm_mxfp8(
        hidden_states, residual, norm.weight, norm.variance_epsilon
    )
    return x_q, x_scale, None, res_out

fused_allreduce_gemma_rms_norm(hidden_states, residual, norm)

All-reduce hidden_states + add residual + GemmaRMSNorm, fused.

hidden_states is the per-rank partial (un-reduced) output of a row-parallel linear; norm is the GemmaRMSNorm applied right after. Returns (normed_output, new_residual), equivalent to norm(all_reduce(hidden_states), residual).

Source code in vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py

def fused_allreduce_gemma_rms_norm(
    hidden_states: torch.Tensor,
    residual: torch.Tensor,
    norm: GemmaRMSNorm,
) -> tuple[torch.Tensor, torch.Tensor]:
    """All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused.

    ``hidden_states`` is the per-rank *partial* (un-reduced) output of a
    row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after.
    Returns ``(normed_output, new_residual)``, equivalent to
    ``norm(all_reduce(hidden_states), residual)``.
    """
    tp_size = get_tensor_model_parallel_world_size()
    if tp_size == 1:
        # No all-reduce needed; identical to the unfused path.
        return norm(hidden_states, residual)

    ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
    if ok:
        norm_out = torch.empty_like(hidden_states)
        # With norm_out provided, the kernel writes the new residual
        # (all_reduce(hidden_states) + residual) into the hidden_states buffer
        # and the normalized result into norm_out, leaving `residual` untouched.
        flashinfer_trtllm_fused_allreduce_norm(
            allreduce_in=hidden_states,
            residual=residual,
            rms_gamma=norm.weight,
            rms_eps=norm.variance_epsilon,
            world_size=tp_size,
            weight_bias=1.0,  # GemmaRMSNorm-style
            launch_with_pdl=True,
            fp32_acc=True,
            max_token_num=max_token_num,
            pattern_code=_AR_RESIDUAL_RMS_NORM,
            norm_out=norm_out,
        )
        return norm_out, hidden_states

    if _can_use_aiter_fused_ar_rms(hidden_states):
        from vllm._aiter_ops import rocm_aiter_ops

        return rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()(
            input_=hidden_states,
            residual=residual,
            weight=norm.weight,
            epsilon=norm.variance_epsilon,
            gemma_norm=True,
        )

    # Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
    reduced = tensor_model_parallel_all_reduce(hidden_states)
    return norm(reduced, residual)