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

[Feature] Support JIT set kv cache (#16273)

Browse Source
This commit is contained in:
DarkSharpness
2026-01-11 09:34:09 +08:00
committed by GitHub
parent a2c2c09d7d
commit d112f6a25b
9 changed files with 516 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,12 +17,12 @@ if TYPE_CHECKING:
@cache_once
def _jit_norm_module(head_dims: int, dtype: torch.dtype) -> Module:
args = make_cpp_args(head_dims, is_arch_support_pdl(), dtype)
def _jit_qknorm_module(head_dim: int, dtype: torch.dtype) -> Module:
args = make_cpp_args(head_dim, is_arch_support_pdl(), dtype)
return load_jit(
"norm",
"qknorm",
*args,
cuda_files=["norm.cuh"],
cuda_files=["elementwise/qknorm.cuh"],
cuda_wrappers=[("qknorm", f"QKNormKernel<{args}>::run")],
)
@@ -34,7 +34,7 @@ def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool:
logger.warning(f"Unsupported head_dim={head_dim} for JIT QK-Norm kernel")
return False
try:
_jit_norm_module(head_dim, dtype)
_jit_qknorm_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, q.dtype)
module = _jit_qknorm_module(head_dim, q.dtype)
module.qknorm(q, k, q_weight, k_weight, eps)