diff --git a/python/sglang/srt/model_executor/cpu_graph_runner.py b/python/sglang/srt/model_executor/cpu_graph_runner.py
index 6d7843dc0..000b960b5 100644
--- a/python/sglang/srt/model_executor/cpu_graph_runner.py
+++ b/python/sglang/srt/model_executor/cpu_graph_runner.py
@@ -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):