sgl-project/sglang

[Feature] Support PDL on norm in sgl-kernel

Fermée

#5 946 ouverte le 1 mai 2025

 (4 commentaires) (2 réactions) (1 personne assignée)Python (6 216 forks)auto 404
good first issuesgl-kernel

Métriques du dépôt

Stars
 (28 442 étoiles)
Métriques de merge PR
 (Merge moyen 2j 1h) (1 000 PRs mergées en 30 j)

Description

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.

Guide contributeur