Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
af7380d8
Unverified
Commit
af7380d8
authored
Nov 01, 2024
by
youkaichao
Committed by
GitHub
Nov 01, 2024
Browse files
[torch.compile] fix cpu broken code (#9947)
Signed-off-by:
youkaichao
<
youkaichao@gmail.com
>
parent
a78dd330
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
8 additions
and
1 deletion
+8
-1
vllm/utils.py
vllm/utils.py
+8
-1
No files found.
vllm/utils.py
View file @
af7380d8
...
...
@@ -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"
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment