sgl-project/sglang

[Feature] Support PDL on norm in sgl-kernel

已关闭

#5,946 创建于 2025年5月1日

 (4 条评论) (2 个反应) (1 位负责人)Python (6,216 个派生)auto 404
good first issuesgl-kernel

仓库指标

星标
 (28,442 个星标)
PR 合并指标
 (平均合并 2天 1小时) (30 天内合并 1,000 个 PR)

描述

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.

贡献者指南