From 3bbd6cb0a87650fb6408634c2d7839d2a0b8f8b6 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sat, 8 Aug 2026 00:17:04 +0800 Subject: [PATCH] feat(linked): add vLLM moe wna16 gemm provider --- .../torch/nvidia/ops/moe_wna16_gemm/vllm.cc | 53 ++++++++ .../torch/nvidia/ops/moe_wna16_gemm/vllm.h | 46 +++++++ .../torch/nvidia/ops/moe_wna16_gemm/vllm.yaml | 8 ++ src/linked/torch/ops/moe_wna16_gemm.h | 115 ++++++++++++++++++ tests/test_moe_wna16_gemm.py | 58 +++++++++ 5 files changed, 280 insertions(+) create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.cc create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.yaml create mode 100644 src/linked/torch/ops/moe_wna16_gemm.h diff --git a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.cc b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.cc new file mode 100644 index 000000000..031909602 --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.cc @@ -0,0 +1,53 @@ +#include "linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h" + +#include +#include +#include + +#include + +namespace infini::ops::linked::torch::nvidia { + +void VllmMoeWna16Gemm::Call(at::Tensor input, at::Tensor output, + at::Tensor b_qweight, at::Tensor b_scales, + std::optional b_qzeros, + std::optional topk_weights, + at::Tensor sorted_token_ids, at::Tensor expert_ids, + at::Tensor num_tokens_post_pad, int64_t top_k, + int64_t block_size_m, int64_t block_size_n, + int64_t block_size_k, int64_t bit) { + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "_moe_C::moe_wna16_gemm", ""); + c10::Stack stack; + stack.reserve(14); + stack.emplace_back(std::move(input)); + stack.emplace_back(output); + stack.emplace_back(std::move(b_qweight)); + stack.emplace_back(std::move(b_scales)); + stack.emplace_back(b_qzeros ? c10::IValue(std::move(*b_qzeros)) + : c10::IValue()); + stack.emplace_back(topk_weights ? c10::IValue(std::move(*topk_weights)) + : c10::IValue()); + stack.emplace_back(std::move(sorted_token_ids)); + stack.emplace_back(std::move(expert_ids)); + stack.emplace_back(std::move(num_tokens_post_pad)); + stack.emplace_back(top_k); + stack.emplace_back(block_size_m); + stack.emplace_back(block_size_n); + stack.emplace_back(block_size_k); + stack.emplace_back(bit); + op.callBoxed(&stack); + + TORCH_CHECK(stack.size() == 1, + "`moe_wna16_gemm` returned an unexpected number of values"); + (void)std::move(stack.front()).toTensor(); +} + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +template class TorchMoeWna16Gemm< + ::infini::ops::linked::torch::nvidia::VllmMoeWna16Gemm>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h new file mode 100644 index 000000000..de361ad18 --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h @@ -0,0 +1,46 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_ + +#include + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/moe_wna16_gemm.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmMoeWna16Gemm : C10 { + static void Call(at::Tensor input, at::Tensor output, at::Tensor b_qweight, + at::Tensor b_scales, std::optional b_qzeros, + std::optional topk_weights, + at::Tensor sorted_token_ids, at::Tensor expert_ids, + at::Tensor num_tokens_post_pad, int64_t top_k, + int64_t block_size_m, int64_t block_size_n, + int64_t block_size_k, int64_t bit); +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchMoeWna16Gemm< + ::infini::ops::linked::torch::nvidia::VllmMoeWna16Gemm>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchMoeWna16Gemm< + linked::torch::nvidia::VllmMoeWna16Gemm> { + public: + using linked::torch::TorchMoeWna16Gemm< + linked::torch::nvidia::VllmMoeWna16Gemm>::TorchMoeWna16Gemm; + + using linked::torch::TorchMoeWna16Gemm< + linked::torch::nvidia::VllmMoeWna16Gemm>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_GEMM_VLLM_H_ diff --git a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.yaml b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.yaml new file mode 100644 index 000000000..b41586f6c --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.yaml @@ -0,0 +1,8 @@ +library: vllm_moe +operator_schema: >- + _moe_C::moe_wna16_gemm(Tensor input, Tensor! output, + Tensor b_qweight, Tensor b_scales, Tensor? b_qzeros, Tensor? topk_weights, + Tensor sorted_token_ids, Tensor expert_ids, Tensor num_tokens_post_pad, + int top_k, int BLOCK_SIZE_M, int BLOCK_SIZE_N, int BLOCK_SIZE_K, int bit) + -> Tensor +dispatch_key: CUDA diff --git a/src/linked/torch/ops/moe_wna16_gemm.h b/src/linked/torch/ops/moe_wna16_gemm.h new file mode 100644 index 000000000..796d5d19f --- /dev/null +++ b/src/linked/torch/ops/moe_wna16_gemm.h @@ -0,0 +1,115 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_ + +#include +#include + +#include "base/moe_wna16_gemm.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchMoeWna16Gemm : public ::infini::ops::MoeWna16Gemm { + public: + TorchMoeWna16Gemm(const Tensor input, const Tensor b_qweight, + const Tensor b_scales, std::optional b_qzeros, + std::optional topk_weights, + const Tensor sorted_token_ids, const Tensor expert_ids, + const Tensor num_tokens_post_pad, const int64_t top_k, + const int64_t block_size_m, const int64_t block_size_n, + const int64_t block_size_k, const int64_t bit, + Tensor output) + : ::infini::ops::MoeWna16Gemm{input, b_qweight, + b_scales, b_qzeros, + topk_weights, sorted_token_ids, + expert_ids, num_tokens_post_pad, + top_k, block_size_m, + block_size_n, block_size_k, + bit, output}, + input_metadata_{input}, + b_qweight_metadata_{b_qweight}, + b_scales_metadata_{b_scales}, + b_qzeros_metadata_{b_qzeros}, + topk_weights_metadata_{topk_weights}, + sorted_token_ids_metadata_{sorted_token_ids}, + expert_ids_metadata_{expert_ids}, + num_tokens_post_pad_metadata_{num_tokens_post_pad}, + output_metadata_{output} {} + + using ::infini::ops::MoeWna16Gemm::operator(); + + void operator()(const Tensor input, const Tensor b_qweight, + const Tensor b_scales, std::optional b_qzeros, + std::optional topk_weights, + const Tensor sorted_token_ids, const Tensor expert_ids, + const Tensor num_tokens_post_pad, const int64_t top_k, + const int64_t block_size_m, const int64_t block_size_n, + const int64_t block_size_k, const int64_t bit, + Tensor output) const override { + ValidateCallMetadata(input, b_qweight, b_scales, b_qzeros, topk_weights, + sorted_token_ids, expert_ids, num_tokens_post_pad, + top_k, block_size_m, block_size_n, block_size_k, bit, + output); + + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + auto at_input = ToAten(input, input_metadata_); + auto at_b_qweight = ToAten(b_qweight, b_qweight_metadata_); + auto at_b_scales = ToAten(b_scales, b_scales_metadata_); + auto at_b_qzeros = ToOptionalAten(b_qzeros, b_qzeros_metadata_); + auto at_topk_weights = ToOptionalAten(topk_weights, topk_weights_metadata_); + auto at_sorted_token_ids = + ToAten(sorted_token_ids, sorted_token_ids_metadata_); + auto at_expert_ids = ToAten(expert_ids, expert_ids_metadata_); + auto at_num_tokens_post_pad = + ToAten(num_tokens_post_pad, num_tokens_post_pad_metadata_); + auto at_output = ToAten(output, output_metadata_); + + Backend::Call(std::move(at_input), std::move(at_output), + std::move(at_b_qweight), std::move(at_b_scales), + std::move(at_b_qzeros), std::move(at_topk_weights), + std::move(at_sorted_token_ids), std::move(at_expert_ids), + std::move(at_num_tokens_post_pad), top_k, block_size_m, + block_size_n, block_size_k, bit); + } + + private: + at::Tensor ToAten(const Tensor tensor, const Tensor metadata) const { + return ToAtenTensor( + const_cast(tensor.data()), metadata.shape(), metadata.strides(), + metadata.dtype(), device_index_); + } + + std::optional ToOptionalAten( + const std::optional& tensor, + const std::optional& metadata) const { + if (!tensor) { + return std::nullopt; + } + + return ToAten(*tensor, *metadata); + } + + Tensor input_metadata_; + + Tensor b_qweight_metadata_; + + Tensor b_scales_metadata_; + + std::optional b_qzeros_metadata_; + + std::optional topk_weights_metadata_; + + Tensor sorted_token_ids_metadata_; + + Tensor expert_ids_metadata_; + + Tensor num_tokens_post_pad_metadata_; + + Tensor output_metadata_; +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_GEMM_H_ diff --git a/tests/test_moe_wna16_gemm.py b/tests/test_moe_wna16_gemm.py index fbeff05fb..6c3e98713 100644 --- a/tests/test_moe_wna16_gemm.py +++ b/tests/test_moe_wna16_gemm.py @@ -229,6 +229,64 @@ def test_moe_wna16_gemm_skips_padding_and_unmapped_expert(device, implementation assert torch.count_nonzero(output.flatten(0, 1)[3:]).item() == 0 +@pytest.mark.parametrize("device, implementation_index", (("cuda", 16),)) +def test_moe_wna16_gemm_linked_vllm(device, implementation_index): + if ( + implementation_index + not in infini.ops.MoeWna16Gemm.active_implementation_indices(device) + ): + pytest.skip("vLLM linked implementation is not active") + + tensors = _make_case(4, True, True, torch.float16, device) + ( + input, + qweight, + scales, + qzeros, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_pad, + output, + quantized_weight, + zero_points, + ) = tensors + + result = infini.ops.moe_wna16_gemm( + input, + qweight, + scales, + qzeros, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_pad, + 2, + 2, + 16, + 32, + 4, + output, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + assert result is None + expected = _reference( + input, + quantized_weight, + scales, + zero_points, + 4, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_pad, + block_size_m=2, + ) + torch.testing.assert_close(output, expected, rtol=2e-2, atol=2e-2) + + def _make_case( bit, has_zero_point,