From 3173d65199dbb4d4d6e275d5567fb8609ae50400 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Sat, 8 Aug 2026 20:46:49 +0800 Subject: [PATCH] feat(nvidia): link vLLM `moe_wna16_marlin_gemm` --- src/base/moe_wna16_marlin_gemm.h | 362 +++++++++++++++++ .../nvidia/ops/moe_wna16_marlin_gemm/vllm.cc | 8 + .../nvidia/ops/moe_wna16_marlin_gemm/vllm.h | 130 +++++++ .../ops/moe_wna16_marlin_gemm/vllm.yaml | 12 + src/linked/torch/ops/moe_wna16_marlin_gemm.h | 81 ++++ tests/test_moe_wna16_marlin_gemm.py | 365 ++++++++++++++++++ 6 files changed, 958 insertions(+) create mode 100644 src/base/moe_wna16_marlin_gemm.h create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.cc create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h create mode 100644 src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.yaml create mode 100644 src/linked/torch/ops/moe_wna16_marlin_gemm.h create mode 100644 tests/test_moe_wna16_marlin_gemm.py diff --git a/src/base/moe_wna16_marlin_gemm.h b/src/base/moe_wna16_marlin_gemm.h new file mode 100644 index 000000000..cec5c6038 --- /dev/null +++ b/src/base/moe_wna16_marlin_gemm.h @@ -0,0 +1,362 @@ +#ifndef INFINI_OPS_BASE_MOE_WNA16_MARLIN_GEMM_H_ +#define INFINI_OPS_BASE_MOE_WNA16_MARLIN_GEMM_H_ + +#include +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// Aligned with vLLM `_moe_C::moe_wna16_marlin_gemm` at commit +// 568afb3a13806beb53bb2e6bd518269357b237c0. +class MoeWna16MarlinGemm : public Operator { + public: + MoeWna16MarlinGemm(const Tensor a, const Tensor b_q_weight, + std::optional b_bias_or_none, + const Tensor b_scales, std::optional a_scales, + std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, + std::optional perm_or_none, const Tensor workspace, + const Tensor sorted_token_ids, const Tensor expert_ids, + const Tensor num_tokens_past_padded, + const Tensor topk_weights, const int64_t moe_block_size, + const int64_t top_k, const bool mul_topk_weights, + const int64_t b_type_id, const int64_t size_m, + const int64_t size_n, const int64_t size_k, + const bool is_full_k, const bool use_atomic_add, + const bool use_fp32_reduce, const bool is_zp_float, + const int64_t thread_k, const int64_t thread_n, + const int64_t blocks_per_sm, Tensor out) + : a_metadata_{a}, + b_q_weight_metadata_{b_q_weight}, + b_bias_or_none_metadata_{b_bias_or_none}, + b_scales_metadata_{b_scales}, + a_scales_metadata_{a_scales}, + global_scale_metadata_{global_scale}, + b_zeros_or_none_metadata_{b_zeros_or_none}, + g_idx_or_none_metadata_{g_idx_or_none}, + perm_or_none_metadata_{perm_or_none}, + workspace_metadata_{workspace}, + sorted_token_ids_metadata_{sorted_token_ids}, + expert_ids_metadata_{expert_ids}, + num_tokens_past_padded_metadata_{num_tokens_past_padded}, + topk_weights_metadata_{topk_weights}, + out_metadata_{out}, + moe_block_size_{moe_block_size}, + top_k_{top_k}, + mul_topk_weights_{mul_topk_weights}, + b_type_id_{b_type_id}, + size_m_{size_m}, + size_n_{size_n}, + size_k_{size_k}, + is_full_k_{is_full_k}, + use_atomic_add_{use_atomic_add}, + use_fp32_reduce_{use_fp32_reduce}, + is_zp_float_{is_zp_float}, + thread_k_{thread_k}, + thread_n_{thread_n}, + blocks_per_sm_{blocks_per_sm}, + device_index_{a.device().index()} { + Validate(a, b_q_weight, b_bias_or_none, b_scales, a_scales, global_scale, + b_zeros_or_none, g_idx_or_none, perm_or_none, workspace, + sorted_token_ids, expert_ids, num_tokens_past_padded, topk_weights, + out); + } + + virtual void operator()( + const Tensor a, const Tensor b_q_weight, + std::optional b_bias_or_none, const Tensor b_scales, + std::optional a_scales, std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, std::optional perm_or_none, + const Tensor workspace, const Tensor sorted_token_ids, + const Tensor expert_ids, const Tensor num_tokens_past_padded, + const Tensor topk_weights, const int64_t moe_block_size, + const int64_t top_k, const bool mul_topk_weights, const int64_t b_type_id, + const int64_t size_m, const int64_t size_n, const int64_t size_k, + const bool is_full_k, const bool use_atomic_add, + const bool use_fp32_reduce, const bool is_zp_float, + const int64_t thread_k, const int64_t thread_n, + const int64_t blocks_per_sm, Tensor out) const = 0; + + protected: + void ValidateCallMetadata( + const Tensor a, const Tensor b_q_weight, + std::optional b_bias_or_none, const Tensor b_scales, + std::optional a_scales, std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, std::optional perm_or_none, + const Tensor workspace, const Tensor sorted_token_ids, + const Tensor expert_ids, const Tensor num_tokens_past_padded, + const Tensor topk_weights, const int64_t moe_block_size, + const int64_t top_k, const bool mul_topk_weights, const int64_t b_type_id, + const int64_t size_m, const int64_t size_n, const int64_t size_k, + const bool is_full_k, const bool use_atomic_add, + const bool use_fp32_reduce, const bool is_zp_float, + const int64_t thread_k, const int64_t thread_n, + const int64_t blocks_per_sm, Tensor out) const { + assert(moe_block_size == moe_block_size_ && top_k == top_k_ && + mul_topk_weights == mul_topk_weights_ && b_type_id == b_type_id_ && + size_m == size_m_ && size_n == size_n_ && size_k == size_k_ && + is_full_k == is_full_k_ && use_atomic_add == use_atomic_add_ && + use_fp32_reduce == use_fp32_reduce_ && is_zp_float == is_zp_float_ && + thread_k == thread_k_ && thread_n == thread_n_ && + blocks_per_sm == blocks_per_sm_ && + "`MoeWna16MarlinGemm` attributes changed after descriptor " + "creation"); + + const std::equal_to same_metadata; + const auto optional_matches = [&](const std::optional& expected, + const std::optional& actual) { + return expected.has_value() == actual.has_value() && + (!expected || same_metadata(*expected, *actual)); + }; + const auto matches = + same_metadata(a_metadata_, a) && + same_metadata(b_q_weight_metadata_, b_q_weight) && + optional_matches(b_bias_or_none_metadata_, b_bias_or_none) && + same_metadata(b_scales_metadata_, b_scales) && + optional_matches(a_scales_metadata_, a_scales) && + optional_matches(global_scale_metadata_, global_scale) && + optional_matches(b_zeros_or_none_metadata_, b_zeros_or_none) && + optional_matches(g_idx_or_none_metadata_, g_idx_or_none) && + optional_matches(perm_or_none_metadata_, perm_or_none) && + same_metadata(workspace_metadata_, workspace) && + same_metadata(sorted_token_ids_metadata_, sorted_token_ids) && + same_metadata(expert_ids_metadata_, expert_ids) && + same_metadata(num_tokens_past_padded_metadata_, + num_tokens_past_padded) && + same_metadata(topk_weights_metadata_, topk_weights) && + same_metadata(out_metadata_, out); + assert(matches && + "`MoeWna16MarlinGemm` tensor metadata must match descriptor"); + } + + private: + void Validate(const Tensor a, const Tensor b_q_weight, + std::optional b_bias_or_none, const Tensor b_scales, + std::optional a_scales, + std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, + std::optional perm_or_none, const Tensor workspace, + const Tensor sorted_token_ids, const Tensor expert_ids, + const Tensor num_tokens_past_padded, const Tensor topk_weights, + const Tensor out) const { + assert(a.ndim() == 2 && a.size(0) == size_m_ && a.size(1) == size_k_ && + "`MoeWna16MarlinGemm` `a` shape must match `size_m` and " + "`size_k`"); + const auto is_a_8bit = a.dtype() == DataType::kInt8; + const auto output_dtype = is_a_8bit ? b_scales.dtype() : a.dtype(); + assert( + (a.dtype() == DataType::kFloat16 || a.dtype() == DataType::kBFloat16 || + is_a_8bit) && + a.IsContiguous() && + "`MoeWna16MarlinGemm` requires contiguous float16, bfloat16, or int8 " + "`a`"); + assert(size_m_ > 0 && size_n_ > 0 && size_k_ > 0 && top_k_ > 0 && + size_k_ % 16 == 0 && size_n_ % 64 == 0 && + "`MoeWna16MarlinGemm` received unsupported dimensions"); + assert((moe_block_size_ == 8 || + (moe_block_size_ >= 16 && moe_block_size_ <= 64 && + moe_block_size_ % 16 == 0)) && + "`MoeWna16MarlinGemm` received an unsupported `moe_block_size`"); + assert(size_m_ <= std::numeric_limits::max() / top_k_ && + "`MoeWna16MarlinGemm` output dimensions overflow"); + + constexpr int64_t kUint4B8 = 1125899907892224; + constexpr int64_t kUint8B128 = 1125899923621888; + constexpr int64_t kUint4 = 1125899906843648; + constexpr int64_t kUint8 = 1125899906844672; + constexpr int64_t kInt4 = 1125899906908928; + constexpr int64_t kInt8 = 1125899906909952; + constexpr int64_t kFloat8E4M3Fn = 2814749767172868; + constexpr int64_t kFloat4E2M1F = 562949953487106; + const auto supported_qtype = + b_type_id_ == kUint4B8 || b_type_id_ == kUint8B128 || + b_type_id_ == kUint4 || b_type_id_ == kUint8 || b_type_id_ == kInt4 || + b_type_id_ == kInt8 || b_type_id_ == kFloat8E4M3Fn || + b_type_id_ == kFloat4E2M1F; + const auto has_zero_points = b_zeros_or_none.has_value(); + assert(supported_qtype && + has_zero_points == (b_type_id_ == kUint4 || b_type_id_ == kUint8) && + (!is_zp_float_ || + (has_zero_points && a.dtype() == DataType::kFloat16)) && + "`MoeWna16MarlinGemm` received an unsupported quantization " + "configuration"); + + const auto pack_factor = + (b_type_id_ == kUint8B128 || b_type_id_ == kUint8 || + b_type_id_ == kInt8 || b_type_id_ == kFloat8E4M3Fn) + ? 4 + : 8; + assert(size_n_ <= std::numeric_limits::max() / 16 && + b_q_weight.ndim() == 3 && b_q_weight.size(1) == size_k_ / 16 && + b_q_weight.size(2) == size_n_ * 16 / pack_factor && + b_q_weight.dtype() == DataType::kInt32 && + b_q_weight.IsContiguous() && + "`MoeWna16MarlinGemm` received invalid packed weights"); + assert(b_scales.ndim() == 3 && b_scales.size(0) == b_q_weight.size(0) && + b_scales.size(1) > 0 && b_scales.size(2) == size_n_ && + size_k_ % b_scales.size(1) == 0 && + (output_dtype == DataType::kFloat16 || + output_dtype == DataType::kBFloat16) && + b_scales.dtype() == output_dtype && b_scales.IsContiguous() && + "`MoeWna16MarlinGemm` received invalid weight scales"); + assert( + a_scales.has_value() == is_a_8bit && + "`MoeWna16MarlinGemm` requires activation scales exactly for int8 `a`"); + if (a_scales) { + assert(a_scales->shape() == + Tensor::Shape({static_cast(size_m_), 1}) && + a_scales->dtype() == DataType::kFloat32 && + a_scales->IsContiguous() && + "`MoeWna16MarlinGemm` received invalid activation scales"); + } + + if (b_bias_or_none) { + assert(b_bias_or_none->ndim() == 2 && + b_bias_or_none->size(0) == b_q_weight.size(0) && + b_bias_or_none->size(1) == size_n_ && + b_bias_or_none->dtype() == output_dtype && + b_bias_or_none->IsContiguous() && + "`MoeWna16MarlinGemm` received invalid bias"); + } + + if (b_zeros_or_none) { + assert(b_zeros_or_none->ndim() == 3 && + b_zeros_or_none->size(0) == b_q_weight.size(0) && + b_zeros_or_none->size(1) == b_scales.size(1) && + b_zeros_or_none->size(2) == + (is_zp_float_ ? size_n_ : size_n_ / pack_factor) && + (is_zp_float_ ? b_zeros_or_none->dtype() == output_dtype + : b_zeros_or_none->dtype() == DataType::kInt32) && + "`MoeWna16MarlinGemm` received invalid zero points"); + } + + const auto same_device = [&](const Tensor tensor) { + return tensor.device().type() == a.device().type() && + tensor.device().index() == a.device().index(); + }; + const auto valid_optional = [&](const std::optional& tensor) { + return !tensor || (tensor->IsContiguous() && same_device(*tensor)); + }; + assert(same_device(b_q_weight) && same_device(b_scales) && + valid_optional(b_bias_or_none) && valid_optional(a_scales) && + valid_optional(global_scale) && valid_optional(b_zeros_or_none) && + valid_optional(g_idx_or_none) && valid_optional(perm_or_none) && + same_device(workspace) && same_device(sorted_token_ids) && + same_device(expert_ids) && same_device(num_tokens_past_padded) && + same_device(topk_weights) && same_device(out) && + "`MoeWna16MarlinGemm` requires all tensors on the input device"); + + assert(g_idx_or_none.has_value() == perm_or_none.has_value() && + "`MoeWna16MarlinGemm` requires `g_idx_or_none` and `perm_or_none` " + "together"); + if (g_idx_or_none) { + assert(g_idx_or_none->ndim() > 0 && perm_or_none->ndim() > 0 && + g_idx_or_none->size(-1) == perm_or_none->size(-1) && + (g_idx_or_none->size(-1) == 0 || + g_idx_or_none->size(-1) == size_k_) && + g_idx_or_none->dtype() == DataType::kInt32 && + perm_or_none->dtype() == DataType::kInt32 && + (!is_full_k_ || b_scales.size(1) > 1) && + "`MoeWna16MarlinGemm` received invalid activation-order " + "metadata"); + } + + assert(workspace.ndim() == 1 && workspace.numel() > 0 && + workspace.dtype() == DataType::kInt32 && workspace.IsContiguous() && + "`MoeWna16MarlinGemm` requires a non-empty int32 workspace"); + assert(sorted_token_ids.ndim() == 1 && + sorted_token_ids.dtype() == DataType::kInt32 && + sorted_token_ids.IsContiguous() && expert_ids.ndim() == 1 && + expert_ids.dtype() == DataType::kInt32 && + expert_ids.IsContiguous() && num_tokens_past_padded.numel() == 1 && + num_tokens_past_padded.dtype() == DataType::kInt32 && + num_tokens_past_padded.IsContiguous() && + "`MoeWna16MarlinGemm` received invalid routing metadata"); + assert(topk_weights.numel() == size_m_ * top_k_ && + ((!mul_topk_weights_ && + (topk_weights.dtype() == DataType::kFloat16 || + topk_weights.dtype() == DataType::kBFloat16)) || + topk_weights.dtype() == DataType::kFloat32) && + topk_weights.IsContiguous() && + "`MoeWna16MarlinGemm` received invalid top-k weights"); + assert(out.shape() == + Tensor::Shape({static_cast(size_m_ * top_k_), + static_cast(size_n_)}) && + out.dtype() == output_dtype && out.IsContiguous() && + "`MoeWna16MarlinGemm` output metadata is invalid"); + } + + Tensor a_metadata_; + + Tensor b_q_weight_metadata_; + + std::optional b_bias_or_none_metadata_; + + Tensor b_scales_metadata_; + + std::optional a_scales_metadata_; + + std::optional global_scale_metadata_; + + std::optional b_zeros_or_none_metadata_; + + std::optional g_idx_or_none_metadata_; + + std::optional perm_or_none_metadata_; + + Tensor workspace_metadata_; + + Tensor sorted_token_ids_metadata_; + + Tensor expert_ids_metadata_; + + Tensor num_tokens_past_padded_metadata_; + + Tensor topk_weights_metadata_; + + Tensor out_metadata_; + + int64_t moe_block_size_{0}; + + int64_t top_k_{0}; + + bool mul_topk_weights_{false}; + + int64_t b_type_id_{0}; + + int64_t size_m_{0}; + + int64_t size_n_{0}; + + int64_t size_k_{0}; + + bool is_full_k_{false}; + + bool use_atomic_add_{false}; + + bool use_fp32_reduce_{false}; + + bool is_zp_float_{false}; + + int64_t thread_k_{0}; + + int64_t thread_n_{0}; + + int64_t blocks_per_sm_{0}; + + protected: + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_MOE_WNA16_MARLIN_GEMM_H_ diff --git a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.cc b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.cc new file mode 100644 index 000000000..e6331d401 --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.cc @@ -0,0 +1,8 @@ +#include "linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h" + +namespace infini::ops::linked::torch { + +template class TorchMoeWna16MarlinGemm< + ::infini::ops::linked::torch::nvidia::VllmMoeWna16MarlinGemm>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h new file mode 100644 index 000000000..f64e1ef4a --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h @@ -0,0 +1,130 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_MARLIN_GEMM_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_MARLIN_GEMM_VLLM_H_ + +#include +#include +#include + +#include +#include + +#include "linked/torch/nvidia/c10.h" +#include "linked/torch/ops/moe_wna16_marlin_gemm.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmMoeWna16MarlinGemm : C10 { + static void Validate(const DataType a_dtype, const int64_t b_type_id, + const bool has_a_scales, const bool has_global_scale) { + constexpr int64_t kFloat4E2M1F = 562949953487106; + const auto has_int8_activation = a_dtype == DataType::kInt8; + TORCH_CHECK( + a_dtype == DataType::kFloat16 || a_dtype == DataType::kBFloat16 || + has_int8_activation, + "Linked `moe_wna16_marlin_gemm` supports float16, bfloat16, and " + "int8 activations only; InfiniRT cannot represent vLLM's FP8 " + "activation dtype."); + TORCH_CHECK(has_a_scales == has_int8_activation, + "Linked `moe_wna16_marlin_gemm` requires `a_scales` " + "exactly when `a` is int8."); + TORCH_CHECK( + b_type_id != kFloat4E2M1F, + "Linked `moe_wna16_marlin_gemm` does not support `float4_e2m1f` " + "because InfiniRT cannot represent its float8 scales."); + TORCH_CHECK( + !has_global_scale, + "Linked `moe_wna16_marlin_gemm` does not support `global_scale`."); + } + + static void Call(at::Tensor a, at::Tensor out, at::Tensor b_q_weight, + std::optional b_bias_or_none, + at::Tensor b_scales, std::optional a_scales, + std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, + std::optional perm_or_none, at::Tensor workspace, + at::Tensor sorted_token_ids, at::Tensor expert_ids, + at::Tensor num_tokens_past_padded, at::Tensor topk_weights, + int64_t moe_block_size, int64_t top_k, bool mul_topk_weights, + int64_t b_type_id, int64_t size_m, int64_t size_n, + int64_t size_k, bool is_full_k, bool use_atomic_add, + bool use_fp32_reduce, bool is_zp_float, int64_t thread_k, + int64_t thread_n, int64_t blocks_per_sm) { + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "_moe_C::moe_wna16_marlin_gemm", ""); + c10::Stack stack; + stack.reserve(29); + stack.emplace_back(std::move(a)); + stack.emplace_back(out); + stack.emplace_back(std::move(b_q_weight)); + stack.emplace_back(b_bias_or_none ? c10::IValue(std::move(*b_bias_or_none)) + : c10::IValue()); + stack.emplace_back(std::move(b_scales)); + stack.emplace_back(a_scales ? c10::IValue(std::move(*a_scales)) + : c10::IValue()); + stack.emplace_back(global_scale ? c10::IValue(std::move(*global_scale)) + : c10::IValue()); + stack.emplace_back(b_zeros_or_none + ? c10::IValue(std::move(*b_zeros_or_none)) + : c10::IValue()); + stack.emplace_back(g_idx_or_none ? c10::IValue(std::move(*g_idx_or_none)) + : c10::IValue()); + stack.emplace_back(perm_or_none ? c10::IValue(std::move(*perm_or_none)) + : c10::IValue()); + stack.emplace_back(std::move(workspace)); + stack.emplace_back(std::move(sorted_token_ids)); + stack.emplace_back(std::move(expert_ids)); + stack.emplace_back(std::move(num_tokens_past_padded)); + stack.emplace_back(std::move(topk_weights)); + stack.emplace_back(moe_block_size); + stack.emplace_back(top_k); + stack.emplace_back(mul_topk_weights); + stack.emplace_back(b_type_id); + stack.emplace_back(size_m); + stack.emplace_back(size_n); + stack.emplace_back(size_k); + stack.emplace_back(is_full_k); + stack.emplace_back(use_atomic_add); + stack.emplace_back(use_fp32_reduce); + stack.emplace_back(is_zp_float); + stack.emplace_back(thread_k); + stack.emplace_back(thread_n); + stack.emplace_back(blocks_per_sm); + op.callBoxed(&stack); + + TORCH_CHECK(stack.size() == 1, + "Linked `moe_wna16_marlin_gemm` returned an unexpected " + "number of values."); + auto result = std::move(stack.front()).toTensor(); + TORCH_CHECK(result.unsafeGetTensorImpl() == out.unsafeGetTensorImpl(), + "Linked `moe_wna16_marlin_gemm` did not return the provided " + "output tensor."); + } +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchMoeWna16MarlinGemm< + ::infini::ops::linked::torch::nvidia::VllmMoeWna16MarlinGemm>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchMoeWna16MarlinGemm< + linked::torch::nvidia::VllmMoeWna16MarlinGemm> { + public: + using linked::torch::TorchMoeWna16MarlinGemm< + linked::torch::nvidia::VllmMoeWna16MarlinGemm>::TorchMoeWna16MarlinGemm; + + using linked::torch::TorchMoeWna16MarlinGemm< + linked::torch::nvidia::VllmMoeWna16MarlinGemm>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_MOE_WNA16_MARLIN_GEMM_VLLM_H_ diff --git a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.yaml b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.yaml new file mode 100644 index 000000000..4cf863bfb --- /dev/null +++ b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.yaml @@ -0,0 +1,12 @@ +library: vllm_moe +operator_schema: >- + _moe_C::moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none, + Tensor! b_q_weight, Tensor? b_bias_or_none, Tensor! b_scales, + Tensor? a_scales, Tensor? global_scale, Tensor? b_zeros_or_none, + Tensor? g_idx_or_none, Tensor? perm_or_none, Tensor! workspace, + Tensor sorted_token_ids, Tensor! expert_ids, Tensor! num_tokens_past_padded, + Tensor! topk_weights, int moe_block_size, int top_k, + bool mul_topk_weights, int b_type_id, int size_m, int size_n, int size_k, + bool is_full_k, bool use_atomic_add, bool use_fp32_reduce, + bool is_zp_float, int thread_k, int thread_n, int blocks_per_sm) -> Tensor +dispatch_key: CUDA diff --git a/src/linked/torch/ops/moe_wna16_marlin_gemm.h b/src/linked/torch/ops/moe_wna16_marlin_gemm.h new file mode 100644 index 000000000..a261d7027 --- /dev/null +++ b/src/linked/torch/ops/moe_wna16_marlin_gemm.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_MARLIN_GEMM_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_MARLIN_GEMM_H_ + +#include +#include + +#include "base/moe_wna16_marlin_gemm.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchMoeWna16MarlinGemm : public ::infini::ops::MoeWna16MarlinGemm { + public: + using ::infini::ops::MoeWna16MarlinGemm::MoeWna16MarlinGemm; + + using ::infini::ops::MoeWna16MarlinGemm::operator(); + + void operator()(const Tensor a, const Tensor b_q_weight, + std::optional b_bias_or_none, const Tensor b_scales, + std::optional a_scales, + std::optional global_scale, + std::optional b_zeros_or_none, + std::optional g_idx_or_none, + std::optional perm_or_none, const Tensor workspace, + const Tensor sorted_token_ids, const Tensor expert_ids, + const Tensor num_tokens_past_padded, + const Tensor topk_weights, const int64_t moe_block_size, + const int64_t top_k, const bool mul_topk_weights, + const int64_t b_type_id, const int64_t size_m, + const int64_t size_n, const int64_t size_k, + const bool is_full_k, const bool use_atomic_add, + const bool use_fp32_reduce, const bool is_zp_float, + const int64_t thread_k, const int64_t thread_n, + const int64_t blocks_per_sm, Tensor out) const override { + ValidateCallMetadata(a, b_q_weight, b_bias_or_none, b_scales, a_scales, + global_scale, b_zeros_or_none, g_idx_or_none, + perm_or_none, workspace, sorted_token_ids, expert_ids, + num_tokens_past_padded, topk_weights, moe_block_size, + top_k, mul_topk_weights, b_type_id, size_m, size_n, + size_k, is_full_k, use_atomic_add, use_fp32_reduce, + is_zp_float, thread_k, thread_n, blocks_per_sm, out); + + Backend::Validate(a.dtype(), b_type_id, a_scales.has_value(), + global_scale.has_value()); + + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + Backend::Call(ToAten(a), ToAten(out), ToAten(b_q_weight), + ToOptionalAten(b_bias_or_none), ToAten(b_scales), + ToOptionalAten(a_scales), ToOptionalAten(global_scale), + ToOptionalAten(b_zeros_or_none), + ToOptionalAten(g_idx_or_none), ToOptionalAten(perm_or_none), + ToAten(workspace), ToAten(sorted_token_ids), + ToAten(expert_ids), ToAten(num_tokens_past_padded), + ToAten(topk_weights), moe_block_size, top_k, mul_topk_weights, + b_type_id, size_m, size_n, size_k, is_full_k, use_atomic_add, + use_fp32_reduce, is_zp_float, thread_k, thread_n, + blocks_per_sm); + } + + private: + at::Tensor ToAten(const Tensor tensor) const { + return ToAtenTensor(const_cast(tensor.data()), + tensor.shape(), tensor.strides(), + tensor.dtype(), device_index_); + } + + std::optional ToOptionalAten( + const std::optional& tensor) const { + if (!tensor) { + return std::nullopt; + } + + return ToAten(*tensor); + } +}; + +} // namespace infini::ops::linked::torch + +#endif // INFINI_OPS_LINKED_TORCH_OPS_MOE_WNA16_MARLIN_GEMM_H_ diff --git a/tests/test_moe_wna16_marlin_gemm.py b/tests/test_moe_wna16_marlin_gemm.py new file mode 100644 index 000000000..8ebfd33a9 --- /dev/null +++ b/tests/test_moe_wna16_marlin_gemm.py @@ -0,0 +1,365 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +if not hasattr(infini.ops, "MoeWna16MarlinGemm"): + pytest.skip( + "`MoeWna16MarlinGemm` is not available on this platform", + allow_module_level=True, + ) + + +@pytest.mark.parametrize( + "dtype, b_type_id, num_bits, rtol, atol", + ( + (torch.float16, 1125899907892224, 4, 2e-2, 2e-2), + (torch.bfloat16, 1125899907892224, 4, 5e-2, 5e-2), + (torch.float16, 1125899923621888, 8, 2e-2, 2e-2), + ), +) +def test_moe_wna16_marlin_gemm( + dtype, + b_type_id, + num_bits, + rtol, + atol, + device, + implementation_index, +): + if device != "cuda": + pytest.skip("`moe_wna16_marlin_gemm` requires the NVIDIA backend") + + provider_case = _make_case(device, dtype, b_type_id, num_bits) + case = _make_case(device, dtype, b_type_id, num_bits) + expected = provider_case["out"] + provider_result = _call_provider(provider_case, expected) + + assert provider_result.data_ptr() == expected.data_ptr() + result = _call_infini(case, implementation_index, get_stream(case["a"].device)) + + assert result is None + torch.testing.assert_close(case["out"], expected, rtol=rtol, atol=atol) + + +def test_moe_wna16_marlin_gemm_non_default_stream(device, implementation_index): + if device != "cuda": + pytest.skip("non-default CUDA streams require the NVIDIA backend") + + provider_case = _make_case(device, torch.float16, 1125899907892224, 4) + case = _make_case(device, torch.float16, 1125899907892224, 4) + expected = provider_case["out"] + _call_provider(provider_case, expected) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + + _call_infini(case, implementation_index, stream.cuda_stream) + + stream.synchronize() + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def test_moe_wna16_marlin_gemm_optional_act_order(device, implementation_index): + if device != "cuda": + pytest.skip("activation-order metadata requires the NVIDIA backend") + + provider_case = _make_case( + device, torch.float16, 1125899907892224, 4, has_act_order=True + ) + case = _make_case(device, torch.float16, 1125899907892224, 4, has_act_order=True) + expected = provider_case["out"] + _call_provider(provider_case, expected) + + _call_infini(case, implementation_index, get_stream(case["a"].device)) + + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def test_moe_wna16_marlin_gemm_zero_points(device, implementation_index): + if device != "cuda": + pytest.skip("zero-point metadata requires the NVIDIA backend") + + provider_case = _make_case( + device, torch.float16, 1125899906843648, 4, has_zero_points=True + ) + case = _make_case(device, torch.float16, 1125899906843648, 4, has_zero_points=True) + expected = provider_case["out"] + provider_result = _call_provider(provider_case, expected) + + assert provider_result.data_ptr() == expected.data_ptr() + _call_infini(case, implementation_index, get_stream(case["a"].device)) + + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def test_moe_wna16_marlin_gemm_nonzero_bias_and_explicit_tuning( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("bias and tuning require the NVIDIA backend") + + case_kwargs = { + "has_bias": True, + "mul_topk_weights": False, + "size_n": 256, + "size_k": 256, + "thread_config": (128, 128, 1), + "top_k": 1, + "topk_weight_dtype": torch.float16, + "use_atomic_add": False, + "zero_scales": True, + } + provider_case = _make_case( + device, torch.float16, 1125899907892224, 4, **case_kwargs + ) + case = _make_case(device, torch.float16, 1125899907892224, 4, **case_kwargs) + expected = provider_case["out"] + _call_provider(provider_case, expected) + + assert torch.count_nonzero(expected).item() > 0 + _call_infini(case, implementation_index, get_stream(case["a"].device)) + + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def test_moe_wna16_marlin_gemm_int8_activation_scales( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("int8 activation scales require the NVIDIA backend") + + case_kwargs = { + "activation_dtype": torch.int8, + "mul_topk_weights": False, + "size_n": 1024, + "size_k": 1024, + "top_k": 1, + "topk_weight_dtype": torch.float16, + "use_atomic_add": False, + } + provider_case = _make_case( + device, torch.float16, 1125899907892224, 4, **case_kwargs + ) + case = _make_case(device, torch.float16, 1125899907892224, 4, **case_kwargs) + expected = provider_case["out"] + _call_provider(provider_case, expected) + + assert torch.count_nonzero(expected).item() > 0 + _call_infini(case, implementation_index, get_stream(case["a"].device)) + + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def test_moe_wna16_marlin_gemm_accepts_flat_topk_weights( + device, + implementation_index, +): + if device != "cuda": + pytest.skip("top-k weights require the NVIDIA backend") + + provider_case = _make_case(device, torch.float16, 1125899907892224, 4) + case = _make_case(device, torch.float16, 1125899907892224, 4) + provider_case["topk_weights"] = provider_case["topk_weights"].reshape(-1) + case["topk_weights"] = case["topk_weights"].reshape(-1) + expected = provider_case["out"] + _call_provider(provider_case, expected) + + _call_infini(case, implementation_index, get_stream(case["a"].device)) + + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +def _make_case( + device, + dtype, + b_type_id, + num_bits, + *, + activation_dtype=None, + has_act_order=False, + has_bias=False, + has_zero_points=False, + mul_topk_weights=True, + size_n=128, + size_k=256, + thread_config=(-1, -1, -1), + top_k=2, + topk_weight_dtype=torch.float32, + use_atomic_add=None, + zero_scales=False, +): + torch.manual_seed(0) + size_m = 1 + num_experts = 4 + moe_block_size = 16 + route_count = size_m * top_k + padding = route_count + sorted_token_ids = torch.full( + (route_count + num_experts * (moe_block_size - 1),), + padding, + dtype=torch.int32, + device=device, + ) + sorted_token_ids[torch.arange(route_count, device=device) * moe_block_size] = ( + torch.arange(route_count, dtype=torch.int32, device=device) + ) + expert_ids = torch.zeros((num_experts,), dtype=torch.int32, device=device) + expert_ids[:route_count] = torch.arange( + route_count, dtype=torch.int32, device=device + ) + pack_factor = 32 // num_bits + num_groups = 8 if has_act_order else 1 + scale = 0.0 if has_act_order or zero_scales else 0.02 + g_idx_or_none = None + perm_or_none = None + thread_k, thread_n, blocks_per_sm = thread_config + activation_dtype = activation_dtype or dtype + a_scales = None + if activation_dtype == torch.int8: + a = torch.ones((size_m, size_k), dtype=activation_dtype, device=device) + a_scales = torch.full((size_m, 1), 0.01, dtype=torch.float32, device=device) + else: + a = torch.randn((size_m, size_k), dtype=activation_dtype, device=device) + + b_bias_or_none = None + if has_bias: + b_bias_or_none = torch.full( + (num_experts, size_n), 0.125, dtype=dtype, device=device + ) + + if has_act_order: + g_idx_or_none = ( + torch.arange(num_groups, dtype=torch.int32, device=device) + .repeat_interleave(size_k // num_groups) + .repeat(num_experts, 1) + ) + perm_or_none = torch.arange(size_k, dtype=torch.int32, device=device).repeat( + num_experts, 1 + ) + + b_zeros_or_none = None + if has_zero_points: + b_zeros_or_none = torch.zeros( + (num_experts, num_groups, size_n // pack_factor), + dtype=torch.int32, + device=device, + ) + + if use_atomic_add is None: + use_atomic_add = dtype == torch.float16 + + return { + "a": a, + "b_q_weight": torch.zeros( + (num_experts, size_k // 16, size_n * 16 // pack_factor), + dtype=torch.int32, + device=device, + ), + "b_bias_or_none": b_bias_or_none, + "b_scales": torch.full( + (num_experts, num_groups, size_n), scale, dtype=dtype, device=device + ), + "a_scales": a_scales, + "global_scale": None, + "b_zeros_or_none": b_zeros_or_none, + "g_idx_or_none": g_idx_or_none, + "perm_or_none": perm_or_none, + "workspace": torch.zeros(432, dtype=torch.int32, device=device), + "sorted_token_ids": sorted_token_ids, + "expert_ids": expert_ids, + "num_tokens_past_padded": torch.tensor( + (route_count * moe_block_size,), dtype=torch.int32, device=device + ), + "topk_weights": torch.full( + (size_m, top_k), 1.0 / top_k, dtype=topk_weight_dtype, device=device + ), + "moe_block_size": moe_block_size, + "top_k": top_k, + "mul_topk_weights": mul_topk_weights, + "b_type_id": b_type_id, + "size_m": size_m, + "size_n": size_n, + "size_k": size_k, + "is_full_k": True, + "use_atomic_add": use_atomic_add, + "use_fp32_reduce": True, + "is_zp_float": False, + "thread_k": thread_k, + "thread_n": thread_n, + "blocks_per_sm": blocks_per_sm, + "out": torch.zeros((route_count, size_n), dtype=dtype, device=device), + } + + +def _call_provider(case, out): + return torch.ops._moe_C.moe_wna16_marlin_gemm( + case["a"], + out, + case["b_q_weight"], + case["b_bias_or_none"], + case["b_scales"], + case["a_scales"], + case["global_scale"], + case["b_zeros_or_none"], + case["g_idx_or_none"], + case["perm_or_none"], + case["workspace"], + case["sorted_token_ids"], + case["expert_ids"], + case["num_tokens_past_padded"], + case["topk_weights"], + case["moe_block_size"], + case["top_k"], + case["mul_topk_weights"], + case["b_type_id"], + case["size_m"], + case["size_n"], + case["size_k"], + case["is_full_k"], + case["use_atomic_add"], + case["use_fp32_reduce"], + case["is_zp_float"], + case["thread_k"], + case["thread_n"], + case["blocks_per_sm"], + ) + + +def _call_infini(case, implementation_index, stream): + return infini.ops.moe_wna16_marlin_gemm( + case["a"], + case["b_q_weight"], + case["b_bias_or_none"], + case["b_scales"], + case["a_scales"], + case["global_scale"], + case["b_zeros_or_none"], + case["g_idx_or_none"], + case["perm_or_none"], + case["workspace"], + case["sorted_token_ids"], + case["expert_ids"], + case["num_tokens_past_padded"], + case["topk_weights"], + case["moe_block_size"], + case["top_k"], + case["mul_topk_weights"], + case["b_type_id"], + case["size_m"], + case["size_n"], + case["size_k"], + case["is_full_k"], + case["use_atomic_add"], + case["use_fp32_reduce"], + case["is_zp_float"], + case["thread_k"], + case["thread_n"], + case["blocks_per_sm"], + case["out"], + stream=stream, + implementation_index=implementation_index, + )