chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[MUSA][2/N] sgl-kernel build (#17053)

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-01-24 06:41:47 +08:00
committed by GitHub
parent a77729a276
commit 628ab5d57b
5 changed files with 328 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
.gitignore vendored
View File

@@ -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_*

25
docs/platforms/mthreads_gpu.md Normal file
View File

@@ -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]"
```

51
sgl-kernel/csrc/common_extension_musa.cc Normal file
View File

@@ -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 <ATen/core/dispatch/Dispatcher.h>
#include <torch/library.h>
#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)

33
sgl-kernel/pyproject_musa.toml Normal file
View File

@@ -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*",
]

205
sgl-kernel/setup_musa.py Normal file
View File

@@ -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"}},
)