[CPU] Fix rotary_embedding_cpu fake for in-place layouts (#36443)

2026-08-27 09:48:07 +08:00
parent d2b3e9051d
commit b658d64d5b
1 changed files with 3 additions and 3 deletions
@@ -281,9 +281,9 @@ def register_fake_ops(tp_size: int):

    @register_cpu_compile_fake("rotary_embedding_cpu")
    def _(positions, query, key, head_size, cos_sin_cache, is_neox):
-        # TODO: the kernel aliases query/key for 2D and 4D but allocates for 3D,
-        # which no schema expresses; an accurate fake needs it to pick one
-        return torch.empty_like(query), torch.empty_like(key)
+        if query.ndim == 3:
+            return torch.empty_like(query), torch.empty_like(key)
+        return query, key

    @register_cpu_compile_fake("apply_rotary_pos_emb_cpu")
    def _(query, key, cos, sin):