feat(thead): add linked vLLM rms_norm and silu_and_mul

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

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<Device::Type::kThead> {
+  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<RmsNorm, Device::Type::kThead, 16>
+    : public linked::torch::TorchRmsNorm<linked::torch::thead::VllmRmsNorm> {
+ 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_

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)

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

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<Device::Type::kThead> {
+  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<SiluAndMul, Device::Type::kThead, 16>
+    : 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_

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&)

src/linked/torch/thead/vllm.yaml

@@ -0,0 +1,2 @@
+python_distribution_package: vllm
+library_glob: vllm/_C*.so

src/torch/thead/c10.h

@@ -0,0 +1,28 @@
+#ifndef INFINI_OPS_TORCH_THEAD_C10_H_
+#define INFINI_OPS_TORCH_THEAD_C10_H_
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime_api.h>
+
+#include "torch/c10.h"
+
+namespace infini::ops {
+
+template <>
+struct C10<Device::Type::kThead> {
+  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<cudaStream_t>(stream),
+        static_cast<c10::DeviceIndex>(device_index));
+  }
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_TORCH_THEAD_C10_H_