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

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

Browse Source
This commit is contained in:
陈一涵
2026-01-08 13:54:33 +08:00
committed by GitHub
parent 41b434a7e6
commit 48b8dcd42e
9 changed files with 91 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/jit_kernel/norm.py
View File

@@ -17,8 +17,8 @@ if TYPE_CHECKING:
@cache_once
def _jit_norm_module(head_dims: int) -> Module:
args = make_cpp_args(head_dims, is_arch_support_pdl())
def _jit_norm_module(head_dims: int, dtype: torch.dtype) -> Module:
args = make_cpp_args(head_dims, is_arch_support_pdl(), dtype)
return load_jit(
"norm",
*args,
@@ -28,13 +28,13 @@ def _jit_norm_module(head_dims: int) -> Module:
@cache_once
def can_use_fused_inplace_qknorm(head_dim: int) -> bool:
def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
logger = logging.getLogger(__name__)
if head_dim not in [64, 128, 256, 512, 1024]:
logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel")
return False
try:
_jit_norm_module(head_dim)
_jit_norm_module(head_dim, dtype)
return True
except Exception as e:
logger.warning(f"Failed to load JIT QK-Norm kernel: {e}")
@@ -51,5 +51,5 @@ def fused_inplace_qknorm(
head_dim: int = 0,
) -> None:
head_dim = head_dim or q.size(-1)
module = _jit_norm_module(head_dim)
module = _jit_norm_module(head_dim, q.dtype)
module.qknorm(q, k, q_weight, k_weight, eps)