From b658d64d5b13af095e74bb5014392937c6398459 Mon Sep 17 00:00:00 2001 From: Haotong Zou Date: Thu, 27 Aug 2026 01:48:07 +0000 Subject: [PATCH] [CPU] Fix rotary_embedding_cpu fake for in-place layouts (#36443) --- python/sglang/srt/model_executor/cpu_graph_runner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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):