diff --git a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.cc b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.cc new file mode 100644 index 000000000..b2dc2bb79 --- /dev/null +++ b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.cc @@ -0,0 +1,50 @@ +#include "linked/torch/nvidia/ops/topk_sigmoid/vllm.h" + +#include +#include +#include + +#include + +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 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 diff --git a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h new file mode 100644 index 000000000..c825f7b1c --- /dev/null +++ b/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 + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/topk_sigmoid.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmTopkSigmoid : C10 { + 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 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 + : 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_ diff --git a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.yaml b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.yaml new file mode 100644 index 000000000..48d9996ab --- /dev/null +++ b/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 diff --git a/src/linked/torch/ops/topk_sigmoid.h b/src/linked/torch/ops/topk_sigmoid.h new file mode 100644 index 000000000..c7f0e7219 --- /dev/null +++ b/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 +#include + +#include "base/topk_sigmoid.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +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 e_score_correction_bias, + std::optional 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( + const_cast(gating_output.data()), gating_output.shape(), + gating_output.strides(), gating_output.dtype(), device_index_); + std::optional at_e_score_correction_bias; + if (e_score_correction_bias.has_value()) { + at_e_score_correction_bias.emplace(ToAtenTensor( + const_cast(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( + topk_weights.data(), topk_weights.shape(), topk_weights.strides(), + topk_weights.dtype(), device_index_); + auto at_topk_ids = ToAtenTensor( + topk_ids.data(), topk_ids.shape(), topk_ids.strides(), topk_ids.dtype(), + device_index_); + auto at_token_expert_indices = ToAtenTensor( + 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_ diff --git a/tests/test_topk_sigmoid.py b/tests/test_topk_sigmoid.py index b9ac3fbc2..54f3439d0 100644 --- a/tests/test_topk_sigmoid.py +++ b/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(