diff --git a/.gitignore b/.gitignore index 9ad1bf599..3ecff6e63 100644 --- a/.gitignore +++ b/.gitignore @@ -266,3 +266,17 @@ python/sglang/_version.py python/sglang/srt/grpc/*_pb2.py python/sglang/srt/grpc/*_pb2_grpc.py python/sglang/srt/grpc/*_pb2.pyi + +# MUSA section +# Generated source files by torchada +sgl-kernel/csrc_musa/ +sgl-kernel/include_musa/ +sgl-kernel/csrc/**/*_musa/ +sgl-kernel/third_party/*/csrc_musa/ +sgl-kernel/third_party/*/include_musa/ + +# Third-party libraries source code +sgl-kernel/third_party/ + +# MUSA core dump files +core_* diff --git a/docs/platforms/mthreads_gpu.md b/docs/platforms/mthreads_gpu.md new file mode 100644 index 000000000..3538df318 --- /dev/null +++ b/docs/platforms/mthreads_gpu.md @@ -0,0 +1,25 @@ +# Moore Threads GPUs + +This document describes how run SGLang on Moore Threads GPUs. If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues). + +## Install SGLang + +You can install SGLang using one of the methods below. + +### Install from Source + +```bash +# Use the default branch +git clone https://github.com/sgl-project/sglang.git +cd sglang + +# Compile sgl-kernel +pip install --upgrade pip +cd sgl-kernel +python setup_musa.py install + +# Install sglang python package +cd .. +rm -f python/pyproject.toml && mv python/pyproject_other.toml python/pyproject.toml +pip install -e "python[all_musa]" +``` diff --git a/sgl-kernel/csrc/common_extension_musa.cc b/sgl-kernel/csrc/common_extension_musa.cc new file mode 100644 index 000000000..33bc639a1 --- /dev/null +++ b/sgl-kernel/csrc/common_extension_musa.cc @@ -0,0 +1,51 @@ +/* Copyright 2025 SGLang Team. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include "sgl_kernel_ops.h" +#include "torch_musa/csrc/aten/musa/MUSAContext.h" + +TORCH_LIBRARY_EXPAND(sgl_kernel, m) { + /* + * From FlashInfer + */ + 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_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()"); + m.impl("top_k_renorm_probs", torch::kMUSA, &top_k_renorm_probs); + + 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( + "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); + + m.def( + "top_k_top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_k_arr, " + "float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()"); + m.impl("top_k_top_p_sampling_from_probs", torch::kMUSA, &top_k_top_p_sampling_from_probs); + + m.def("top_k_mask_logits(Tensor logits, Tensor mask_logits, Tensor? maybe_top_k_arr, int top_k_val) -> ()"); + m.impl("top_k_mask_logits", torch::kMUSA, &top_k_mask_logits); +} + +REGISTER_EXTENSION(common_ops) diff --git a/sgl-kernel/pyproject_musa.toml b/sgl-kernel/pyproject_musa.toml new file mode 100644 index 000000000..b7d7a781b --- /dev/null +++ b/sgl-kernel/pyproject_musa.toml @@ -0,0 +1,33 @@ +[build-system] +requires = [ + "setuptools>=75.0", + "scikit-build-core>=0.10", + "torch", + "torchada>=0.1.14", + "wheel", +] +build-backend = "setuptools.build_meta" + +[project] +name = "sgl-kernel" +version = "0.3.20" +description = "Kernel Library for SGLang" +readme = "README.md" +requires-python = ">=3.10" +license = { file = "LICENSE" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + "Environment :: GPU :: MTHREADS MUSA" +] +dependencies = [] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang/tree/main/sgl-kernel" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.wheel] +exclude = [ + "dist*", + "tests*", +] diff --git a/sgl-kernel/setup_musa.py b/sgl-kernel/setup_musa.py new file mode 100644 index 000000000..a1df4d90f --- /dev/null +++ b/sgl-kernel/setup_musa.py @@ -0,0 +1,205 @@ +# Copyright 2025 SGLang Team. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import os +import platform +import subprocess +import sys +from pathlib import Path + +# isort: off +import torch +import torchada # noqa: F401 + +# isort: on +from setuptools import find_packages, setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +root = Path(__file__).parent.resolve() +third_party = Path("third_party") +arch = platform.machine().lower() + + +class _RepoInfo: + """Configuration for a third-party git repository.""" + + def __init__(self, name, git_repository, git_tag, git_shallow=False): + self.name = name + self.git_repository = git_repository + self.git_tag = git_tag + self.git_shallow = git_shallow + self.source_dir = third_party / name + + +_FLASHINFER_REPO = _RepoInfo( + name="flashinfer", + git_repository="https://github.com/flashinfer-ai/flashinfer.git", + git_tag="bc29697ba20b7e6bdb728ded98f04788e16ee021", + git_shallow=False, +) + +_MUTLASS_REPO = _RepoInfo( + name="mutlass", + git_repository="https://github.com/MooreThreads/mutlass.git", + git_tag="3abd6a728aacd190df0d922514aca8a8bc3c46b7", + git_shallow=False, +) + + +def _get_version(): + with open(root / "pyproject.toml") as f: + for line in f: + if line.startswith("version"): + return line.split("=")[1].strip().strip('"') + + +operator_namespace = "sgl_kernel" +include_dirs = [ + root / "include", + root / "include" / "impl", + root / "csrc", + root / _FLASHINFER_REPO.source_dir / "include", + root / _FLASHINFER_REPO.source_dir / "csrc", + root / _MUTLASS_REPO.source_dir / "include", +] + +sources = [ + "csrc/common_extension_musa.cc", + str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"), + str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"), + str(_FLASHINFER_REPO.source_dir / "csrc/sampling.cu"), +] + +cxx_flags = ["force_mcc"] +libraries = ["c10", "torch", "torch_python"] +extra_link_args = [ + "-Wl,-rpath,$ORIGIN/../../torch/lib", + f"-L/usr/lib/{arch}-linux-gnu", + "-lmublasLt", +] + +default_target = "mp_31" +mtgpu_target = os.environ.get("MTGPU_TARGET", default_target) + +if torch.musa.is_available(): + try: + prop = torch.musa.get_device_properties(0) + mtgpu_target = f"mp_{prop.major}{prop.minor}" + except Exception as e: + print(f"Warning: Failed to detect GPU properties: {e}") +else: + print(f"Warning: torch.musa not available. Using default target: {mtgpu_target}") + +if mtgpu_target not in ["mp_22", "mp_31"]: + print( + f"Warning: Unsupported GPU architecture detected '{mtgpu_target}'. Expected 'mp_22' or 'mp_31'." + ) + sys.exit(1) + +mcc_flags = [ + "-DNDEBUG", + f"-DOPERATOR_NAMESPACE={operator_namespace}", + "-O3", + "-fPIC", + "-std=c++17", + f"--cuda-gpu-arch={mtgpu_target}", + "-x", + "musa", + "-mtgpu", + "-Od3", + "-ffast-math", + "-fmusa-flush-denormals-to-zero", + "-fno-strict-aliasing", + "-DUSE_MUSA", + "-DENABLE_BF16", + "-DFLASHINFER_ENABLE_F16", + "-DFLASHINFER_ENABLE_BF16", +] + +if mtgpu_target == "mp_31": + mcc_flags.extend( + [ + "-DENABLE_FP8", + "-DFLASHINFER_ENABLE_FP8", + "-DFLASHINFER_ENABLE_FP8_E4M3", + "-DFLASHINFER_ENABLE_FP8_E5M2", + ] + ) + +ext_modules = [ + CUDAExtension( + name="sgl_kernel.common_ops", + sources=sources, + include_dirs=include_dirs, + extra_compile_args={ + "mcc": mcc_flags, + "cxx": cxx_flags, + }, + libraries=libraries, + extra_link_args=extra_link_args, + py_limited_api=False, + ), +] + + +class _CustomBuildExt(BuildExtension): + """Custom build extension that clones third-party repositories before building.""" + + @staticmethod + def _clone_and_checkout(repo_path, repo_url, git_tag, git_shallow): + """Clone a git repository and checkout a specific tag/commit.""" + repo_path.parent.mkdir(exist_ok=True) + if not repo_path.exists(): + clone_cmd = ["git", "clone"] + if git_shallow: + clone_cmd += ["--depth", "1"] + clone_cmd += [repo_url, str(repo_path)] + subprocess.check_call(clone_cmd) + subprocess.check_call(["git", "checkout", git_tag], cwd=repo_path) + else: + subprocess.check_call(["git", "fetch", "--all"], cwd=repo_path) + subprocess.check_call(["git", "checkout", git_tag], cwd=repo_path) + + def run(self): + if os.environ.get("SKIP_THIRD_PARTY", "0") == "1": + print("Skipping third-party repositories cloning (SKIP_THIRD_PARTY=1)") + else: + print("Cloning third-party repositories...") + self._clone_and_checkout( + _MUTLASS_REPO.source_dir, + _MUTLASS_REPO.git_repository, + _MUTLASS_REPO.git_tag, + _MUTLASS_REPO.git_shallow, + ) + self._clone_and_checkout( + _FLASHINFER_REPO.source_dir, + _FLASHINFER_REPO.git_repository, + _FLASHINFER_REPO.git_tag, + _FLASHINFER_REPO.git_shallow, + ) + print("Third-party repositories ready.") + + super().run() + + +setup( + name="sgl-kernel", + version=_get_version(), + packages=find_packages(where="python"), + package_dir={"": "python"}, + ext_modules=ext_modules, + cmdclass={"build_ext": _CustomBuildExt.with_options(use_ninja=True)}, + options={"bdist_wheel": {"py_limited_api": "cp39"}}, +)