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_dual_rmsnorm_cat(
    hidden: torch.Tensor,
    hidden_weight: torch.Tensor,
    embed_weight: torch.Tensor,
    eps: float,
    *,
    embeds: torch.Tensor | None = None,
    input_ids: torch.Tensor | None = None,
    embed_table: torch.Tensor | None = None,
    pre_norm_weight: torch.Tensor | None = None,
) -> torch.Tensor:
    """The MTP depth-layer input in one launch:
    ``cat([rmsnorm(hidden, w_h), rmsnorm(pre?(emb), w_e)], -1)``.

    The embed side is either a fused row gather ``embed_table[input_ids]``
    (draft decode steps) or precomputed ``embeds`` ([T, N], the target-merged
    multimodal embeddings at draft prefill); ``pre_norm_weight`` chains the
    backbone embed_norm in front of the depth embed_norm (bit-exact vs the
    unfused sequence). The concat copies collapse into direct writes."""
    T, n = hidden.shape
    if embeds is not None:
        assert embeds.shape == hidden.shape
        src, ids, src_stride = embeds, embeds, embeds.stride(0)
        gather = False
    else:
        assert input_ids is not None and embed_table is not None
        assert input_ids.shape == (T,) and embed_table.shape[1] == n
        src, ids, src_stride = embed_table, input_ids, embed_table.stride(0)
        gather = True
    out = torch.empty((T, 2 * n), dtype=hidden.dtype, device=hidden.device)
    if T == 0:
        return out
    block_size_n = triton.next_power_of_2(n)
    _embed_dual_rmsnorm_cat_kernel[(T, 2)](
        hidden,
        src,
        ids,
        hidden_weight,
        pre_norm_weight if pre_norm_weight is not None else embed_weight,
        embed_weight,
        out,
        eps,
        hidden.stride(0),
        src_stride,
        n,
        block_size_n,
        GATHER=gather,
        HAS_PRE_NORM=pre_norm_weight is not None,
        num_warps=_get_num_warps_from_block_size(block_size_n),
    )
    return 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