sgl-project/sglang

[Feature] Support PDL on norm in sgl-kernel

Chiusa

#5946 aperta il 1 mag 2025

 (4 commenti) (2 reazioni) (1 assegnatario)Python (6216 fork)auto 404
good first issuesgl-kernel

Metriche repository

Star
 (28.442 stelle)
Metriche merge PR
 (Merge medio 2g 1h) (1000 PR mergiate in 30 g)

Descrizione

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.

Guida contributor