[VLM] Adopt jit qk_norm kernel in VLM (#16171)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
2026-01-01 10:10:36 +08:00
parent 12b89e51d8
commit 3a42c5e341
4 changed files with 26 additions and 7 deletions
--- a/python/sglang/jit_kernel/norm.py
+++ b/python/sglang/jit_kernel/norm.py
@@ -30,7 +30,7 @@ def _jit_norm_module(head_dims: int) -> Module:
@cache_once
 def can_use_fused_inplace_qknorm(head_dim: int) -> bool:
    logger = logging.getLogger(__name__)
-    if head_dim not in [64, 128, 256]:
+    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: