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

[Refactor] Clean up JIT kernel utilites (#16884)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
DarkSharpness
2026-01-13 17:54:16 +08:00
committed by GitHub
parent 740d3c0b39
commit ba9f6d8f26
30 changed files with 928 additions and 513 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/jit_kernel/utils.py
View File

@@ -42,7 +42,7 @@ 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]
CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool, torch.dtype]
class CPPArgList(list[str]):
@@ -50,6 +50,13 @@ class CPPArgList(list[str]):
return ", ".join(self)
CPP_DTYPE_MAP = {
torch.float: "fp32_t",
torch.float16: "fp16_t",
torch.bfloat16: "bf16_t",
}
def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
def _convert(arg: CPP_TEMPLATE_TYPE) -> str:
if isinstance(arg, bool):
@@ -57,13 +64,7 @@ def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
if isinstance(arg, (int, float)):
return str(arg)
if isinstance(arg, torch.dtype):
if arg == torch.float:
return "float"
if arg == torch.float16:
return "__half"
if arg == torch.bfloat16:
return "nv_bfloat16"
raise TypeError(f"Not implement this arg wrapper yet: {arg}")
return CPP_DTYPE_MAP[arg]
raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}")
return CPPArgList(_convert(arg) for arg in args)