From 75017c3fa0c9bfb1c528cad3b7698354e4f77ce9 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 28 Jul 2026 19:51:00 +0800 Subject: [PATCH] [diffusion] fix: keep fused qk-norm-rope out of dynamo tracing (#31849) --- python/sglang/multimodal_gen/runtime/layers/layernorm.py | 1 + .../sglang/multimodal_gen/runtime/models/dits/qwen_image.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/layernorm.py b/python/sglang/multimodal_gen/runtime/layers/layernorm.py index d57e86373..3b27c3222 100755 --- a/python/sglang/multimodal_gen/runtime/layers/layernorm.py +++ b/python/sglang/multimodal_gen/runtime/layers/layernorm.py @@ -1021,6 +1021,7 @@ def apply_qk_norm_rope( if ( fused_enabled and _is_cuda + and not torch.compiler.is_compiling() and allow_inplace and (q_eps == k_eps) and q.dtype in (torch.float16, torch.bfloat16) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py index 104eb3bfa..68396b3e8 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py @@ -737,7 +737,7 @@ class QwenImageCrossAttention(nn.Module): k=img_key, q_norm=self.norm_q, k_norm=self.norm_k, - head_dim=img_query.shape[-1], + head_dim=self.head_dim, cos_sin_cache=img_cache, is_neox=False, allow_inplace=True, @@ -747,7 +747,7 @@ class QwenImageCrossAttention(nn.Module): k=txt_key, q_norm=self.norm_added_q, k_norm=self.norm_added_k, - head_dim=txt_query.shape[-1], + head_dim=self.head_dim, cos_sin_cache=txt_cache, is_neox=False, allow_inplace=True,