diff --git a/src/linked/torch/thead/ops/rms_norm/vllm.cc b/src/linked/torch/thead/ops/rms_norm/vllm.cc new file mode 100644 index 000000000..dcc1ab49c --- /dev/null +++ b/src/linked/torch/thead/ops/rms_norm/vllm.cc @@ -0,0 +1,20 @@ +#include "linked/torch/thead/ops/rms_norm/vllm.h" + +// T-Head PPU vLLM `_C` exports this global C++ symbol. +void rms_norm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, + double epsilon); + +namespace infini::ops::linked::torch::thead { + +void VllmRmsNorm::Call(at::Tensor& out, at::Tensor& input, at::Tensor& weight, + double epsilon) { + ::rms_norm(out, input, weight, epsilon); +} + +} // namespace infini::ops::linked::torch::thead + +namespace infini::ops::linked::torch { + +template class TorchRmsNorm< ::infini::ops::linked::torch::thead::VllmRmsNorm>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/thead/ops/rms_norm/vllm.h b/src/linked/torch/thead/ops/rms_norm/vllm.h new file mode 100644 index 000000000..99b79b875 --- /dev/null +++ b/src/linked/torch/thead/ops/rms_norm/vllm.h @@ -0,0 +1,37 @@ +#ifndef INFINI_OPS_LINKED_TORCH_THEAD_OPS_RMS_NORM_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_THEAD_OPS_RMS_NORM_VLLM_H_ + +#include "linked/torch/ops/rms_norm.h" +#include "torch/thead/c10.h" + +namespace infini::ops::linked::torch::thead { + +struct VllmRmsNorm : C10 { + static void Call(at::Tensor& out, at::Tensor& input, at::Tensor& weight, + double epsilon); +}; + +} // namespace infini::ops::linked::torch::thead + +namespace infini::ops::linked::torch { + +extern template class TorchRmsNorm< + ::infini::ops::linked::torch::thead::VllmRmsNorm>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchRmsNorm { + public: + using linked::torch::TorchRmsNorm< + linked::torch::thead::VllmRmsNorm>::TorchRmsNorm; + using linked::torch::TorchRmsNorm< + linked::torch::thead::VllmRmsNorm>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_THEAD_OPS_RMS_NORM_VLLM_H_ diff --git a/src/linked/torch/thead/ops/rms_norm/vllm.yaml b/src/linked/torch/thead/ops/rms_norm/vllm.yaml new file mode 100644 index 000000000..cfed9bbc1 --- /dev/null +++ b/src/linked/torch/thead/ops/rms_norm/vllm.yaml @@ -0,0 +1,3 @@ +library: vllm +required_symbols: + - rms_norm(at::Tensor&, at::Tensor&, at::Tensor&, double) diff --git a/src/linked/torch/thead/ops/silu_and_mul/vllm.cc b/src/linked/torch/thead/ops/silu_and_mul/vllm.cc new file mode 100644 index 000000000..4526d4d11 --- /dev/null +++ b/src/linked/torch/thead/ops/silu_and_mul/vllm.cc @@ -0,0 +1,19 @@ +#include "linked/torch/thead/ops/silu_and_mul/vllm.h" + +// T-Head PPU vLLM `_C` exports this global C++ symbol. +void silu_and_mul(at::Tensor& out, at::Tensor& input); + +namespace infini::ops::linked::torch::thead { + +void VllmSiluAndMul::Call(at::Tensor& out, at::Tensor& input) { + ::silu_and_mul(out, input); +} + +} // namespace infini::ops::linked::torch::thead + +namespace infini::ops::linked::torch { + +template class TorchSiluAndMul< + ::infini::ops::linked::torch::thead::VllmSiluAndMul>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/thead/ops/silu_and_mul/vllm.h b/src/linked/torch/thead/ops/silu_and_mul/vllm.h new file mode 100644 index 000000000..276709fa1 --- /dev/null +++ b/src/linked/torch/thead/ops/silu_and_mul/vllm.h @@ -0,0 +1,37 @@ +#ifndef INFINI_OPS_LINKED_TORCH_THEAD_OPS_SILU_AND_MUL_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_THEAD_OPS_SILU_AND_MUL_VLLM_H_ + +#include "linked/torch/ops/silu_and_mul.h" +#include "torch/thead/c10.h" + +namespace infini::ops::linked::torch::thead { + +struct VllmSiluAndMul : C10 { + static void Call(at::Tensor& out, at::Tensor& input); +}; + +} // namespace infini::ops::linked::torch::thead + +namespace infini::ops::linked::torch { + +extern template class TorchSiluAndMul< + ::infini::ops::linked::torch::thead::VllmSiluAndMul>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchSiluAndMul< + linked::torch::thead::VllmSiluAndMul> { + public: + using linked::torch::TorchSiluAndMul< + linked::torch::thead::VllmSiluAndMul>::TorchSiluAndMul; + using linked::torch::TorchSiluAndMul< + linked::torch::thead::VllmSiluAndMul>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_THEAD_OPS_SILU_AND_MUL_VLLM_H_ diff --git a/src/linked/torch/thead/ops/silu_and_mul/vllm.yaml b/src/linked/torch/thead/ops/silu_and_mul/vllm.yaml new file mode 100644 index 000000000..56fafbd6e --- /dev/null +++ b/src/linked/torch/thead/ops/silu_and_mul/vllm.yaml @@ -0,0 +1,3 @@ +library: vllm +required_symbols: + - silu_and_mul(at::Tensor&, at::Tensor&) diff --git a/src/linked/torch/thead/vllm.yaml b/src/linked/torch/thead/vllm.yaml new file mode 100644 index 000000000..7f9882f27 --- /dev/null +++ b/src/linked/torch/thead/vllm.yaml @@ -0,0 +1,2 @@ +python_distribution_package: vllm +library_glob: vllm/_C*.so diff --git a/src/torch/thead/c10.h b/src/torch/thead/c10.h new file mode 100644 index 000000000..14d2bfe4c --- /dev/null +++ b/src/torch/thead/c10.h @@ -0,0 +1,28 @@ +#ifndef INFINI_OPS_TORCH_THEAD_C10_H_ +#define INFINI_OPS_TORCH_THEAD_C10_H_ + +#include +#include +#include + +#include "torch/c10.h" + +namespace infini::ops { + +template <> +struct C10 { + static constexpr Device::Type kDeviceType = Device::Type::kThead; + + using StreamGuard = c10::cuda::CUDAStreamGuard; + + static c10::cuda::CUDAStream GetStreamFromExternal(void* stream, + int device_index) { + return c10::cuda::getStreamFromExternal( + reinterpret_cast(stream), + static_cast(device_index)); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_TORCH_THEAD_C10_H_