From af7380d83b0d67726a4a6c7a86766423bed6a7a8 Mon Sep 17 00:00:00 2001 From: youkaichao Date: Fri, 1 Nov 2024 23:35:47 -0700 Subject: [PATCH] [torch.compile] fix cpu broken code (#9947) Signed-off-by: youkaichao --- vllm/utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/utils.py b/vllm/utils.py index 1041120a24b3f..a742ec8d76908 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -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")