100 lines
2.9 KiB
Python
100 lines
2.9 KiB
Python
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():
|
|
logger = logging.getLogger()
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate .clangd file for sglang jit kernel development."
|
|
)
|
|
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_str}
|
|
]
|
|
"""
|
|
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:
|
|
f.write(clangd_content)
|
|
logger.info(".clangd file generated.")
|
|
|
|
|
|
assert __name__ == "__main__"
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
generate_clangd()
|