feat(linked): add vLLM topk_sigmoid provider

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

@@ -0,0 +1,50 @@
+#include "linked/torch/nvidia/ops/topk_sigmoid/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 VllmTopkSigmoid::Validate(bool has_is_padding,
+                               double routed_scaling_factor) {
+  TORCH_CHECK(!has_is_padding,
+              "Linked `topk_sigmoid` does not support `is_padding`.");
+  TORCH_CHECK(routed_scaling_factor == 1.0,
+              "Linked `topk_sigmoid` requires "
+              "`routed_scaling_factor == 1.0`.");
+}
+
+void VllmTopkSigmoid::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_sigmoid", "");
+  c10::Stack stack;
+  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);
+
+  TORCH_CHECK(stack.empty(),
+              "Linked `topk_sigmoid` returned unexpected values.");
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchTopkSigmoid<
+    ::infini::ops::linked::torch::nvidia::VllmTopkSigmoid>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,44 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SIGMOID_VLLM_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SIGMOID_VLLM_H_
+
+#include <optional>
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/topk_sigmoid.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct VllmTopkSigmoid : C10<Device::Type::kNvidia> {
+  static void Validate(bool has_is_padding, double routed_scaling_factor);
+
+  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 TorchTopkSigmoid<
+    ::infini::ops::linked::torch::nvidia::VllmTopkSigmoid>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<TopkSigmoid, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchTopkSigmoid<
+          linked::torch::nvidia::VllmTopkSigmoid> {
+ public:
+  using linked::torch::TorchTopkSigmoid<
+      linked::torch::nvidia::VllmTopkSigmoid>::TorchTopkSigmoid;
+
+  using linked::torch::TorchTopkSigmoid<
+      linked::torch::nvidia::VllmTopkSigmoid>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_TOPK_SIGMOID_VLLM_H_

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

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

src/linked/torch/ops/topk_sigmoid.h

@@ -0,0 +1,64 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_TOPK_SIGMOID_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_TOPK_SIGMOID_H_
+
+#include <optional>
+#include <utility>
+
+#include "base/topk_sigmoid.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchTopkSigmoid : public ::infini::ops::TopkSigmoid {
+ public:
+  using ::infini::ops::TopkSigmoid::TopkSigmoid;
+
+  using ::infini::ops::TopkSigmoid::operator();
+
+  void operator()(const Tensor gating_output,
+                  std::optional<Tensor> e_score_correction_bias,
+                  std::optional<Tensor> is_padding, const bool renormalize,
+                  const double routed_scaling_factor, Tensor topk_weights,
+                  Tensor topk_ids, Tensor token_expert_indices) const override {
+    ValidateCallMetadata(gating_output, e_score_correction_bias, is_padding,
+                         renormalize, routed_scaling_factor, topk_weights,
+                         topk_ids, token_expert_indices);
+    Backend::Validate(is_padding.has_value(), routed_scaling_factor);
+    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_);
+    std::optional<at::Tensor> at_e_score_correction_bias;
+    if (e_score_correction_bias.has_value()) {
+      at_e_score_correction_bias.emplace(ToAtenTensor<Backend::kDeviceType>(
+          const_cast<void*>(e_score_correction_bias->data()),
+          e_score_correction_bias->shape(), e_score_correction_bias->strides(),
+          e_score_correction_bias->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_ids = ToAtenTensor<Backend::kDeviceType>(
+        topk_ids.data(), topk_ids.shape(), topk_ids.strides(), topk_ids.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_);
+
+    Backend::Call(std::move(at_topk_weights), std::move(at_topk_ids),
+                  std::move(at_token_expert_indices),
+                  std::move(at_gating_output), renormalize,
+                  std::move(at_e_score_correction_bias));
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_TOPK_SIGMOID_H_

tests/test_topk_sigmoid.py

@@ -5,6 +5,14 @@
 from tests.utils import get_stream
 
 
+if not hasattr(infini.ops, "TopkSigmoid"):
+    pytest.skip(
+        "`TopkSigmoid` is not available on this platform",
+        allow_module_level=True,
+    )
+
+
+_VLLM_IMPLEMENTATION_INDEX = 16
 _INDEX_DTYPES = (torch.int32, torch.int64)
 if hasattr(torch, "uint32"):
     _INDEX_DTYPES += (torch.uint32,)
@@ -193,6 +201,139 @@ def test_topk_sigmoid_empty_tokens():
     assert all(output.shape == (0, 2) for output in outputs)
 
 
+@pytest.mark.parametrize("renormalize", (False, True))
+@pytest.mark.parametrize("has_bias", (False, True))
+@pytest.mark.parametrize("index_dtype", _INDEX_DTYPES)
+@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16, torch.float32))
+def test_topk_sigmoid_vllm_provider(dtype, index_dtype, has_bias, renormalize):
+    _require_vllm_implementation()
+
+    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="cuda",
+    )
+    bias = None
+    if has_bias:
+        bias = torch.tensor(
+            (0.0, 0.75, -0.5, -1.0, 1.25),
+            dtype=torch.float32,
+            device="cuda",
+        )
+    outputs = _make_outputs(gating_output, topk=2, index_dtype=index_dtype)
+
+    result = infini.ops.topk_sigmoid(
+        gating_output,
+        bias,
+        None,
+        renormalize,
+        1.0,
+        *outputs,
+        stream=get_stream(gating_output.device),
+        implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+    )
+
+    assert result is None
+    expected = _reference(gating_output, bias, None, 2, renormalize, 1.0)
+    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)
+
+
+@pytest.mark.parametrize("unsupported", ("is_padding", "scaling"))
+def test_topk_sigmoid_vllm_provider_rejects_unsupported_arguments(unsupported):
+    _require_vllm_implementation()
+
+    gating_output = torch.randn((2, 4), dtype=torch.float32, device="cuda")
+    outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.int32)
+    is_padding = (
+        torch.zeros(2, dtype=torch.bool, device="cuda")
+        if unsupported == "is_padding"
+        else None
+    )
+    routed_scaling_factor = 0.5 if unsupported == "scaling" else 1.0
+    message = "does not support `is_padding`" if is_padding is not None else "requires"
+
+    with pytest.raises(RuntimeError, match=message):
+        infini.ops.topk_sigmoid(
+            gating_output,
+            None,
+            is_padding,
+            False,
+            routed_scaling_factor,
+            *outputs,
+            stream=get_stream(gating_output.device),
+            implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+        )
+
+
+def test_topk_sigmoid_vllm_provider_non_default_stream():
+    _require_vllm_implementation()
+
+    gating_output = torch.tensor(
+        ((1.0, -0.5, 2.0, 0.25), (-1.0, 1.5, 0.5, 2.25)),
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.int32)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+
+    infini.ops.topk_sigmoid(
+        gating_output,
+        None,
+        None,
+        False,
+        1.0,
+        *outputs,
+        stream=stream.cuda_stream,
+        implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+    )
+
+    stream.synchronize()
+    expected = _reference(gating_output, None, None, 2, False, 1.0)
+    torch.testing.assert_close(outputs[0], expected[0], rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(outputs[1], expected[1].to(torch.int32), rtol=0, atol=0)
+    torch.testing.assert_close(outputs[2], expected[2], rtol=0, atol=0)
+
+
+def test_topk_sigmoid_vllm_provider_empty_tokens():
+    _require_vllm_implementation()
+
+    gating_output = torch.empty((0, 4), dtype=torch.float16, device="cuda")
+    outputs = _make_outputs(gating_output, topk=2, index_dtype=torch.int64)
+
+    result = infini.ops.topk_sigmoid(
+        gating_output,
+        None,
+        None,
+        False,
+        1.0,
+        *outputs,
+        stream=get_stream(gating_output.device),
+        implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+    )
+
+    assert result is None
+    assert all(output.shape == (0, 2) for output in outputs)
+    probe = torch.ones((1,), dtype=torch.float32, device="cuda") + 1
+    torch.cuda.synchronize()
+    assert probe.item() == 2
+
+
+def _require_vllm_implementation():
+    if not torch.cuda.is_available():
+        pytest.skip("`topk_sigmoid` vLLM provider requires the NVIDIA backend")
+    if _VLLM_IMPLEMENTATION_INDEX not in (
+        infini.ops.TopkSigmoid.active_implementation_indices("nvidia")
+    ):
+        pytest.skip("vLLM `topk_sigmoid` provider is not active")
+
+
 def _make_outputs(gating_output, topk, index_dtype):
     shape = (gating_output.size(0), topk)
     topk_weights = torch.full(