[NPU] adapt_fused_rope_qk_mqa_optimize (#28872)

2026-06-25 08:55:52 +08:00
parent d0524d6433
commit 9305d10099
1 changed files with 1 additions and 4 deletions
@@ -294,10 +294,7 @@ class RotaryEmbedding(MultiPlatformOp):
            else:
                cos_sin = self.cos_sin_cache.index_select(0, positions)
-            if (
+            if fused_rope_qk_mqa is not None and query.shape[0] < 65535:
                fused_rope_qk_mqa is not None
                and query.shape[0] * query.shape[1] < 65535
            ):
                return fused_rope_qk_mqa(
                    query,
                    key,