feat(linked): add vLLM moe_wna16_gemm provider

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

@@ -0,0 +1,53 @@
+#include "linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/stack.h>
+#include <c10/util/Exception.h>
+
+#include <utility>
+
+namespace infini::ops::linked::torch::nvidia {
+
+void VllmMoeWna16Gemm::Call(at::Tensor input, at::Tensor output,
+                            at::Tensor b_qweight, at::Tensor b_scales,
+                            std::optional<at::Tensor> b_qzeros,
+                            std::optional<at::Tensor> topk_weights,
+                            at::Tensor sorted_token_ids, at::Tensor expert_ids,
+                            at::Tensor num_tokens_post_pad, int64_t top_k,
+                            int64_t block_size_m, int64_t block_size_n,
+                            int64_t block_size_k, int64_t bit) {
+  static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow(
+      "_moe_C::moe_wna16_gemm", "");
+  c10::Stack stack;
+  stack.reserve(14);
+  stack.emplace_back(std::move(input));
+  stack.emplace_back(output);
+  stack.emplace_back(std::move(b_qweight));
+  stack.emplace_back(std::move(b_scales));
+  stack.emplace_back(b_qzeros ? c10::IValue(std::move(*b_qzeros))
+                              : c10::IValue());
+  stack.emplace_back(topk_weights ? c10::IValue(std::move(*topk_weights))
+                                  : c10::IValue());
+  stack.emplace_back(std::move(sorted_token_ids));
+  stack.emplace_back(std::move(expert_ids));
+  stack.emplace_back(std::move(num_tokens_post_pad));
+  stack.emplace_back(top_k);
+  stack.emplace_back(block_size_m);
+  stack.emplace_back(block_size_n);
+  stack.emplace_back(block_size_k);
+  stack.emplace_back(bit);
+  op.callBoxed(&stack);
+
+  TORCH_CHECK(stack.size() == 1,
+              "`moe_wna16_gemm` returned an unexpected number of values");
+  (void)std::move(stack.front()).toTensor();
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchMoeWna16Gemm<
+    ::infini::ops::linked::torch::nvidia::VllmMoeWna16Gemm>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,46 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_
+
+#include <optional>
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/moe_wna16_gemm.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct VllmMoeWna16Gemm : C10<Device::Type::kNvidia> {
+  static void Call(at::Tensor input, at::Tensor output, at::Tensor b_qweight,
+                   at::Tensor b_scales, std::optional<at::Tensor> b_qzeros,
+                   std::optional<at::Tensor> topk_weights,
+                   at::Tensor sorted_token_ids, at::Tensor expert_ids,
+                   at::Tensor num_tokens_post_pad, int64_t top_k,
+                   int64_t block_size_m, int64_t block_size_n,
+                   int64_t block_size_k, int64_t bit);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchMoeWna16Gemm<
+    ::infini::ops::linked::torch::nvidia::VllmMoeWna16Gemm>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<MoeWna16Gemm, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchMoeWna16Gemm<
+          linked::torch::nvidia::VllmMoeWna16Gemm> {
+ public:
+  using linked::torch::TorchMoeWna16Gemm<
+      linked::torch::nvidia::VllmMoeWna16Gemm>::TorchMoeWna16Gemm;
+
+  using linked::torch::TorchMoeWna16Gemm<
+      linked::torch::nvidia::VllmMoeWna16Gemm>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_

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

@@ -0,0 +1,8 @@
+library: vllm_moe
+operator_schema: >-
+  _moe_C::moe_wna16_gemm(Tensor input, Tensor! output,
+  Tensor b_qweight, Tensor b_scales, Tensor? b_qzeros, Tensor? topk_weights,
+  Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_post_pad,
+  int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, int bit)
+  -> Tensor
+dispatch_key: CUDA

src/linked/torch/ops/moe_wna16_gemm.h

@@ -0,0 +1,115 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_
+
+#include <optional>
+#include <utility>
+
+#include "base/moe_wna16_gemm.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchMoeWna16Gemm : public ::infini::ops::MoeWna16Gemm {
+ public:
+  TorchMoeWna16Gemm(const Tensor input, const Tensor b_qweight,
+                    const Tensor b_scales, std::optional<Tensor> b_qzeros,
+                    std::optional<Tensor> topk_weights,
+                    const Tensor sorted_token_ids, const Tensor expert_ids,
+                    const Tensor num_tokens_post_pad, const int64_t top_k,
+                    const int64_t block_size_m, const int64_t block_size_n,
+                    const int64_t block_size_k, const int64_t bit,
+                    Tensor output)
+      : ::infini::ops::MoeWna16Gemm{input,        b_qweight,
+                                    b_scales,     b_qzeros,
+                                    topk_weights, sorted_token_ids,
+                                    expert_ids,   num_tokens_post_pad,
+                                    top_k,        block_size_m,
+                                    block_size_n, block_size_k,
+                                    bit,          output},
+        input_metadata_{input},
+        b_qweight_metadata_{b_qweight},
+        b_scales_metadata_{b_scales},
+        b_qzeros_metadata_{b_qzeros},
+        topk_weights_metadata_{topk_weights},
+        sorted_token_ids_metadata_{sorted_token_ids},
+        expert_ids_metadata_{expert_ids},
+        num_tokens_post_pad_metadata_{num_tokens_post_pad},
+        output_metadata_{output} {}
+
+  using ::infini::ops::MoeWna16Gemm::operator();
+
+  void operator()(const Tensor input, const Tensor b_qweight,
+                  const Tensor b_scales, std::optional<Tensor> b_qzeros,
+                  std::optional<Tensor> topk_weights,
+                  const Tensor sorted_token_ids, const Tensor expert_ids,
+                  const Tensor num_tokens_post_pad, const int64_t top_k,
+                  const int64_t block_size_m, const int64_t block_size_n,
+                  const int64_t block_size_k, const int64_t bit,
+                  Tensor output) const override {
+    ValidateCallMetadata(input, b_qweight, b_scales, b_qzeros, topk_weights,
+                         sorted_token_ids, expert_ids, num_tokens_post_pad,
+                         top_k, block_size_m, block_size_n, block_size_k, bit,
+                         output);
+
+    const typename Backend::StreamGuard stream_guard{
+        Backend::GetStreamFromExternal(stream_, device_index_)};
+    auto at_input = ToAten(input, input_metadata_);
+    auto at_b_qweight = ToAten(b_qweight, b_qweight_metadata_);
+    auto at_b_scales = ToAten(b_scales, b_scales_metadata_);
+    auto at_b_qzeros = ToOptionalAten(b_qzeros, b_qzeros_metadata_);
+    auto at_topk_weights = ToOptionalAten(topk_weights, topk_weights_metadata_);
+    auto at_sorted_token_ids =
+        ToAten(sorted_token_ids, sorted_token_ids_metadata_);
+    auto at_expert_ids = ToAten(expert_ids, expert_ids_metadata_);
+    auto at_num_tokens_post_pad =
+        ToAten(num_tokens_post_pad, num_tokens_post_pad_metadata_);
+    auto at_output = ToAten(output, output_metadata_);
+
+    Backend::Call(std::move(at_input), std::move(at_output),
+                  std::move(at_b_qweight), std::move(at_b_scales),
+                  std::move(at_b_qzeros), std::move(at_topk_weights),
+                  std::move(at_sorted_token_ids), std::move(at_expert_ids),
+                  std::move(at_num_tokens_post_pad), top_k, block_size_m,
+                  block_size_n, block_size_k, bit);
+  }
+
+ private:
+  at::Tensor ToAten(const Tensor tensor, const Tensor metadata) const {
+    return ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(tensor.data()), metadata.shape(), metadata.strides(),
+        metadata.dtype(), device_index_);
+  }
+
+  std::optional<at::Tensor> ToOptionalAten(
+      const std::optional<Tensor>& tensor,
+      const std::optional<Tensor>& metadata) const {
+    if (!tensor) {
+      return std::nullopt;
+    }
+
+    return ToAten(*tensor, *metadata);
+  }
+
+  Tensor input_metadata_;
+
+  Tensor b_qweight_metadata_;
+
+  Tensor b_scales_metadata_;
+
+  std::optional<Tensor> b_qzeros_metadata_;
+
+  std::optional<Tensor> topk_weights_metadata_;
+
+  Tensor sorted_token_ids_metadata_;
+
+  Tensor expert_ids_metadata_;
+
+  Tensor num_tokens_post_pad_metadata_;
+
+  Tensor output_metadata_;
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_

tests/test_moe_wna16_gemm.py

@@ -229,6 +229,64 @@ def test_moe_wna16_gemm_skips_padding_and_unmapped_expert(device, implementation
     assert torch.count_nonzero(output.flatten(0, 1)[3:]).item() == 0
 
 
+@pytest.mark.parametrize("device, implementation_index", (("cuda", 16),))
+def test_moe_wna16_gemm_linked_vllm(device, implementation_index):
+    if (
+        implementation_index
+        not in infini.ops.MoeWna16Gemm.active_implementation_indices(device)
+    ):
+        pytest.skip("vLLM linked implementation is not active")
+
+    tensors = _make_case(4, True, True, torch.float16, device)
+    (
+        input,
+        qweight,
+        scales,
+        qzeros,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        output,
+        quantized_weight,
+        zero_points,
+    ) = tensors
+
+    result = infini.ops.moe_wna16_gemm(
+        input,
+        qweight,
+        scales,
+        qzeros,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        2,
+        2,
+        16,
+        32,
+        4,
+        output,
+        stream=get_stream(input.device),
+        implementation_index=implementation_index,
+    )
+
+    assert result is None
+    expected = _reference(
+        input,
+        quantized_weight,
+        scales,
+        zero_points,
+        4,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        block_size_m=2,
+    )
+    torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2)
+
+
 def _make_case(
     bit,
     has_zero_point,