diff --git a/python/sglang/jit_kernel/__main__.py b/python/sglang/jit_kernel/__main__.py index bacf4f84e..b626fde78 100644 --- a/python/sglang/jit_kernel/__main__.py +++ b/python/sglang/jit_kernel/__main__.py @@ -1,43 +1,91 @@ -assert __name__ == "__main__" +import argparse +import logging +import os + +from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path + +from sglang.jit_kernel.utils import ( + _REGISTERED_DEPENDENCIES, + DEFAULT_INCLUDE, + _get_default_target_flags, + get_jit_cuda_arch, + override_jit_cuda_arch, +) def generate_clangd(): - import logging - import os - import subprocess - - from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path - - from sglang.jit_kernel.utils import DEFAULT_INCLUDE - logger = logging.getLogger() - logger.info("Generating .clangd file...") - include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE - status = subprocess.run( - args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], - capture_output=True, - check=True, + parser = argparse.ArgumentParser( + description="Generate .clangd file for sglang jit kernel development." ) - compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] - major, minor = compute_cap.split(".") - compile_flags = ",\n ".join( - [ - "-xcuda", - f"--cuda-gpu-arch=sm_{major}{minor}", - "-std=c++20", - "-Wall", - "-Wextra", - ] - + [f"-isystem{path}" for path in include_paths] + parser.add_argument( + "--overwrite", + action="store_true", + help="Overwrite existing .clangd file if it exists.", ) + parser.add_argument( + "--dependencies", + "--dep", + nargs="*", + default=[], + choices=_REGISTERED_DEPENDENCIES.keys(), + help="Extra dependency libraries to include.", + ) + parser.add_argument( + "--cuda-target", + "--cuda", + default=None, + type=str, + help="Target architecture to generate compile flags for.", + ) + args = parser.parse_args() + + dep_include_paths = [] + for dep in args.dependencies: + if dep not in _REGISTERED_DEPENDENCIES: + raise ValueError(f"Dependency {dep} is not registered.") + dep_include_paths += _REGISTERED_DEPENDENCIES[dep]() + + include_paths = [ + *DEFAULT_INCLUDE, + find_include_path(), + find_dlpack_include_path(), + *dep_include_paths, + ] + if args.cuda_target: + assert args.cuda_target.count(".") == 1 + major, minor = args.cuda_target.split(".") + major, minor = int(major), int(minor) + context = override_jit_cuda_arch(major, minor) + context.__enter__() + else: + arch = get_jit_cuda_arch() + major, minor = arch.major, f"{arch.minor}{arch.suffix}" + assert ( + major > 0 + ), "Cannot detect CUDA architecture, please specify --cuda-target explicitly." + + compile_flags = [ + "-xcuda", + f"--cuda-gpu-arch=sm_{major}{minor}", + "-Wall", + "-Wextra", + *_get_default_target_flags(), + *[f"-isystem{path}" for path in include_paths], + ] + # NOTE: skip these flags because clangd don't recognize them + UNSUPPORTED_FLAGS = {"--expt-relaxed-constexpr"} + compile_flags = [flag for flag in compile_flags if flag not in UNSUPPORTED_FLAGS] + compile_flags_str = ",\n ".join(compile_flags) clangd_content = f""" CompileFlags: Add: [ - {compile_flags} + {compile_flags_str} ] """ - if os.path.exists(".clangd"): + if os.path.exists(".clangd") and not args.overwrite: logger.warning(".clangd file already exists, nothing done.") + logger.warning("Use --overwrite to force overwrite the existing .clangd file.") logger.warning(f"suggested content: {clangd_content}") else: with open(".clangd", "w") as f: @@ -45,4 +93,7 @@ CompileFlags: logger.info(".clangd file generated.") +assert __name__ == "__main__" + +logging.basicConfig(level=logging.INFO) generate_clangd() diff --git a/python/sglang/jit_kernel/nvfp4.py b/python/sglang/jit_kernel/nvfp4.py index 6ca86c47e..f54e6414b 100644 --- a/python/sglang/jit_kernel/nvfp4.py +++ b/python/sglang/jit_kernel/nvfp4.py @@ -1,14 +1,11 @@ from __future__ import annotations -import importlib.util import os -import pathlib -from contextlib import contextmanager from typing import TYPE_CHECKING, Optional, Tuple import torch -from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.jit_kernel.utils import cache_once, load_jit, override_jit_cuda_arch from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op @@ -20,43 +17,6 @@ _FLOAT4_E2M1_MAX = 6.0 _FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max -def _find_package_root(package: str) -> Optional[pathlib.Path]: - spec = importlib.util.find_spec(package) - if spec is None or spec.origin is None: - return None - return pathlib.Path(spec.origin).resolve().parent - - -def _resolve_cutlass_include_paths() -> list[str]: - include_paths: list[str] = [] - - flashinfer_root = _find_package_root("flashinfer") - if flashinfer_root is not None: - candidates = [ - flashinfer_root / "data" / "cutlass" / "include", - flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include", - ] - for path in candidates: - if path.exists(): - include_paths.append(str(path)) - - deep_gemm_root = _find_package_root("deep_gemm") - if deep_gemm_root is not None: - candidate = deep_gemm_root / "include" - if candidate.exists(): - include_paths.append(str(candidate)) - - # De-duplicate while preserving order. - unique_paths = [] - seen = set() - for path in include_paths: - if path in seen: - continue - seen.add(path) - unique_paths.append(path) - return unique_paths - - def _nvfp4_cuda_flags() -> list[str]: return [ "-DNDEBUG", @@ -71,7 +31,7 @@ def _nvfp4_cuda_flags() -> list[str]: ] -def _get_nvfp4_cuda_arch_list() -> str: +def _nvfp4_arch_env(): if not torch.cuda.is_available(): raise RuntimeError("NVFP4 JIT kernels require CUDA.") major, minor = torch.cuda.get_device_capability() @@ -84,32 +44,11 @@ def _get_nvfp4_cuda_arch_list() -> str: # JIT compilation targets only the current device, unlike AOT fat-binaries; # adding extra architectures here would clash with the single SGL_CUDA_ARCH # value injected by load_jit(). - return f"{major}.{minor}a" - - -@contextmanager -def _nvfp4_arch_env(): - key = "TVM_FFI_CUDA_ARCH_LIST" - old_val = os.environ.get(key) - os.environ[key] = _get_nvfp4_cuda_arch_list() - try: - yield - finally: - if old_val is None: - os.environ.pop(key, None) - else: - os.environ[key] = old_val + return override_jit_cuda_arch(major, minor, suffix="a") @cache_once def _jit_nvfp4_quant_module() -> Module: - extra_include_paths = _resolve_cutlass_include_paths() - if not extra_include_paths: - raise RuntimeError( - "Cannot find CUTLASS headers required for NVFP4 JIT quantization. " - "Please install flashinfer or deep_gemm with CUTLASS headers." - ) - with _nvfp4_arch_env(): return load_jit( "nvfp4_quant", @@ -119,20 +58,13 @@ def _jit_nvfp4_quant_module() -> Module: cuda_wrappers=[ ("scaled_fp4_quant", "scaled_fp4_quant_sm100a_sm120a"), ], - extra_include_paths=extra_include_paths, extra_cuda_cflags=_nvfp4_cuda_flags(), + extra_dependencies=["cutlass"], ) @cache_once def _jit_nvfp4_expert_quant_module() -> Module: - extra_include_paths = _resolve_cutlass_include_paths() - if not extra_include_paths: - raise RuntimeError( - "Cannot find CUTLASS headers required for NVFP4 JIT expert quantization. " - "Please install flashinfer or deep_gemm with CUTLASS headers." - ) - with _nvfp4_arch_env(): return load_jit( "nvfp4_expert_quant", @@ -146,20 +78,13 @@ def _jit_nvfp4_expert_quant_module() -> Module: "silu_and_mul_scaled_fp4_experts_quant_sm100a", ), ], - extra_include_paths=extra_include_paths, + extra_dependencies=["cutlass"], extra_cuda_cflags=_nvfp4_cuda_flags(), ) @cache_once def _jit_nvfp4_scaled_mm_module() -> Module: - extra_include_paths = _resolve_cutlass_include_paths() - if not extra_include_paths: - raise RuntimeError( - "Cannot find CUTLASS headers required for NVFP4 JIT GEMM. " - "Please install flashinfer or deep_gemm with CUTLASS headers." - ) - with _nvfp4_arch_env(): return load_jit( "nvfp4_scaled_mm", @@ -168,20 +93,13 @@ def _jit_nvfp4_scaled_mm_module() -> Module: "gemm/nvfp4/nvfp4_scaled_mm_entry.cuh", ], cuda_wrappers=[("cutlass_scaled_fp4_mm", "cutlass_scaled_fp4_mm")], - extra_include_paths=extra_include_paths, + extra_dependencies=["cutlass"], extra_cuda_cflags=_nvfp4_cuda_flags(), ) @cache_once def _jit_nvfp4_blockwise_moe_module() -> Module: - extra_include_paths = _resolve_cutlass_include_paths() - if not extra_include_paths: - raise RuntimeError( - "Cannot find CUTLASS headers required for NVFP4 JIT MoE grouped GEMM. " - "Please install flashinfer or deep_gemm with CUTLASS headers." - ) - with _nvfp4_arch_env(): return load_jit( "nvfp4_blockwise_moe", @@ -191,7 +109,7 @@ def _jit_nvfp4_blockwise_moe_module() -> Module: cuda_wrappers=[ ("cutlass_fp4_group_mm", "cutlass_fp4_group_mm_sm100a_sm120a") ], - extra_include_paths=extra_include_paths, + extra_dependencies=["cutlass"], extra_cuda_cflags=_nvfp4_cuda_flags(), ) diff --git a/python/sglang/jit_kernel/tests/test_dependency.py b/python/sglang/jit_kernel/tests/test_dependency.py new file mode 100644 index 000000000..9fc981956 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_dependency.py @@ -0,0 +1,13 @@ +import pytest + +from sglang.jit_kernel.utils import _REGISTERED_DEPENDENCIES +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=30, suite="nightly-kernel-1-gpu", nightly=True) + + +@pytest.mark.parametrize("name", _REGISTERED_DEPENDENCIES.keys()) +def test_availability(name: str) -> None: + # NOTE: the path resolution should not fail + _REGISTERED_DEPENDENCIES[name]() diff --git a/python/sglang/jit_kernel/utils.py b/python/sglang/jit_kernel/utils.py index fdcbc8e8a..63a0ba99d 100644 --- a/python/sglang/jit_kernel/utils.py +++ b/python/sglang/jit_kernel/utils.py @@ -1,9 +1,24 @@ from __future__ import annotations import functools +import importlib.util +import logging import os import pathlib -from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union +from contextlib import contextmanager +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + TypeAlias, + TypeVar, + Union, +) import torch @@ -15,6 +30,8 @@ if TYPE_CHECKING: F = TypeVar("F", bound=Callable[..., Any]) _FULL_TEST_ENV_VAR = "SGLANG_JIT_KERNEL_RUN_FULL_TESTS" +logger = logging.getLogger(__name__) + def should_run_full_tests() -> bool: return os.getenv(_FULL_TEST_ENV_VAR, "false").lower() == "true" @@ -72,10 +89,6 @@ def _resolve_kernel_path() -> pathlib.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_HIP_CFLAGS = [ - flag for flag in DEFAULT_CUDA_CFLAGS if flag != "--expt-relaxed-constexpr" -] DEFAULT_LDFLAGS = [] CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool, torch.dtype] @@ -125,6 +138,7 @@ def load_jit( extra_cuda_cflags: List[str] | None = None, extra_ldflags: List[str] | None = None, extra_include_paths: List[str] | None = None, + extra_dependencies: List[str] | None = None, build_directory: str | None = None, ) -> Module: """ @@ -151,6 +165,8 @@ def load_jit( :type extra_ldflags: List[str] | None :param extra_include_paths: Extra include paths. :type extra_include_paths: List[str] | None + :param extra_dependencies: Extra dependencies for the JIT module, e.g., cutlass. + :type extra_dependencies: List[str] | None :param build_directory: The build directory for JIT compilation. :type build_directory: str | None :return: A just-in-time(JIT) compiled module. @@ -168,6 +184,11 @@ def load_jit( extra_ldflags = extra_ldflags or [] extra_include_paths = extra_include_paths or [] + for dep in set(extra_dependencies or []): + if dep not in _REGISTERED_DEPENDENCIES: + raise ValueError(f"Dependency {dep} is not registered.") + extra_include_paths += _REGISTERED_DEPENDENCIES[dep]() + # 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] @@ -178,55 +199,196 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] - # Override TVM_FFI_CUDA_ARCH_LIST if it does not exist. - env_key = "TVM_FFI_CUDA_ARCH_LIST" - env_existed = env_key in os.environ - selected_cuda_cflags = DEFAULT_CUDA_CFLAGS - if is_hip_runtime(): - selected_cuda_cflags = DEFAULT_HIP_CFLAGS - extra_cuda_cflags = ["-DUSE_ROCM"] + extra_cuda_cflags - else: - extra_cuda_cflags = [ - f"-DSGL_CUDA_ARCH={_get_cuda_arch_value()}" - ] + extra_cuda_cflags - if not env_existed: - os.environ[env_key] = _get_cuda_arch_list() - try: + with _jit_compile_context(): 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=selected_cuda_cflags + extra_cuda_cflags, + extra_cuda_cflags=_get_default_target_flags() + extra_cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) + + +@dataclass +class ArchInfo: + major: int + minor: int + suffix: str + + @property + def target_name(self) -> str: + return f"{self.major}.{self.minor}{self.suffix}" + + @property + def jit_flag(self) -> str: + return f"-DSGL_CUDA_ARCH={self.major * 100 + self.minor * 10}" + + +@cache_once +def _init_jit_cuda_arch_once(): + global _CUDA_ARCH + try: + device = torch.cuda.current_device() + major, minor = torch.cuda.get_device_capability(device) + except Exception: + logger.warning("Cannot detect CUDA architecture.") + major, minor = 0, 0 # invalid value to trigger compile error if used + _CUDA_ARCH = ArchInfo(major, minor, "") + + +@contextmanager +def _jit_compile_context(): + if is_hip_runtime(): + yield # TODO: support ROCm `TVM_FFI_ROCM_ARCH_LIST` if needed + return + env_key = "TVM_FFI_CUDA_ARCH_LIST" + old_value = os.environ.get(env_key, None) + os.environ[env_key] = get_jit_cuda_arch().target_name + try: + yield finally: - # Reset TVM_FFI_CUDA_ARCH_LIST to original state (not exist) - if not env_existed: - del os.environ[env_key] + if old_value is None: + os.environ.pop(env_key, None) + else: + os.environ[env_key] = old_value + + +# NOTE: this might also be used in __main__.py for compile flags export +def _get_default_target_flags() -> List[str]: + if is_hip_runtime(): + return ["-DUSE_ROCM", "-std=c++20", "-O3"] + else: + return [ + get_jit_cuda_arch().jit_flag, + "-std=c++20", + "-O3", + "--expt-relaxed-constexpr", + ] + + +@contextmanager +def override_jit_cuda_arch(major: int, minor: int, suffix: str = ""): + """A context manager to temporarily override CUDA architecture.""" + global _CUDA_ARCH + old_value = get_jit_cuda_arch() + _CUDA_ARCH = ArchInfo(major, minor, suffix) + try: + yield + finally: + _CUDA_ARCH = old_value + + +def get_jit_cuda_arch() -> ArchInfo: + """Get the current CUDA architecture info.""" + _init_jit_cuda_arch_once() + return _CUDA_ARCH -@cache_once def is_arch_support_pdl() -> bool: - import torch - - device = torch.cuda.current_device() - return torch.cuda.get_device_capability(device)[0] >= 9 + if is_hip_runtime(): + return False + return get_jit_cuda_arch().major >= 9 -@cache_once -def _get_cuda_arch_value() -> int: - """Get CUDA arch value for -DSGL_CUDA_ARCH (e.g. 900 for SM 9.0).""" - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - return major * 100 + minor * 10 +def _find_package_root(package: str) -> Optional[pathlib.Path]: + spec = importlib.util.find_spec(package) + if spec is None or spec.origin is None: + return None + return pathlib.Path(spec.origin).resolve().parent -@cache_once -def _get_cuda_arch_list() -> str: - """Get the correct CUDA architecture string for TVM_FFI_CUDA_ARCH_LIST.""" - device = torch.cuda.current_device() - major, minor = torch.cuda.get_device_capability(device) - return f"{major}.{minor}" +# NOTE: this might also be used in __main__.py for compile flags export +_REGISTERED_DEPENDENCIES: Dict[str, Callable[[], List[str]]] = {} + + +def register_dependency(name: str): + def decorator(f: Callable[[], List[str]]) -> Callable[[], List[str]]: + if name in _REGISTERED_DEPENDENCIES: + raise ValueError(f"Dependency {name} already registered") + _REGISTERED_DEPENDENCIES[name] = f + return f + + return decorator + + +@register_dependency("flashinfer") +def get_flashinfer_include_paths() -> List[str]: + include_paths: List[str] = [] + flashinfer_root = _find_package_root("flashinfer") + if flashinfer_root is None: + raise RuntimeError( + "Cannot find flashinfer package. Please install flashinfer to get" + "the required headers for JIT compilation." + ) + + flashinfer_data = flashinfer_root / "data" + candidates = [ + flashinfer_data / "include", + flashinfer_data / "csrc", + flashinfer_data / "cutlass" / "include", + flashinfer_data / "cutlass" / "tools" / "util" / "include", + flashinfer_data / "spdlog" / "include", + ] + + for path in candidates: + if not path.exists(): + raise RuntimeError( + f"Required header path {path} for flashinfer dependency not found." + " Please check your flashinfer installation." + ) + include_paths.append(str(path)) + return include_paths + + +@register_dependency("cutlass") +def get_cutlass_include_paths() -> List[str]: + include_paths: List[str] = [] + + flashinfer_root = _find_package_root("flashinfer") + if flashinfer_root is not None: + candidates = [ + flashinfer_root / "data" / "cutlass" / "include", + flashinfer_root / "data" / "cutlass" / "tools" / "util" / "include", + ] + for path in candidates: + if path.exists(): + include_paths.append(str(path)) + + deep_gemm_root = _find_package_root("deep_gemm") + if deep_gemm_root is not None: + candidate = deep_gemm_root / "include" + if candidate.exists(): + include_paths.append(str(candidate)) + + # De-duplicate while preserving order. + unique_paths = [] + seen = set() + for path in include_paths: + if path in seen: + continue + seen.add(path) + unique_paths.append(path) + + if not unique_paths: + raise RuntimeError( + "Cannot find CUTLASS headers required for JIT compilation. " + "Please install flashinfer or deep_gemm with CUTLASS headers." + ) + return unique_paths + + +__all__ = [ + "should_run_full_tests", + "get_ci_test_range", + "cache_once", + "is_hip_runtime", + "make_cpp_args", + "load_jit", + "override_jit_cuda_arch", + "get_jit_cuda_arch", + "is_arch_support_pdl", + "register_dependency", +]