[diffusion] fix: keep fused qk-norm-rope out of dynamo tracing (#31849)

2026-07-28 19:51:00 +08:00
parent a24906a091
commit 75017c3fa0
2 changed files with 3 additions and 2 deletions
@@ -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)
@@ -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,