Skip to content

Commit

Permalink
[torch.compile] fix cpu broken code (vllm-project#9947)
Browse files Browse the repository at this point in the history
Signed-off-by: youkaichao <[email protected]>
  • Loading branch information
youkaichao authored Nov 2, 2024
1 parent a78dd33 commit af7380d
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion vllm/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1551,7 +1551,14 @@ def direct_register_custom_op(
"""
if is_in_doc_build():
return
schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
import torch.library
if hasattr(torch.library, "infer_schema"):
schema_str = torch.library.infer_schema(op_func,
mutates_args=mutates_args)
else:
# for pytorch 2.4
import torch._custom_op.impl
schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
my_lib = target_lib or vllm_lib
my_lib.define(op_name + schema_str)
my_lib.impl(op_name, op_func, "CUDA")
Expand Down

0 comments on commit af7380d

Please sign in to comment.