minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Haotong Zou
2026-08-27 09:48:07 +08:00
committed by GitHub
parent d2b3e9051d
commit b658d64d5b
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cpu_graph_runner.py
+3 -3
View File
@@ -281,9 +281,9 @@ def register_fake_ops(tp_size: int):
@register_cpu_compile_fake("rotary_embedding_cpu") @register_cpu_compile_fake("rotary_embedding_cpu")
def _(positions, query, key, head_size, cos_sin_cache, is_neox): 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, if query.ndim == 3:
# which no schema expresses; an accurate fake needs it to pick one return torch.empty_like(query), torch.empty_like(key)
return torch.empty_like(query), torch.empty_like(key) return query, key
@register_cpu_compile_fake("apply_rotary_pos_emb_cpu") @register_cpu_compile_fake("apply_rotary_pos_emb_cpu")
def _(query, key, cos, sin): def _(query, key, cos, sin):