feat(nvidia): link vLLM topk_softmax provider

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

@@ -0,0 +1,42 @@
+#include "linked/torch/nvidia/ops/topk_softmax/vllm.h"
+
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <ATen/core/stack.h>
+
+#include <cassert>
+#include <utility>
+
+namespace infini::ops::linked::torch::nvidia {
+
+void VllmTopkSoftmax::Call(at::Tensor topk_weights, at::Tensor topk_indices,
+                           at::Tensor token_expert_indices,
+                           at::Tensor gating_output, bool renormalize,
+                           std::optional<at::Tensor> bias) {
+  static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow(
+      "_moe_C::topk_softmax", "");
+  c10::Stack stack;
+  stack.reserve(6);
+  stack.emplace_back(std::move(topk_weights));
+  stack.emplace_back(std::move(topk_indices));
+  stack.emplace_back(std::move(token_expert_indices));
+  stack.emplace_back(std::move(gating_output));
+  stack.emplace_back(renormalize);
+  if (bias.has_value()) {
+    stack.emplace_back(std::move(*bias));
+  } else {
+    stack.emplace_back();
+  }
+  op.callBoxed(&stack);
+
+  assert(stack.empty() &&
+         "`topk_softmax` returned an unexpected number of values");
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchTopkSoftmax<
+    ::infini::ops::linked::torch::nvidia::VllmTopkSoftmax>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,42 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SOFTMAX_VLLM_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SOFTMAX_VLLM_H_
+
+#include <optional>
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/topk_softmax.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct VllmTopkSoftmax : C10<Device::Type::kNvidia> {
+  static void Call(at::Tensor topk_weights, at::Tensor topk_indices,
+                   at::Tensor token_expert_indices, at::Tensor gating_output,
+                   bool renormalize, std::optional<at::Tensor> bias);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchTopkSoftmax<
+    ::infini::ops::linked::torch::nvidia::VllmTopkSoftmax>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<TopkSoftmax, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchTopkSoftmax<
+          linked::torch::nvidia::VllmTopkSoftmax> {
+ public:
+  using linked::torch::TorchTopkSoftmax<
+      linked::torch::nvidia::VllmTopkSoftmax>::TorchTopkSoftmax;
+
+  using linked::torch::TorchTopkSoftmax<
+      linked::torch::nvidia::VllmTopkSoftmax>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SOFTMAX_VLLM_H_

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

@@ -0,0 +1,6 @@
+library: vllm_moe
+operator_schema: >-
+  _moe_C::topk_softmax(Tensor! topk_weights, Tensor! topk_indices,
+  Tensor! token_expert_indices, Tensor gating_output, bool renormalize,
+  Tensor? bias) -> ()
+dispatch_key: CUDA

src/linked/torch/nvidia/vllm_moe.yaml

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

src/linked/torch/ops/topk_softmax.h

@@ -0,0 +1,64 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_TOPK_SOFTMAX_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_TOPK_SOFTMAX_H_
+
+#include <c10/util/Exception.h>
+
+#include <optional>
+#include <utility>
+
+#include "base/topk_softmax.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchTopkSoftmax : public ::infini::ops::TopkSoftmax {
+ public:
+  using ::infini::ops::TopkSoftmax::TopkSoftmax;
+  using ::infini::ops::TopkSoftmax::operator();
+
+  void operator()(const Tensor gating_output, std::optional<Tensor> bias,
+                  std::optional<Tensor> is_padding, const bool renormalize,
+                  Tensor topk_weights, Tensor topk_indices,
+                  Tensor token_expert_indices) const override {
+    ValidateCallMetadata(gating_output, bias, is_padding, renormalize,
+                         topk_weights, topk_indices, token_expert_indices);
+    TORCH_CHECK(!is_padding.has_value(),
+                "Linked `topk_softmax` does not support `is_padding`");
+
+    if (num_tokens_ == 0) {
+      return;
+    }
+
+    const typename Backend::StreamGuard stream_guard{
+        Backend::GetStreamFromExternal(stream_, device_index_)};
+    auto at_gating_output = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(gating_output.data()), gating_output.shape(),
+        gating_output.strides(), gating_output.dtype(), device_index_);
+    auto at_topk_weights = ToAtenTensor<Backend::kDeviceType>(
+        topk_weights.data(), topk_weights.shape(), topk_weights.strides(),
+        topk_weights.dtype(), device_index_);
+    auto at_topk_indices = ToAtenTensor<Backend::kDeviceType>(
+        topk_indices.data(), topk_indices.shape(), topk_indices.strides(),
+        topk_indices.dtype(), device_index_);
+    auto at_token_expert_indices = ToAtenTensor<Backend::kDeviceType>(
+        token_expert_indices.data(), token_expert_indices.shape(),
+        token_expert_indices.strides(), token_expert_indices.dtype(),
+        device_index_);
+
+    std::optional<at::Tensor> at_bias;
+    if (bias.has_value()) {
+      at_bias.emplace(ToAtenTensor<Backend::kDeviceType>(
+          const_cast<void*>(bias->data()), bias->shape(), bias->strides(),
+          bias->dtype(), device_index_));
+    }
+
+    Backend::Call(std::move(at_topk_weights), std::move(at_topk_indices),
+                  std::move(at_token_expert_indices),
+                  std::move(at_gating_output), renormalize, std::move(at_bias));
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_TOPK_SOFTMAX_H_

tests/test_topk_softmax.py

@@ -1,3 +1,7 @@
+import subprocess
+import sys
+import textwrap
+
 import infini.ops
 import pytest
 import torch
@@ -318,3 +322,93 @@ def _reference(gating_output, bias, is_padding, topk, renormalize):
     ).unsqueeze(-1)
 
     return weights, indices, token_expert_indices
+
+
+_LINKED_IMPLEMENTATION_INDEX = 16
+
+
+@pytest.mark.parametrize("renormalize", (False, True))
+@pytest.mark.parametrize("has_bias", (False, True))
+@pytest.mark.parametrize("index_dtype", (torch.int32, torch.int64, torch.uint32))
+@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32))
+def test_topk_softmax_linked(dtype, index_dtype, has_bias, renormalize, device):
+    if device != "cuda":
+        pytest.skip("linked `topk_softmax` requires the NVIDIA backend")
+    if _LINKED_IMPLEMENTATION_INDEX not in (
+        infini.ops.TopkSoftmax.active_implementation_indices("nvidia")
+    ):
+        pytest.skip("linked `topk_softmax` provider is not active")
+
+    gating_output = torch.tensor(
+        (
+            (1.25, -0.5, 0.75, 2.0, -1.0),
+            (-0.25, 1.5, 0.5, -1.25, 2.25),
+            (0.125, 0.75, 2.5, 1.0, -0.75),
+        ),
+        dtype=dtype,
+        device=device,
+    )
+    bias = None
+    if has_bias:
+        bias = torch.tensor(
+            (0.0, 0.75, -0.5, -1.0, 1.25),
+            dtype=torch.float32,
+            device=device,
+        )
+    outputs = _make_outputs(gating_output, topk=2, index_dtype=index_dtype)
+
+    infini.ops.topk_softmax(
+        gating_output,
+        bias,
+        None,
+        renormalize,
+        *outputs,
+        stream=get_stream(gating_output.device),
+        implementation_index=_LINKED_IMPLEMENTATION_INDEX,
+    )
+
+    expected = _reference(gating_output, bias, None, 2, renormalize)
+    torch.testing.assert_close(outputs[0], expected[0], rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(outputs[1], expected[1].to(index_dtype), rtol=0, atol=0)
+    torch.testing.assert_close(outputs[2], expected[2], rtol=0, atol=0)
+
+
+def test_topk_softmax_linked_rejects_is_padding(device):
+    if device != "cuda":
+        pytest.skip("linked `topk_softmax` requires the NVIDIA backend")
+    if _LINKED_IMPLEMENTATION_INDEX not in (
+        infini.ops.TopkSoftmax.active_implementation_indices("nvidia")
+    ):
+        pytest.skip("linked `topk_softmax` provider is not active")
+
+    result = subprocess.run(
+        [sys.executable, "-c", _IS_PADDING_SCRIPT], capture_output=True, text=True
+    )
+
+    assert result.returncode != 0
+    assert "does not support `is_padding`" in result.stderr
+
+
+_IS_PADDING_SCRIPT = textwrap.dedent(
+    """
+    import infini.ops
+    import torch
+
+
+    gating_output = torch.randn((2, 4), dtype=torch.float16, device="cuda")
+    is_padding = torch.zeros((2,), dtype=torch.bool, device="cuda")
+    outputs = (
+        torch.empty((2, 2), dtype=torch.float32, device="cuda"),
+        torch.empty((2, 2), dtype=torch.int32, device="cuda"),
+        torch.empty((2, 2), dtype=torch.int32, device="cuda"),
+    )
+    infini.ops.topk_softmax(
+        gating_output,
+        None,
+        is_padding,
+        False,
+        *outputs,
+        implementation_index=16,
+    )
+    """
+)