mirror of
https://github.com/vllm-project/vllm-ascend.git
synced 2025-10-21 06:10:33 +08:00
c0f0b70813
### What this PR does / why we need it?
Thanks to the PR https://github.com/vllm-project/vllm-ascend/pull/426
make vllm-ascend support the aclgraph inference to reduce the host
overhead. However, the capability of aclgraph strongly relies on the
functionality provided by `torch.compile`, which is the key feature
supported in torch 2.x . Therefore, capture custom op into aclgraph is
only possible when it can be recognize and captured by `torch.compile`.
In this PR, we register the meta implementation of current custom ops to
enable the fx graph capture. And by doing that, insert those custom ops
into aclgraph become a natural thing to the ascend runtime.
### Does this PR introduce _any_ user-facing change?
No user face change.
### How was this patch tested?
Tested in unittest, we will integrate the `rotary_embedding` op into a
small custom model and use `torch.compile` and aclgraph to capture and
replay it to verify its functionality.
- vLLM version: v0.10.0
- vLLM main:
1b99028069
---------
Signed-off-by: ganyi <pleaplusone.gy@gmail.com>
32 lines
1.1 KiB
C
32 lines
1.1 KiB
C
#pragma once
|
|
|
|
#include "kernels/types.h"
|
|
#include <c10/core/ScalarType.h>
|
|
#include <Python.h>
|
|
|
|
#define _CONCAT(A, B) A##B
|
|
#define CONCAT(A, B) _CONCAT(A, B)
|
|
|
|
#define _STRINGIFY(A) #A
|
|
#define STRINGIFY(A) _STRINGIFY(A)
|
|
|
|
// A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
|
|
// could be a macro instead of a literal token.
|
|
#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
|
|
|
|
// A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
|
|
// could be a macro instead of a literal token.
|
|
#define TORCH_LIBRARY_IMPL_EXPAND(NAME, DEVICE, MODULE) \
|
|
TORCH_LIBRARY_IMPL(NAME, DEVICE, MODULE)
|
|
|
|
// REGISTER_EXTENSION allows the shared library to be loaded and initialized
|
|
// via python's import statement.
|
|
#define REGISTER_EXTENSION(NAME) \
|
|
PyMODINIT_FUNC CONCAT(PyInit_, NAME)() { \
|
|
static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, \
|
|
STRINGIFY(NAME), nullptr, 0, nullptr}; \
|
|
return PyModule_Create(&module); \
|
|
}
|
|
|
|
|