sgl-project/sglang

[Feature] Support PDL on norm in sgl-kernel

クローズ

#5,946 opened on 2025/05/01

 (4 件のコメント) (2 件のリアクション) (1 人の担当者)Python (6,216 件のフォーク)auto 404
good first issuesgl-kernel

Repository metrics

Stars
 (28,442 個のスター)
PR merge metrics
 (平均マージ 2d 1h) (30d で 1,000 merged PRs)

説明

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.

コントリビューターガイド