vllm.kernels.triton.fused_pre_fc_norms

Functions:

• fused_pre_fc_norms –

  Returns cat([norm_e(inputs_embeds), norm_h(hidden_states)]) -> [N, 2H].

fused_pre_fc_norms(inputs_embeds, hidden_states, emb_weight, hid_weight, eps)

Returns cat([norm_e(inputs_embeds), norm_h(hidden_states)]) -> [N, 2H].

Both norms are Gemma-style RMSNorm (scale by 1 + weight, fp32 accumulation, multiply in fp32 before the cast back to the input dtype).

Source code in vllm/kernels/triton/fused_pre_fc_norms.py

def fused_pre_fc_norms(
    inputs_embeds: torch.Tensor,
    hidden_states: torch.Tensor,
    emb_weight: torch.Tensor,
    hid_weight: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """Returns cat([norm_e(inputs_embeds), norm_h(hidden_states)]) -> [N, 2H].

    Both norms are Gemma-style RMSNorm (scale by ``1 + weight``, fp32
    accumulation, multiply in fp32 before the cast back to the input dtype).
    """
    n, h = inputs_embeds.shape
    out = torch.empty(n, 2 * h, dtype=inputs_embeds.dtype, device=inputs_embeds.device)
    _fused_pre_fc_norms_kernel[(n,)](
        inputs_embeds,
        inputs_embeds.stride(0),
        hidden_states,
        hidden_states.stride(0),
        emb_weight,
        hid_weight,
        eps,
        out,
        out.stride(0),
        h,
        triton.next_power_of_2(h),
    )
    return out