[Chore] Clean up JIT compilation flags (#21022)

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...")
+    parser = argparse.ArgumentParser(
-    include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
+        description="Generate .clangd file for sglang jit kernel development."
-    status = subprocess.run(
-        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
-        capture_output=True,
-        check=True,
     )
-    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+    parser.add_argument(
-    major, minor = compute_cap.split(".")
+        "--overwrite",
-    compile_flags = ",\n    ".join(
+        action="store_true",
-        [
+        help="Overwrite existing .clangd file if it exists.",
-            "-xcuda",
-            f"--cuda-gpu-arch=sm_{major}{minor}",
-            "-std=c++20",
-            "-Wall",
-            "-Wextra",
-        ]
-        + [f"-isystem{path}" for path in include_paths]
     )
+    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()

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"
+    return override_jit_cuda_arch(major, minor, suffix="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
 
 
 @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(),
         )
 

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]()

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.
+    with _jit_compile_context():
-    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:
         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 old_value is None:
-        if not env_existed:
+            os.environ.pop(env_key, None)
-            del os.environ[env_key]
+        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
+    if is_hip_runtime():
-
+        return False
-    device = torch.cuda.current_device()
+    return get_jit_cuda_arch().major >= 9
-    return torch.cuda.get_device_capability(device)[0] >= 9
 
 
-@cache_once
+def _find_package_root(package: str) -> Optional[pathlib.Path]:
-def _get_cuda_arch_value() -> int:
+    spec = importlib.util.find_spec(package)
-    """Get CUDA arch value for -DSGL_CUDA_ARCH (e.g. 900 for SM 9.0)."""
+    if spec is None or spec.origin is None:
-    device = torch.cuda.current_device()
+        return None
-    major, minor = torch.cuda.get_device_capability(device)
+    return pathlib.Path(spec.origin).resolve().parent
-    return major * 100 + minor * 10
 
 
-@cache_once
+# NOTE: this might also be used in __main__.py for compile flags export
-def _get_cuda_arch_list() -> str:
+_REGISTERED_DEPENDENCIES: Dict[str, Callable[[], 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)
+def register_dependency(name: str):
-    return f"{major}.{minor}"
+    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",
+]