diff --git a/src/linked/torch/nvidia/ops/topk_softmax/vllm.cc b/src/linked/torch/nvidia/ops/topk_softmax/vllm.cc new file mode 100644 index 000000000..9093acc2e --- /dev/null +++ b/src/linked/torch/nvidia/ops/topk_softmax/vllm.cc @@ -0,0 +1,42 @@ +#include "linked/torch/nvidia/ops/topk_softmax/vllm.h" + +#include +#include + +#include +#include + +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 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 diff --git a/src/linked/torch/nvidia/ops/topk_softmax/vllm.h b/src/linked/torch/nvidia/ops/topk_softmax/vllm.h new file mode 100644 index 000000000..1463dd5a2 --- /dev/null +++ b/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 + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/topk_softmax.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmTopkSoftmax : C10 { + 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 TorchTopkSoftmax< + ::infini::ops::linked::torch::nvidia::VllmTopkSoftmax>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : 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_ diff --git a/src/linked/torch/nvidia/ops/topk_softmax/vllm.yaml b/src/linked/torch/nvidia/ops/topk_softmax/vllm.yaml new file mode 100644 index 000000000..da39044b3 --- /dev/null +++ b/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 diff --git a/src/linked/torch/nvidia/vllm_moe.yaml b/src/linked/torch/nvidia/vllm_moe.yaml new file mode 100644 index 000000000..115e7e47c --- /dev/null +++ b/src/linked/torch/nvidia/vllm_moe.yaml @@ -0,0 +1,2 @@ +python_distribution_package: vllm +library_glob: vllm/_moe_C.*.so diff --git a/src/linked/torch/ops/topk_softmax.h b/src/linked/torch/ops/topk_softmax.h new file mode 100644 index 000000000..2f9f56a06 --- /dev/null +++ b/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 + +#include +#include + +#include "base/topk_softmax.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +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 bias, + std::optional 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( + const_cast(gating_output.data()), gating_output.shape(), + gating_output.strides(), gating_output.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_indices = ToAtenTensor( + topk_indices.data(), topk_indices.shape(), topk_indices.strides(), + topk_indices.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_); + + std::optional at_bias; + if (bias.has_value()) { + at_bias.emplace(ToAtenTensor( + const_cast(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_ diff --git a/tests/test_topk_softmax.py b/tests/test_topk_softmax.py index 0e1ce64cb..825731c99 100644 --- a/tests/test_topk_softmax.py +++ b/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, + ) + """ +)