diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py
index 5c20732e0..39cb8f7e2 100644
--- a/python/sglang/srt/layers/rotary_embedding/base.py
+++ b/python/sglang/srt/layers/rotary_embedding/base.py
@@ -294,10 +294,7 @@ class RotaryEmbedding(MultiPlatformOp):
             else:
                 cos_sin = self.cos_sin_cache.index_select(0, positions)
 
-            if (
-                fused_rope_qk_mqa is not None
-                and query.shape[0] * query.shape[1] < 65535
-            ):
+            if fused_rope_qk_mqa is not None and query.shape[0] < 65535:
                 return fused_rope_qk_mqa(
                     query,
                     key,