mirror of
https://github.com/kvcache-ai/sglang.git
synced 2026-07-13 02:46:57 +00:00
104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
from __future__ import annotations
|
|
|
|
import pathlib
|
|
from functools import lru_cache
|
|
from typing import TYPE_CHECKING, List, Tuple, TypeAlias, Union
|
|
|
|
if TYPE_CHECKING:
|
|
from tvm_ffi import Module
|
|
|
|
|
|
def _make_wrapper(tup: Tuple[str, str]) -> str:
|
|
export_name, kernel_name = tup
|
|
return f"TVM_FFI_DLL_EXPORT_TYPED_FUNC({export_name}, ({kernel_name}));"
|
|
|
|
|
|
@lru_cache()
|
|
def _resolve_kernel_path() -> pathlib.Path:
|
|
cur_dir = pathlib.Path(__file__).parent.resolve()
|
|
|
|
# first, try this directory structure
|
|
def _environment_install():
|
|
candidate = cur_dir.resolve()
|
|
if (candidate / "include").exists() and (candidate / "csrc").exists():
|
|
return candidate
|
|
return None
|
|
|
|
def _package_install():
|
|
# TODO: support find path by package
|
|
return None
|
|
|
|
path = _environment_install() or _package_install()
|
|
if path is None:
|
|
raise RuntimeError("Cannot find sgl-kernel/jit path")
|
|
return path
|
|
|
|
|
|
KERNEL_PATH = _resolve_kernel_path()
|
|
DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")]
|
|
DEFAULT_CFLAGS = ["-std=c++20", "-O3"]
|
|
DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"]
|
|
DEFAULT_LDFLAGS = []
|
|
CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool]
|
|
|
|
|
|
class CPPArgList(list[str]):
|
|
def __str__(self) -> str:
|
|
return ", ".join(self)
|
|
|
|
|
|
def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
|
|
def _convert(arg: CPP_TEMPLATE_TYPE) -> str:
|
|
if isinstance(arg, bool):
|
|
return "true" if arg else "false"
|
|
if isinstance(arg, (int, float)):
|
|
return str(arg)
|
|
raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}")
|
|
|
|
return CPPArgList(_convert(arg) for arg in args)
|
|
|
|
|
|
def load_jit(
|
|
*args: str,
|
|
cpp_files: List[str] | None = None,
|
|
cuda_files: List[str] | None = None,
|
|
cpp_wrappers: List[Tuple[str, str]] | None = None,
|
|
cuda_wrappers: List[Tuple[str, str]] | None = None,
|
|
extra_cflags: List[str] | None = None,
|
|
extra_cuda_cflags: List[str] | None = None,
|
|
extra_ldflags: List[str] | None = None,
|
|
extra_include_paths: List[str] | None = None,
|
|
build_directory: str | None = None,
|
|
) -> Module:
|
|
from tvm_ffi.cpp import load_inline
|
|
|
|
cpp_files = cpp_files or []
|
|
cuda_files = cuda_files or []
|
|
cpp_wrappers = cpp_wrappers or []
|
|
cuda_wrappers = cuda_wrappers or []
|
|
extra_cflags = extra_cflags or []
|
|
extra_cuda_cflags = extra_cuda_cflags or []
|
|
extra_ldflags = extra_ldflags or []
|
|
extra_include_paths = extra_include_paths or []
|
|
|
|
# include cpp files
|
|
cpp_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cpp_files]
|
|
cpp_sources = [f'#include "{path}"' for path in cpp_paths]
|
|
cpp_sources += [_make_wrapper(tup) for tup in cpp_wrappers]
|
|
|
|
# include cuda files
|
|
cuda_paths = [(KERNEL_PATH / "csrc" / f).resolve() for f in cuda_files]
|
|
cuda_sources = [f'#include "{path}"' for path in cuda_paths]
|
|
cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers]
|
|
|
|
return load_inline(
|
|
"sgl_kernel_jit_" + "_".join(str(arg) for arg in args),
|
|
cpp_sources=cpp_sources,
|
|
cuda_sources=cuda_sources,
|
|
extra_cflags=DEFAULT_CFLAGS + extra_cflags,
|
|
extra_cuda_cflags=DEFAULT_CUDA_CFLAGS + extra_cuda_cflags,
|
|
extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags,
|
|
extra_include_paths=DEFAULT_INCLUDE + extra_include_paths,
|
|
build_directory=build_directory,
|
|
)
|