diff --git a/src/linked/torch/nvidia/ops/grouped_topk/vllm.cc b/src/linked/torch/nvidia/ops/grouped_topk/vllm.cc new file mode 100644 index 000000000..96ac58cd7 --- /dev/null +++ b/src/linked/torch/nvidia/ops/grouped_topk/vllm.cc @@ -0,0 +1,49 @@ +#include "linked/torch/nvidia/ops/grouped_topk/vllm.h" + +#include +#include +#include +#include + +#include + +namespace infini::ops::linked::torch::nvidia { + +std::pair 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 diff --git a/src/linked/torch/nvidia/ops/grouped_topk/vllm.h b/src/linked/torch/nvidia/ops/grouped_topk/vllm.h new file mode 100644 index 000000000..da6a1037d --- /dev/null +++ b/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 + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/grouped_topk.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmGroupedTopk : C10 { + static std::pair 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 + : 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_ diff --git a/src/linked/torch/nvidia/ops/grouped_topk/vllm.yaml b/src/linked/torch/nvidia/ops/grouped_topk/vllm.yaml new file mode 100644 index 000000000..9b234d55e --- /dev/null +++ b/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 diff --git a/src/linked/torch/ops/grouped_topk.h b/src/linked/torch/ops/grouped_topk.h new file mode 100644 index 000000000..d71f8f8fa --- /dev/null +++ b/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 + +#include "base/grouped_topk.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +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( + const_cast(scores.data()), scores.shape(), scores.strides(), + scores.dtype(), device_index_); + auto at_bias = ToAtenTensor( + const_cast(bias.data()), bias.shape(), bias.strides(), + bias.dtype(), device_index_); + auto at_topk_values = ToAtenTensor( + topk_values.data(), topk_values.shape(), topk_values.strides(), + topk_values.dtype(), device_index_); + auto at_topk_indices = ToAtenTensor( + 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_ diff --git a/tests/test_grouped_topk.py b/tests/test_grouped_topk.py index 7bc78ed66..4ecceef4d 100644 --- a/tests/test_grouped_topk.py +++ b/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(