[Kernel][Misc] Use TORCH_LIBRARY instead of PYBIND11_MODULE for custom ops

[bnellnm]

This PR makes the following changes

• replaces the uses of PYBIND11_MODULE with TORCH_LIBRARY and adds schemas and meta functions (where) needed for all custom (C++/CUDA) kernels.
• The remainder of the custom operators are wrapped in _custom_ops.py, i.e. cuda_utils, cache_ops, moe, custom_ar and punica.
• The code + libraries are made to use the python stable api. See [Feature]: bind python and c++ through tools other than pybind11 #4694

Motivation

• Using TORCH_LIBRARY is the more official way of defining custom pytorch ops and will also help support use of torch.compile on vllm models. See: [Misc]: Move from using PYBIND11_MODULE macro to bind C++/CUDA kernels to python to using TORCH_LIBRARY macro neuralmagic/nm-vllm#133
• Removing the use of pybind11 enables building and linking with python's stable api. This means that only one version of each .so can be used for any version of python that supports the stable api.
• Using pytorch dispatching would allow multiple implementations to be registered side by side, e.g. cpu/cuda

Details

• The signatures of a number of kernels needed to be modified to adhere to pytorch registration, e.g. int -> int64_t, float -> double
• Explicit function schemas have been added for kernels that modify their inputs, e.g. Tensor! indicates modification of an input argument. I did my best to get these correct but it is possible I missed a few, so these could use careful review. I don't think this will be much of an issue until torch.compile starts to be used.
• Any new kernels will need to be registered properly. Generally, kernels that modify their inputs will require a hand written schema (they can't be properly inferred just by the C++ signature). And kernels that return Tensors will require a meta function. The meta functions can simply return an empty Tensor that has the proper shape, type and device. If a meta function is not provided then the operation will cause a graph break when torch compiling.
• I've wrapped all the kernels in _custom_ops.py and added a function to check for support of a particular kernel based on name (is_custom_op_supported) so optional extensions can be checked at runtime.
• I've kept the ability to import the extension shared libs rather than use torch.ops.load_library since that requires a platform and abi specific filename.
• The implementation and signature of get_graph_buffer_ipc_meta needed to be changed so that it could be registered via TORCH_LIBRARY. The handles were changed to a torch::Tensor rather than a std::vector<uint8_t>. I'm not quite sure of the device registration for the custom_ar ops. I don't know if they should be CPU or CUDA (or both).

[youkaichao]

Is it possible to directly change the bind system to solve #4694 ?

[bnellnm]

Is it possible to directly change the bind system to solve #4694 ?

I'm not sure but I can look into that. I don't think torch uses pybind internally for registering operations so it might be possible.

[bnellnm]

Is it possible to directly change the bind system to solve #4694 ?

I'm not sure but I can look into that. I don't think torch uses pybind internally for registering operations so it might be possible.

I did a little investigating and the custom_ar would be a problem since the api uses types that TORCH_LIBRARY doesn't support, e.g. std::vector of int64_t/std::string/uint8_t. Maybe some of these could be changed to torch::Tensor instead?

[youkaichao]

I did a little investigating and the custom_ar would be a problem since the api uses types that TORCH_LIBRARY doesn't support, e.g. std::vector of int64_t/std::string/uint8_t. Maybe some of these could be changed to torch::Tensor instead?

I can make this change.

Anything else?

[youkaichao]

BTW, what is the initial motivation of changing the binding system to TORCH_LIBRARY?

[bnellnm]

I did a little investigating and the custom_ar would be a problem since the api uses types that TORCH_LIBRARY doesn't support, e.g. std::vector of int64_t/std::string/uint8_t. Maybe some of these could be changed to torch::Tensor instead?

I can make this change.

Anything else?

This wasn't quite as bad as I thought. I was able to get the custom_ar ops registered with the following changes to the API, so I think I can make the implementation changes also.

using fptr_t = int64_t;
fptr_t init_custom_ar(torch::Tensor& meta, torch::Tensor& rank_data,
                      const std::vector<std::string>& handles,
                      const std::vector<int64_t>& offsets, int64_t rank,
                      bool full_nvlink);
bool should_custom_ar(torch::Tensor& inp, int64_t max_size, int64_t world_size,
                      bool full_nvlink);
void all_reduce_reg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out);
void all_reduce_unreg(fptr_t _fa, torch::Tensor& inp, torch::Tensor& reg_buffer,
                      torch::Tensor& out);
void dispose(fptr_t _fa);
int64_t meta_size();
void register_buffer(fptr_t _fa, torch::Tensor& t,
                     const std::vector<std::string>& handles,
                     const std::vector<int64_t>& offsets);
std::tuple<std::vector<std::string>, std::vector<int64_t>> get_graph_buffer_ipc_meta(
    fptr_t _fa);
void register_graph_buffers(fptr_t _fa, const std::vector<std::string>& handles,
                            const std::vector<std::vector<int64_t>>& offsets);

[bnellnm]

BTW, what is the initial motivation of changing the binding system to TORCH_LIBRARY?

I am working on a custom backed for torch.compile that will be used to do optimizations on models, e.g. epilog fusion, quantization, etc. I was running into issues with the custom operators that turned out to be solvable by using TORCH_LIBRARY instead of PYBIND11_MODULE.

[SageMoore]

Looks good, Bill. Thanks for working on this.

[bnellnm]

cc @WoosukKwon , @youkaichao , @pcmoritz , @mgoin

[youkaichao]

LGTM, my initial reuqest was to build python agnostic wheels, and this PR achieves the goal, so I give my approval. Please wait for the rest people to comment on the other pespectives.

[bnellnm]

LGTM, my initial reuqest was to build python agnostic wheels, and this PR achieves the goal, so I give my approval. Please wait for the rest people to comment on the other pespectives.

No problem. Thanks for the review!

[pcmoritz]

This looks great, thanks for doing this! Before merging, can you:

• Rename pybind to something that makes it clearer that TORCH_LIBRARY is being used? E.g. torch_library_bindings
• Investigate if there is an impact on the performance (e.g. by changing float to double and int to int64_t -- btw, why is this needed, do you have some references about this)

[bnellnm]

This looks great, thanks for doing this! Before merging, can you:

• Rename pybind to something that makes it clearer that TORCH_LIBRARY is being used? E.g. torch_library_bindings
• Investigate if there is an impact on the performance (e.g. by changing float to double and int to int64_t -- btw, why is this needed, do you have some references about this)

Pytorch only supports double and int64_t for scalars. The note is buried in this link: https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html. Also, python represents its float types as C/C++ double precision. See https://docs.python.org/3/library/stdtypes.html

I can double check performance but I imagine that pybind11 was already downcasting python double-floats to C floats.

[changyuanzhangchina]

[rank0]: torch.ops._C.rms_norm(out, input, weight, epsilon)
[rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 921, in getattr
[rank0]: raise AttributeError(
[rank0]: AttributeError: '_OpNamespace' '_C' object has no attribute 'rms_norm'

when build the vllm from source, this error happens, how to fix it?

[bnellnm]

[rank0]: torch.ops._C.rms_norm(out, input, weight, epsilon) [rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 921, in getattr [rank0]: raise AttributeError( [rank0]: AttributeError: '_OpNamespace' '_C' object has no attribute 'rms_norm'

when build the vllm from source, this error happens, how to fix it?

Hi @changyuanzhangchina , can you provide some more context? Is this from a clean build area? What were you running when you got this error?

[changyuanzhangchina]

[rank0]: torch.ops._C.rms_norm(out, input, weight, epsilon) [rank0]: File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 921, in getattr [rank0]: raise AttributeError( [rank0]: AttributeError: '_OpNamespace' '_C' object has no attribute 'rms_norm'
when build the vllm from source, this error happens, how to fix it?

Hi @changyuanzhangchina , can you provide some more context? Is this from a clean build area? What were you running when you got this error?

hi, we have found pip install -e . works well.

while python setup.py install (develop) can incur the error.

thanks.

[zhouyuan]

Hi @bnellnm, we find there's a small performance regression(5%) with this commit on CPU backend:

• on 5467ac3, the throughput performance is 502 tok/s
• on 5d7e3d0, the throughput performance is 533 tok/s

The profile data didn't show big time increase in the those kernels so I think the overhead may due to the API change itself. Not sure if this also impact the CUDA backend performance, do you happen to have some ideas on this issue?

CC @bigPYJ1151 @WoosukKwon

thanks, -yuan

[tlrmchlsmth]

@zhouyuan we can try to run this down.

Can you post repro instructions for your performance benchmark (scripts and hardware setup)?
Also if you could file an issue that would be super helpful.

[bnellnm]

Hi @bnellnm, we find there's a small performance regression(5%) with this commit on CPU backend:

• on 5467ac3, the throughput performance is 502 tok/s
• on 5d7e3d0, the throughput performance is 533 tok/s

The profile data didn't show big time increase in the those kernels so I think the overhead may due to the API change itself. Not sure if this also impact the CUDA backend performance, do you happen to have some ideas on this issue?

CC @bigPYJ1151 @WoosukKwon

thanks, -yuan

I didn't see any regressions in the CUDA based benchmarks. I don't know why this would only affect CPU performance.

[zhouyuan]

@tlrmchlsmth @bnellnm thanks for the quick check! The throughput tests(llama2-7b) are running on a 56c SPR server(4th gen Xeon with AMX instructions, link) with benchmark_throughput.py + sharedgpt dataset. the latency test using benchmark_latency.py seems fine and no such regression found in my tests.

Indeed this is an odd issue as from my understanding the tensor does not get changed, only some constants type change(int -> int64_t). Will try to collect more profile data and update back.

thanks, -yuan

[MahmoudAshraf97]

Regarding the speed issue, I have a library that has a c++ extention that was linked to python using TORCH_LIBRARY because it was ported from torchaudio, in this commit I switched TORCH_LIBRARY to PYBIND11_MODULE and I gained around 40% speedup, although the code had no intensive math or operations, only tensor reading and accessing, so my guess is that the slowdown in vllm case might be related to how fast both TORCH_LIBRARY and PYBIND11_MODULE handle memory accessing