refactor(nvidia)!: link vLLM awq_marlin_repack

src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.cc

@@ -0,0 +1,37 @@
+#include "linked/torch/nvidia/ops/awq_marlin_repack/vllm.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/stack.h>
+#include <c10/core/SymInt.h>
+
+#include <cassert>
+#include <utility>
+
+namespace infini::ops::linked::torch::nvidia {
+
+at::Tensor VllmAwqMarlinRepack::Call(at::Tensor b_q_weight, int64_t size_k,
+                                     int64_t size_n, int64_t num_bits,
+                                     bool is_a_8bit) {
+  static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow(
+      "_C::awq_marlin_repack", "");
+  c10::Stack stack;
+  stack.emplace_back(std::move(b_q_weight));
+  stack.emplace_back(c10::SymInt{size_k});
+  stack.emplace_back(c10::SymInt{size_n});
+  stack.emplace_back(num_bits);
+  stack.emplace_back(is_a_8bit);
+  op.callBoxed(&stack);
+
+  assert(stack.size() == 1 &&
+         "`awq_marlin_repack` returned an unexpected number of values");
+  return std::move(stack.front()).toTensor();
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchAwqMarlinRepack<
+    ::infini::ops::linked::torch::nvidia::VllmAwqMarlinRepack>;
+
+}  // namespace infini::ops::linked::torch

src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h

@@ -0,0 +1,39 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/awq_marlin_repack.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct VllmAwqMarlinRepack : C10<Device::Type::kNvidia> {
+  static at::Tensor Call(at::Tensor b_q_weight, int64_t size_k, int64_t size_n,
+                         int64_t num_bits, bool is_a_8bit);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchAwqMarlinRepack<
+    ::infini::ops::linked::torch::nvidia::VllmAwqMarlinRepack>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<AwqMarlinRepack, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchAwqMarlinRepack<
+          linked::torch::nvidia::VllmAwqMarlinRepack> {
+ public:
+  using linked::torch::TorchAwqMarlinRepack<
+      linked::torch::nvidia::VllmAwqMarlinRepack>::TorchAwqMarlinRepack;
+
+  using linked::torch::TorchAwqMarlinRepack<
+      linked::torch::nvidia::VllmAwqMarlinRepack>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_

src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.yaml

@@ -0,0 +1,5 @@
+library: vllm
+operator_schema: >-
+  _C::awq_marlin_repack(Tensor b_q_weight, SymInt size_k, SymInt size_n,
+  int num_bits, bool is_a_8bit) -> Tensor
+dispatch_key: CUDA

src/linked/torch/ops/awq_marlin_repack.h

@@ -0,0 +1,40 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_
+
+#include <utility>
+
+#include "base/awq_marlin_repack.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchAwqMarlinRepack : public ::infini::ops::AwqMarlinRepack {
+ public:
+  using ::infini::ops::AwqMarlinRepack::AwqMarlinRepack;
+  using ::infini::ops::AwqMarlinRepack::operator();
+
+  void operator()(const Tensor b_q_weight, const int64_t size_k,
+                  const int64_t size_n, const int64_t num_bits,
+                  const bool is_a_8bit, Tensor out) const override {
+    ValidateCallMetadata(b_q_weight, size_k, size_n, num_bits, is_a_8bit, out);
+
+    const typename Backend::StreamGuard stream_guard{
+        Backend::GetStreamFromExternal(stream_, device_index_)};
+    auto at_b_q_weight = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(b_q_weight.data()), b_q_weight_metadata_.shape(),
+        b_q_weight_metadata_.strides(), b_q_weight_metadata_.dtype(),
+        device_index_);
+    auto at_out = ToAtenTensor<Backend::kDeviceType>(
+        out.data(), out_metadata_.shape(), out_metadata_.strides(),
+        out_metadata_.dtype(), device_index_);
+
+    auto result = Backend::Call(std::move(at_b_q_weight), size_k, size_n,
+                                num_bits, is_a_8bit);
+    at_out.copy_(result);
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_AWQ_MARLIN_REPACK_H_