From ec01ef9092fbb8467504aa36649ce33bef7a6689 Mon Sep 17 00:00:00 2001 From: Jia Guo Date: Fri, 20 Mar 2026 03:05:09 -0700 Subject: [PATCH] =?UTF-8?q?Fix=20torch.compile/dynamo=20crash=20with=20Qwe?= =?UTF-8?q?n3=20QK-norm=20in=20piecewise=20CUDA=20g=E2=80=A6=20(#19818)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Claude Opus 4.6 (1M context) --- python/sglang/jit_kernel/norm.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index e3b2aee1b..88a4c0b15 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Optional import torch +logger = logging.getLogger(__name__) + from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, @@ -62,9 +64,9 @@ def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module: ) +@torch.compiler.assume_constant_result @cache_once 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