diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index 42c9f0c09..6f38d3952 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -123,7 +123,7 @@ srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", - "torchada>=0.1.55", + "torchada>=0.1.56", "mthreads-ml-py", "mate>=0.2.0", "deep-gemm>=0.1.3", diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 494b1417e..c074868e4 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -115,7 +115,7 @@ srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", - "torchada>=0.1.55", + "torchada>=0.1.56", "mthreads-ml-py", "mate>=0.2.0", "deep-gemm>=0.1.3", diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index caa468b9f..9181fbac5 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -20,6 +20,7 @@ from sglang.srt.utils.common import ( crash_on_warnings, get_bool_env_var, is_cuda, + is_musa, is_npu, ) @@ -33,6 +34,15 @@ if is_cuda(): top_p_renorm_prob, ) +if is_musa(): + from sgl_kernel import ( + min_p_sampling_from_probs, + top_k_renorm_prob, + top_k_top_p_sampling_from_probs, + top_p_renorm_prob, + ) + + if is_npu(): import torch_npu diff --git a/sgl-kernel/csrc/common_extension_musa.cc b/sgl-kernel/csrc/common_extension_musa.cc index bcb3a8893..c59068c8f 100644 --- a/sgl-kernel/csrc/common_extension_musa.cc +++ b/sgl-kernel/csrc/common_extension_musa.cc @@ -278,6 +278,16 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()"); m.impl("top_p_renorm_probs", torch::kMUSA, &top_p_renorm_probs); + m.def( + "min_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_min_p_arr, float " + "min_p_val, bool deterministic, Generator? gen) -> ()"); + m.impl("min_p_sampling_from_probs", torch::kMUSA, &min_p_sampling_from_probs); + + m.def( + "top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_p_arr, " + "float top_p_val, bool deterministic, Generator? gen) -> ()"); + m.impl("top_p_sampling_from_probs", torch::kMUSA, &top_p_sampling_from_probs); + /* * From csrc/musa */ diff --git a/sgl-kernel/include/sgl_kernel_musa_ops.h b/sgl-kernel/include/sgl_kernel_musa_ops.h index a50cd7fc4..6ce91a76b 100644 --- a/sgl-kernel/include/sgl_kernel_musa_ops.h +++ b/sgl-kernel/include/sgl_kernel_musa_ops.h @@ -78,3 +78,21 @@ void musa_top_k_top_p_sampling_from_probs( double top_p_val, bool deterministic, std::optional gen); + +void min_p_sampling_from_probs( + at::Tensor probs, + at::Tensor output, + std::optional maybe_indices, + std::optional maybe_min_p_arr, + double min_p_val, + bool deterministic, + std::optional gen); + +void top_p_sampling_from_probs( + at::Tensor probs, + at::Tensor output, + std::optional maybe_indices, + std::optional maybe_top_p_arr, + double top_p_val, + bool deterministic, + std::optional gen); diff --git a/sgl-kernel/pyproject_musa.toml b/sgl-kernel/pyproject_musa.toml index 58b11253f..c9b790760 100644 --- a/sgl-kernel/pyproject_musa.toml +++ b/sgl-kernel/pyproject_musa.toml @@ -3,7 +3,7 @@ requires = [ "setuptools>=75.0", "scikit-build-core>=0.10", "torch", - "torchada>=0.1.55", + "torchada>=0.1.56", "wheel", ] build-backend = "setuptools.build_meta" diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index 73ec89100..1b97271f2 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -128,11 +128,13 @@ else: if hasattr(torch.version, "musa") and torch.version.musa is not None: from sgl_kernel.musa import ( + min_p_sampling_from_probs, musa_batched_rotary_embedding_contiguous, musa_fused_gemv, musa_fused_moe_gemv, musa_fused_mul_add, musa_rotary_embedding_contiguous, + top_k_top_p_sampling_from_probs, ) _DEBUG_EXPORT_NAMES = [ diff --git a/sgl-kernel/python/sgl_kernel/musa.py b/sgl-kernel/python/sgl_kernel/musa.py index b1bd5eb23..49dd825ca 100644 --- a/sgl-kernel/python/sgl_kernel/musa.py +++ b/sgl-kernel/python/sgl_kernel/musa.py @@ -1,6 +1,7 @@ -from typing import Optional +from typing import Optional, Union import torch +from sgl_kernel.utils import _to_tensor_scalar_tuple def musa_batched_rotary_embedding_contiguous( @@ -167,3 +168,186 @@ def musa_fused_mul_add( torch.ops.sgl_kernel.musa_fused_mul_add(output, self, bias, scale) return output + + +def _top_k_renorm_probs_internal( + probs: torch.Tensor, + maybe_top_k_arr: Optional[torch.Tensor], + top_k_val: int, +) -> torch.Tensor: + probs = probs.float() + maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None + renorm_probs = torch.empty_like(probs) + torch.ops.sgl_kernel.top_k_renorm_probs.default( + probs, renorm_probs, maybe_top_k_arr, top_k_val + ) + return renorm_probs + + +def top_k_renorm_probs( + probs: torch.Tensor, + top_k: Union[torch.Tensor, int], +) -> torch.Tensor: + return _top_k_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_k)) + + +def _top_p_renorm_probs_internal( + probs: torch.Tensor, + maybe_top_p_arr: Optional[torch.Tensor], + top_p_val: float, +) -> torch.Tensor: + probs = probs.float() + maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None + renorm_probs = torch.empty_like(probs) + torch.ops.sgl_kernel.top_p_renorm_probs.default( + probs, renorm_probs, maybe_top_p_arr, top_p_val + ) + return renorm_probs + + +def top_p_renorm_probs( + probs: torch.Tensor, + top_p: Union[torch.Tensor, float], +) -> torch.Tensor: + return _top_p_renorm_probs_internal(probs, *_to_tensor_scalar_tuple(top_p)) + + +def _top_p_sampling_from_probs_internal( + probs: torch.Tensor, + indices: Optional[torch.Tensor], + maybe_top_p_arr: Optional[torch.Tensor], + top_p_val: float, + deterministic: bool, + generator: Optional[torch.Generator], +) -> torch.Tensor: + device = probs.device + probs = probs.float() + maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None + samples = torch.empty(probs.size(0), dtype=torch.int32, device=device) + torch.ops.sgl_kernel.top_p_sampling_from_probs.default( + probs, + samples, + indices, + maybe_top_p_arr, + top_p_val, + deterministic, + generator, + ) + return samples + + +def top_p_sampling_from_probs( + probs: torch.Tensor, + top_p: Union[torch.Tensor, float], + indices: Optional[torch.Tensor] = None, + deterministic: bool = True, + generator: Optional[torch.Generator] = None, + check_nan: bool = False, +) -> torch.Tensor: + if check_nan and torch.any(torch.isnan(probs)): + raise ValueError("Input probs contains NaN.") + return _top_p_sampling_from_probs_internal( + probs, indices, *_to_tensor_scalar_tuple(top_p), deterministic, generator + ) + + +def _top_k_top_p_sampling_from_probs_internal( + probs: torch.Tensor, + indices: Optional[torch.Tensor], + maybe_top_k_arr: Optional[torch.Tensor], + top_k_val: int, + maybe_top_p_arr: Optional[torch.Tensor], + top_p_val: float, + deterministic: bool, + generator: Optional[torch.Generator], +) -> torch.Tensor: + device = probs.device + probs = probs.float() + maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None + maybe_top_p_arr = maybe_top_p_arr.float() if maybe_top_p_arr is not None else None + samples = torch.empty(probs.size(0), dtype=torch.int32, device=device) + torch.ops.sgl_kernel.musa_top_k_top_p_sampling_from_probs.default( + probs, + samples, + indices, + maybe_top_k_arr, + top_k_val, + maybe_top_p_arr, + top_p_val, + deterministic, + generator, + ) + return samples + + +def top_k_top_p_sampling_from_probs( + probs: torch.Tensor, + top_k: Union[torch.Tensor, int], + top_p: Union[torch.Tensor, float], + indices: Optional[torch.Tensor] = None, + filter_apply_order: str = "top_k_first", + deterministic: bool = True, + generator: Optional[torch.Generator] = None, + check_nan: bool = False, +) -> torch.Tensor: + if filter_apply_order == "top_k_first": + renorm_probs = top_k_renorm_probs(probs, top_k) + return top_p_sampling_from_probs( + renorm_probs, + top_p, + indices, + deterministic, + generator=generator, + check_nan=check_nan, + ) + if filter_apply_order == "joint": + if check_nan and torch.any(torch.isnan(probs)): + raise ValueError("Input probs contains NaN.") + return _top_k_top_p_sampling_from_probs_internal( + probs, + indices, + *_to_tensor_scalar_tuple(top_k), + *_to_tensor_scalar_tuple(top_p), + deterministic, + generator, + ) + raise ValueError(f"Invalid filter_apply_order: {filter_apply_order}") + + +def _min_p_sampling_from_probs_internal( + probs: torch.Tensor, + indices: Optional[torch.Tensor], + maybe_min_p_arr: Optional[torch.Tensor], + min_p_val: float, + deterministic: bool, + generator: Optional[torch.Generator], +) -> torch.Tensor: + device = probs.device + probs = probs.float() + maybe_min_p_arr = maybe_min_p_arr.float() if maybe_min_p_arr is not None else None + samples = torch.empty(probs.size(0), dtype=torch.int32, device=device) + torch.ops.sgl_kernel.min_p_sampling_from_probs.default( + probs, + samples, + indices, + maybe_min_p_arr, + min_p_val, + deterministic, + generator, + ) + return samples + + +def min_p_sampling_from_probs( + probs: torch.Tensor, + min_p: Union[torch.Tensor, float], + indices: Optional[torch.Tensor] = None, + deterministic: bool = True, + generator: Optional[torch.Generator] = None, + check_nan: bool = False, +) -> torch.Tensor: + if check_nan and torch.any(torch.isnan(probs)): + raise ValueError("Input probs contains NaN.") + return _min_p_sampling_from_probs_internal( + probs, indices, *_to_tensor_scalar_tuple(min_p), deterministic, generator + ) diff --git a/sgl-kernel/setup_musa.py b/sgl-kernel/setup_musa.py index 4f204dccb..323e9e0b5 100644 --- a/sgl-kernel/setup_musa.py +++ b/sgl-kernel/setup_musa.py @@ -108,6 +108,7 @@ sources = [ "csrc/memory/weak_ref_tensor.cpp", str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"), str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"), + str(_FLASHINFER_REPO.source_dir / "csrc/sampling.cu"), # XXX (MUSA): The following files contain MUSA-specific implementations. "csrc/musa/pos_encoding_contiguous.mu", "csrc/musa/moe_gemv_swiglu.mu",