vllm.models.inkling.nvidia.ops.norm ¶

def add_rmsnorm(
    residual: torch.Tensor,
    delta: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Fused ``res = residual + delta; y = rmsnorm(res)``.

    Returns ``(y, res)``; both are fresh tensors (cudagraph-friendly, no
    in-place update of the inputs).
    """
    assert residual.ndim == 2 and delta.ndim == 2, (residual.shape, delta.shape)
    n_rows, n_cols = residual.shape
    assert weight.shape[0] == n_cols
    y = torch.empty_like(residual)
    res_out = torch.empty_like(residual)
    if n_rows == 0:
        return y, res_out

    block_size_n = triton.next_power_of_2(n_cols)
    max_block_size_n = _MAX_FUSED_SIZE // residual.element_size()
    if max_block_size_n < block_size_n:
        raise RuntimeError(f"Large {n_cols=} is not supported")
    num_warps = _get_num_warps_from_block_size(block_size_n)
    _add_rmsnorm_fwd_kernel[(n_rows,)](
        residual,
        delta,
        weight,
        y,
        res_out,
        eps,
        residual.stride(0),
        delta.stride(0),
        y.stride(0),
        res_out.stride(0),
        n_cols,
        block_size_n,
        num_warps=num_warps,
    )
    return y, res_out

def embed_rmsnorm(
    input_ids: torch.Tensor,
    embed_table: torch.Tensor,
    weight: torch.Tensor | None,
    eps: float,
    chain_weight: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """Fused ``rmsnorm(embed_table[input_ids], weight)`` row gather + norm.

    Requires the full vocab on-rank (replicated or tp_size == 1).
    ``weight=None`` skips the norm (``use_embed_norm=False``), leaving a pure
    embedding-table gather.
    ``chain_weight`` additionally emits ``rmsnorm(out, chain_weight)`` (the
    first decoder layer's pre-attention norm) as a second output, still one
    launch. Bit-exact vs the unfused module sequence."""
    ids = input_ids.view(-1)
    (T,) = ids.shape
    n = embed_table.shape[1]
    out = torch.empty(
        (*input_ids.shape, n), dtype=embed_table.dtype, device=embed_table.device
    )
    chain_out = torch.empty_like(out) if chain_weight is not None else None
    if T > 0:
        block_size_n = triton.next_power_of_2(n)
        _embed_rmsnorm_kernel[(T,)](
            ids,
            embed_table,
            weight if weight is not None else embed_table,
            chain_weight if chain_weight is not None else embed_table,
            out,
            chain_out if chain_out is not None else out,
            eps,
            embed_table.stride(0),
            n,
            block_size_n,
            HAS_NORM=weight is not None,
            HAS_CHAIN=chain_weight is not None,
            num_warps=_get_num_warps_from_block_size(block_size_n),
        )
    if chain_out is not None:
        return out, chain_out
    return out