feat(linked): add vLLM grouped_topk provider

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

@@ -0,0 +1,49 @@
+#include "linked/torch/nvidia/ops/grouped_topk/vllm.h"
+
+#include <ATen/Functions.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 {
+
+std::pair<at::Tensor, at::Tensor> VllmGroupedTopk::Call(
+    at::Tensor scores, at::Tensor bias, int64_t num_expert_group,
+    int64_t topk_group, int64_t topk, bool renormalize,
+    double routed_scaling_factor, int64_t scoring_func) {
+  TORCH_CHECK(scoring_func == 0 || scoring_func == 1,
+              "Linked vLLM `grouped_topk` requires `scoring_func` 0 (none) "
+              "or 1 (sigmoid).");
+
+  auto routed_scores = scoring_func == 0 ? scores : at::sigmoid(scores);
+  auto scores_with_bias =
+      (routed_scores + bias.unsqueeze(0)).to(routed_scores.scalar_type());
+
+  static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow(
+      "_moe_C::grouped_topk", "");
+  c10::Stack stack;
+  stack.emplace_back(std::move(routed_scores));
+  stack.emplace_back(std::move(scores_with_bias));
+  stack.emplace_back(num_expert_group);
+  stack.emplace_back(topk_group);
+  stack.emplace_back(topk);
+  stack.emplace_back(renormalize);
+  stack.emplace_back(routed_scaling_factor);
+  op.callBoxed(&stack);
+
+  TORCH_CHECK(stack.size() == 2,
+              "Linked vLLM `grouped_topk` returned an unexpected number of "
+              "values.");
+  return {std::move(stack[0]).toTensor(), std::move(stack[1]).toTensor()};
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchGroupedTopk<
+    ::infini::ops::linked::torch::nvidia::VllmGroupedTopk>;
+
+}  // namespace infini::ops::linked::torch

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

@@ -0,0 +1,43 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GROUPED_TOPK_VLLM_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GROUPED_TOPK_VLLM_H_
+
+#include <utility>
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/grouped_topk.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct VllmGroupedTopk : C10<Device::Type::kNvidia> {
+  static std::pair<at::Tensor, at::Tensor> Call(
+      at::Tensor scores, at::Tensor bias, int64_t num_expert_group,
+      int64_t topk_group, int64_t topk, bool renormalize,
+      double routed_scaling_factor, int64_t scoring_func);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchGroupedTopk<
+    ::infini::ops::linked::torch::nvidia::VllmGroupedTopk>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<GroupedTopk, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchGroupedTopk<
+          linked::torch::nvidia::VllmGroupedTopk> {
+ public:
+  using linked::torch::TorchGroupedTopk<
+      linked::torch::nvidia::VllmGroupedTopk>::TorchGroupedTopk;
+
+  using linked::torch::TorchGroupedTopk<
+      linked::torch::nvidia::VllmGroupedTopk>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GROUPED_TOPK_VLLM_H_

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

@@ -0,0 +1,6 @@
+library: vllm_moe
+operator_schema: >-
+  _moe_C::grouped_topk(Tensor scores, Tensor scores_with_bias, int n_group,
+  int topk_group, int topk, bool renormalize, float routed_scaling_factor) ->
+  (Tensor, Tensor)
+dispatch_key: CUDA

src/linked/torch/ops/grouped_topk.h

@@ -0,0 +1,53 @@
+#ifndef INFINI_OPS_LINKED_TORCH_OPS_GROUPED_TOPK_H_
+#define INFINI_OPS_LINKED_TORCH_OPS_GROUPED_TOPK_H_
+
+#include <utility>
+
+#include "base/grouped_topk.h"
+#include "torch/tensor_.h"
+
+namespace infini::ops::linked::torch {
+
+template <typename Backend>
+class TorchGroupedTopk : public ::infini::ops::GroupedTopk {
+ public:
+  using ::infini::ops::GroupedTopk::GroupedTopk;
+
+  using ::infini::ops::GroupedTopk::operator();
+
+  void operator()(const Tensor scores, const Tensor bias,
+                  const int64_t num_expert_group, const int64_t topk_group,
+                  const int64_t topk, const bool renormalize,
+                  const double routed_scaling_factor,
+                  const int64_t scoring_func, Tensor topk_values,
+                  Tensor topk_indices) const override {
+    ValidateCallMetadata(scores, bias, num_expert_group, topk_group, topk,
+                         renormalize, routed_scaling_factor, scoring_func,
+                         topk_values, topk_indices);
+
+    const typename Backend::StreamGuard stream_guard{
+        Backend::GetStreamFromExternal(stream_, device_index_)};
+    auto at_scores = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(scores.data()), scores.shape(), scores.strides(),
+        scores.dtype(), device_index_);
+    auto at_bias = ToAtenTensor<Backend::kDeviceType>(
+        const_cast<void*>(bias.data()), bias.shape(), bias.strides(),
+        bias.dtype(), device_index_);
+    auto at_topk_values = ToAtenTensor<Backend::kDeviceType>(
+        topk_values.data(), topk_values.shape(), topk_values.strides(),
+        topk_values.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 [provider_values, provider_indices] = Backend::Call(
+        std::move(at_scores), std::move(at_bias), num_expert_group, topk_group,
+        topk, renormalize, routed_scaling_factor, scoring_func);
+    at_topk_values.copy_(provider_values);
+    at_topk_indices.copy_(provider_indices);
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_OPS_GROUPED_TOPK_H_

tests/test_grouped_topk.py

@@ -9,6 +9,9 @@
 from tests.utils import get_stream
 
 
+_VLLM_IMPLEMENTATION_INDEX = 16
+
+
 if not hasattr(infini.ops, "GroupedTopk"):
     pytest.skip(
         "`GroupedTopk` is not available on this platform", allow_module_level=True
@@ -47,7 +50,6 @@ def test_grouped_topk(
     routed_scaling_factor,
     scoring_func,
     device,
-    implementation_index,
 ):
     if device != "cuda":
         pytest.skip("`grouped_topk` requires the NVIDIA backend")
@@ -68,7 +70,6 @@ def test_grouped_topk(
         topk_values,
         topk_indices,
         stream=get_stream(scores.device),
-        implementation_index=implementation_index,
     )
 
     expected_values, expected_indices = _reference(
@@ -89,7 +90,120 @@ def test_grouped_topk(
     torch.testing.assert_close(scores, original_scores, rtol=0, atol=0)
 
 
-def test_grouped_topk_ties_prefer_smaller_expert_indices(device, implementation_index):
+@pytest.mark.parametrize(
+    (
+        "dtype",
+        "shape",
+        "num_expert_group",
+        "topk_group",
+        "topk",
+        "renormalize",
+        "routed_scaling_factor",
+        "scoring_func",
+    ),
+    (
+        (torch.float32, (2, 16), 4, 2, 4, False, 1.0, 0),
+        (torch.float16, (3, 64), 8, 2, 4, True, 2.5, 0),
+        (torch.bfloat16, (4, 128), 8, 4, 8, True, 1.5, 1),
+    ),
+)
+def test_grouped_topk_matches_vllm_provider(
+    dtype,
+    shape,
+    num_expert_group,
+    topk_group,
+    topk,
+    renormalize,
+    routed_scaling_factor,
+    scoring_func,
+):
+    _require_vllm_implementation()
+    if not torch.cuda.is_available():
+        pytest.skip("vLLM `grouped_topk` requires an NVIDIA device")
+
+    logits = torch.arange(shape[0] * shape[1], device="cuda", dtype=torch.float32)
+    logits = ((logits * 37) % 29 - 14).reshape(shape) / 8
+    scores = logits.to(dtype)
+    bias = (((torch.arange(shape[1], device="cuda") * 11) % 13) - 6).float() / 16
+    original_scores = scores.clone()
+    topk_values = torch.empty((shape[0], topk), dtype=torch.float32, device="cuda")
+    topk_indices = torch.empty((shape[0], topk), dtype=torch.int32, device="cuda")
+
+    routed_scores = scores if scoring_func == 0 else torch.sigmoid(scores)
+    scores_with_bias = (routed_scores + bias.unsqueeze(0)).to(dtype)
+    expected_values, expected_indices = torch.ops._moe_C.grouped_topk(
+        routed_scores,
+        scores_with_bias,
+        num_expert_group,
+        topk_group,
+        topk,
+        renormalize,
+        routed_scaling_factor,
+    )
+
+    infini.ops.grouped_topk(
+        scores,
+        bias,
+        num_expert_group,
+        topk_group,
+        topk,
+        renormalize,
+        routed_scaling_factor,
+        scoring_func,
+        topk_values,
+        topk_indices,
+        stream=get_stream(scores.device),
+        implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+    )
+
+    torch.testing.assert_close(topk_values, expected_values.float(), rtol=0, atol=0)
+    torch.testing.assert_close(topk_indices, expected_indices, rtol=0, atol=0)
+    torch.testing.assert_close(scores, original_scores, rtol=0, atol=0)
+
+
+def test_grouped_topk_uses_vllm_provider_tie_semantics():
+    _require_vllm_implementation()
+    if not torch.cuda.is_available():
+        pytest.skip("vLLM `grouped_topk` requires an NVIDIA device")
+
+    scores = torch.tensor(
+        [[0.25, 0.5, 0.25, 0.5, 0.75, 0.75, 0.125, 0.125]],
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+    bias = torch.zeros(8, dtype=torch.float32, device="cuda")
+    topk_values = torch.empty((1, 4), dtype=torch.float32, device="cuda")
+    topk_indices = torch.empty((1, 4), dtype=torch.int32, device="cuda")
+    expected_values, expected_indices = torch.ops._moe_C.grouped_topk(
+        scores,
+        (scores + bias.unsqueeze(0)).to(scores.dtype),
+        2,
+        2,
+        4,
+        False,
+        1.0,
+    )
+
+    infini.ops.grouped_topk(
+        scores,
+        bias,
+        2,
+        2,
+        4,
+        False,
+        1.0,
+        0,
+        topk_values,
+        topk_indices,
+        stream=get_stream(scores.device),
+        implementation_index=_VLLM_IMPLEMENTATION_INDEX,
+    )
+
+    torch.testing.assert_close(topk_values, expected_values.float(), rtol=0, atol=0)
+    torch.testing.assert_close(topk_indices, expected_indices, rtol=0, atol=0)
+
+
+def test_grouped_topk_ties_prefer_smaller_expert_indices(device):
     if device != "cuda":
         pytest.skip("`grouped_topk` requires the NVIDIA backend")
 
@@ -109,7 +223,6 @@ def test_grouped_topk_ties_prefer_smaller_expert_indices(device, implementation_
         topk_values,
         topk_indices,
         stream=get_stream(scores.device),
-        implementation_index=implementation_index,
     )
 
     expected_indices = torch.tensor(((0, 1, 2, 3),), dtype=torch.int32, device=device)
@@ -118,7 +231,7 @@ def test_grouped_topk_ties_prefer_smaller_expert_indices(device, implementation_
     torch.testing.assert_close(topk_values, expected_values, rtol=0, atol=0)
 
 
-def test_grouped_topk_group_ties_prefer_smaller_group_ids(device, implementation_index):
+def test_grouped_topk_group_ties_prefer_smaller_group_ids(device):
     if device != "cuda":
         pytest.skip("`grouped_topk` requires the NVIDIA backend")
 
@@ -142,7 +255,6 @@ def test_grouped_topk_group_ties_prefer_smaller_group_ids(device, implementation
         topk_values,
         topk_indices,
         stream=get_stream(scores.device),
-        implementation_index=implementation_index,
     )
 
     expected_indices = torch.tensor(((0, 2),), dtype=torch.int32, device=device)
@@ -152,9 +264,7 @@ def test_grouped_topk_group_ties_prefer_smaller_group_ids(device, implementation
 
 
 @pytest.mark.parametrize("scoring_func", (0, 1))
-def test_grouped_topk_excludes_nonfinite_experts(
-    scoring_func, device, implementation_index
-):
+def test_grouped_topk_excludes_nonfinite_experts(scoring_func, device):
     if device != "cuda":
         pytest.skip("`grouped_topk` requires the NVIDIA backend")
 
@@ -178,7 +288,6 @@ def test_grouped_topk_excludes_nonfinite_experts(
         topk_values,
         topk_indices,
         stream=get_stream(scores.device),
-        implementation_index=implementation_index,
     )
 
     expected_values, expected_indices = _reference(
@@ -212,7 +321,7 @@ def test_grouped_topk_descriptor_reuses_matching_metadata(device):
     torch.testing.assert_close(reused_outputs[1], expected[1], rtol=0, atol=0)
 
 
-def test_grouped_topk_non_default_stream(device, implementation_index):
+def test_grouped_topk_non_default_stream(device):
     if device != "cuda":
         pytest.skip("non-default CUDA streams require the NVIDIA backend")
 
@@ -234,7 +343,6 @@ def test_grouped_topk_non_default_stream(device, implementation_index):
         1,
         *outputs,
         stream=stream.cuda_stream,
-        implementation_index=implementation_index,
     )
 
     stream.synchronize()
@@ -314,6 +422,14 @@ def _make_inputs(shape, scores_dtype, bias_dtype, scoring_func, device):
     return scores, bias
 
 
+def _require_vllm_implementation():
+    if (
+        _VLLM_IMPLEMENTATION_INDEX
+        not in infini.ops.GroupedTopk.active_implementation_indices("nvidia")
+    ):
+        pytest.skip("vLLM `grouped_topk` implementation is not available")
+
+
 def _make_outputs(scores, topk):
     shape = (scores.size(0), topk)
     topk_values = torch.full(