sgl-project--sglang/python/sglang/jit_kernel/all_reduce.py

from __future__ import annotations

import enum
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast

import torch
import tvm_ffi
from tvm_ffi import Module

from sglang.jit_kernel.utils import (
    cache_once,
    is_arch_support_pdl,
    load_jit,
    make_cpp_args,
)
from sglang.kernel_api_logging import debug_kernel_api


class ConfigResult(NamedTuple):
    num_blocks: int
    num_threads: int


class AllReduceAlgo(enum.Enum):
    ONE_SHOT_PUSH = enum.auto()
    ONE_SHOT_PULL = enum.auto()
    TWO_SHOT_PULL = enum.auto()

    def is_push(self) -> bool:
        return self == AllReduceAlgo.ONE_SHOT_PUSH

    @property
    def shot(self) -> int:
        return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1


if TYPE_CHECKING:
    CUSTOM_AR_HANDLE = List[int]
    CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE]

    class CustomAllReduceObj:
        def __init__(
            self,
            rank: int,
            world_size: int,
            pull_buffer_bytes: int,
            push_buffer_bytes: int,
            graph_input_count: int,
            *,
            max_pull_blocks: Optional[int] = None,
            max_push_blocks: Optional[int] = None,
        ) -> None:
            """
            Create a CustomAllReduceObj instance.

            :param rank: The rank of the current process.
            :param world_size: The total number of processes in the group.
            :param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce.
            :param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce.
            :param graph_input_count: The maximum number of inputs in all CUDA graphs.
            :param max_pull_blocks: The maximum number of thread blocks to launch for pull-based all-reduce.
                                    If None, it will be determined by the implementation.
            :param max_push_blocks: The maximum number of thread blocks to launch for push-based all-reduce.
                                    If None, it will be determined by the implementation.
            """

        @property
        def world_size(self) -> int: ...
        def share_storage(self) -> CUSTOM_AR_HANDLE: ...
        def share_graph_inputs(self) -> List[CUSTOM_AR_PAIR]: ...
        def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ...
        def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ...
        def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...
        def get_graph_capture_bases(
            self,
        ) -> Tuple[List[Tuple[int, int]], List[List[int]], List[int]]: ...
        def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ...
        def all_reduce(
            self, input: torch.Tensor, algo: AllReduceAlgo
        ) -> tvm_ffi.Tensor: ...
        def config_pull(
            self, num_blocks: int = -1, num_threads: int = -1
        ) -> ConfigResult:
            """
            Configure the CUDA kernel's grid and block dimensions.
            This provides only the upper bound of the configuration,
            and the actual launch configuration may be determined by implementation.
            Note that push-based all-reduce can not be configured currently.

            :param num_blocks: The maximum number of thread blocks to launch. -1 means no limit.
            :param num_threads: The maximum number of threads per block. -1 means no limit.

            :return: The previous configuration as a ConfigResult named tuple.
            """
            ...


@cache_once
def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int) -> Module:
    args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
    return load_jit(
        "custom_all_reduce_pull",
        *args,
        extra_ldflags=["-lcuda"],
        cuda_files=["distributed/custom_all_reduce_pull.cuh"],
        cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
    )


@cache_once
def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int) -> Module:
    args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
    return load_jit(
        "custom_all_reduce_push",
        *args,
        extra_ldflags=["-lcuda"],
        cuda_files=["distributed/custom_all_reduce_push.cuh"],
        cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
    )


@cache_once
def _jit_fused_parallel_qknorm_module(
    dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
) -> Module:
    args = make_cpp_args(dtype, world_size, q_dim, k_dim, is_arch_support_pdl())
    cls_name = f"FusedParallelQKNormAcrossHead<{args}>"
    return load_jit(
        "tp_qknorm",
        *args,
        extra_ldflags=["-lcuda"],
        cuda_files=["distributed/tp_qknorm.cuh"],
        cuda_wrappers=[
            ("fused_parallel_qknorm", f"{cls_name}::run"),
            ("get_max_occupancy", f"{cls_name}::get_max_occupancy"),
        ],
    )


@cache_once
def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
    module = load_jit(
        "custom_all_reduce_base",
        extra_ldflags=["-lcuda"],
        cuda_files=["distributed/custom_all_reduce_base.cuh"],
        cuda_wrappers=[("register_once", "register_custom_all_reduce")],
    )
    module.register_once()
    device = torch.cuda.current_device()
    props = torch.cuda.get_device_properties(device)
    NUM_CTA = props.multi_processor_count
    MAX_THREADS = 512

    @tvm_ffi.register_object("sgl.CustomAllReduce")
    class CustomAllReduceObjReal(tvm_ffi.Object):
        __slots__ = ("__dict__",)

        def __init__(
            self,
            rank: int,
            world_size: int,
            pull_buffer_bytes: int,
            push_buffer_bytes: int,
            graph_input_count: int,
            *,
            max_pull_blocks: Optional[int] = None,
            max_push_blocks: Optional[int] = None,
        ) -> None:
            max_pull_blocks = NUM_CTA if max_pull_blocks is None else max_pull_blocks
            max_push_blocks = NUM_CTA if max_push_blocks is None else max_push_blocks
            self.__ffi_init__(
                rank,
                world_size,
                max_pull_blocks,
                max_push_blocks,
                pull_buffer_bytes,
                push_buffer_bytes,
                graph_input_count,
            )
            self._world_size = world_size
            self._pull_config = ConfigResult(min(NUM_CTA, max_pull_blocks), MAX_THREADS)
            if max_pull_blocks > 0:  # special case: cannot configure 0 blocks
                self.configure_pull(*self._pull_config)  # type: ignore

        @property
        def world_size(self) -> int:
            return self._world_size

        @debug_kernel_api
        def all_reduce(
            self,
            input: torch.Tensor,
            algo: AllReduceAlgo,
        ) -> tvm_ffi.Tensor:
            compile_fn = (
                _jit_custom_all_reduce_push_module
                if algo.is_push()
                else _jit_custom_all_reduce_pull_module
            )
            module = compile_fn(input.dtype, self._world_size)
            return module.all_reduce(self, input, algo.shot)

        def config_pull(
            self, num_blocks: int = -1, num_threads: int = -1
        ) -> ConfigResult:
            old_config = self._pull_config
            num_blocks = num_blocks if num_blocks != -1 else old_config.num_blocks
            num_threads = num_threads if num_threads != -1 else old_config.num_threads
            new_config = ConfigResult(num_blocks, num_threads)
            if new_config != old_config:
                result = ConfigResult(*self.configure_pull(*new_config))  # type: ignore
                assert result == self._pull_config
                self._pull_config = new_config
            return old_config

        def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None:
            self.free_ipc_handles()  # type: ignore
            torch.distributed.barrier(group=tp_cpu_group)
            self.free_storage()  # type: ignore

    return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal)


def get_fused_parallel_qknorm_max_occupancy(
    dtype: torch.dtype, world_size: int, q_dim: int, k_dim: int
) -> int:
    module = _jit_fused_parallel_qknorm_module(dtype, world_size, q_dim, k_dim)
    return module.get_max_occupancy()


def fused_parallel_qknorm(
    custom_ar: CustomAllReduceObj,
    q: torch.Tensor,
    k: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float = 1e-6,
) -> None:
    world_size = custom_ar.world_size
    q_dim = q.shape[-1] * world_size
    k_dim = k.shape[-1] * world_size
    module = _jit_fused_parallel_qknorm_module(q.dtype, world_size, q_dim, k_dim)
    module.fused_parallel_qknorm(custom_ar, q, k, q_weight, k_weight, eps)