[diffusion] fix: only use fused qk_norm on nvidia gpu (#33994)

2026-09-04 09:25:08 +08:00
parent 6c63d678d5
commit 5e81d462c3
1 changed files with 1 additions and 0 deletions
@@ -440,6 +440,7 @@ def _apply_qk_norm(
        and q.stride(-2) == k.stride(-2) == head_dim
        and q_norm.eps == k_norm.eps
        and not torch.compiler.is_compiling()
+        and current_platform.is_cuda()
    ):
        fused_inplace_qknorm(
            q,