sgl-project/sglang

[Feature] Support PDL on norm in sgl-kernel

Geschlossen

#5.946 geöffnet am 01.05.2025

 (4 Kommentare) (2 Reaktionen) (1 zugewiesene Person)Python (6.216 Forks)auto 404
good first issuesgl-kernel

Repository-Metriken

Stars
 (28.442 Sterne)
PR-Merge-Metriken
 (Durchschn. Merge 1T 13h) (1.000 gemergte PRs in 30 T)

Beschreibung

Checklist

Motivation

In previous versions, we updated flashinfer. Flashinfer 0.2.5 supports norm's PDL, but currently, norm's PDL is disabled by default. We would like to modify the code to enable it.

Related resources

We need change code at sgl-kernel/python/sgl_kernel, those who have enable_pdl parameter.

For example:

def rmsnorm(
    input: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    out: Optional[torch.Tensor] = None,
    enable_pdl: bool = False,
) -> torch.Tensor:
    r"""Root mean square normalization.

    ``out[i] = (input[i] / RMS(input)) * weight[i]``

    Parameters
    ----------
    input: torch.Tensor
        Input tensor, shape (batch_size, hidden_size).
    weight: torch.Tensor
        Weight tensor, shape (hidden_size,).
    eps: float
        Epsilon for numerical stability.
    out: Optional[torch.Tensor]
        The output tensor, if specified, the kernel will update this tensor inplace.
    enable_pdl: bool
        Whether to enable `programmatic dependent launch
        <https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#programmatic-dependent-launch-and-synchronization>`_

    Returns
    -------
    output: torch.Tensor
        Normalized tensor, shape (batch_size, hidden_size).
    """
    if out is None:
        out = torch.empty_like(input)
    torch.ops.sgl_kernel.rmsnorm.default(out, input, weight, eps, enable_pdl)
    return out

This is just for example, we have bunch of API need to enhance.

Whats is PDL:

https://github.com/NVIDIA/cutlass/discussions/1791

So we need add a utils function for hopper arch, and use PDL automatically.

Contributor Guide