[jit kernel] support dtype as a cpp template parameter (#16452)

python/sglang/jit_kernel/utils.py

@@ -5,6 +5,8 @@ import pathlib
 from functools import lru_cache
 from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeAlias, TypeVar, Union
 
+import torch
+
 if TYPE_CHECKING:
     from tvm_ffi import Module
 
@@ -54,6 +56,14 @@ def make_cpp_args(*args: CPP_TEMPLATE_TYPE) -> CPPArgList:
             return "true" if arg else "false"
         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}")
         raise TypeError(f"Unsupported argument type for cpp template: {type(arg)}")
 
     return CPPArgList(_convert(arg) for arg in args)