diff --git a/src/base/fused_marlin_moe.h b/src/base/fused_marlin_moe.h new file mode 100644 index 000000000..ac11d4c4c --- /dev/null +++ b/src/base/fused_marlin_moe.h @@ -0,0 +1,360 @@ +#ifndef INFINI_OPS_BASE_FUSED_MARLIN_MOE_H_ +#define INFINI_OPS_BASE_FUSED_MARLIN_MOE_H_ + +#include +#include +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// Aligned with vLLM `fused_marlin_moe` at commit +// bcc0a3cbefe55f99da4821f9d89106e3d71e4867. +class FusedMarlinMoe : public Operator { + public: + FusedMarlinMoe( + const Tensor hidden_states, const Tensor w1, const Tensor w2, + const Tensor w1_scale, const Tensor w2_scale, const Tensor gating_output, + const Tensor topk_weights, const Tensor topk_ids, + const int64_t quant_type_id, const bool apply_router_weight_on_input, + const int64_t global_num_experts, std::optional expert_map, + std::optional global_scale1, std::optional global_scale2, + std::optional g_idx1, std::optional g_idx2, + std::optional sort_indices1, std::optional sort_indices2, + std::optional w1_zeros, std::optional w2_zeros, + std::optional workspace, const bool is_k_full, const bool inplace, + Tensor out) + : hidden_states_metadata_{hidden_states}, + w1_metadata_{w1}, + w2_metadata_{w2}, + w1_scale_metadata_{w1_scale}, + w2_scale_metadata_{w2_scale}, + gating_output_metadata_{gating_output}, + topk_weights_metadata_{topk_weights}, + topk_ids_metadata_{topk_ids}, + expert_map_metadata_{expert_map}, + global_scale1_metadata_{global_scale1}, + global_scale2_metadata_{global_scale2}, + g_idx1_metadata_{g_idx1}, + g_idx2_metadata_{g_idx2}, + sort_indices1_metadata_{sort_indices1}, + sort_indices2_metadata_{sort_indices2}, + w1_zeros_metadata_{w1_zeros}, + w2_zeros_metadata_{w2_zeros}, + workspace_metadata_{workspace}, + out_metadata_{out}, + quant_type_id_{quant_type_id}, + apply_router_weight_on_input_{apply_router_weight_on_input}, + global_num_experts_{global_num_experts}, + is_k_full_{is_k_full}, + inplace_{inplace}, + device_index_{hidden_states.device().index()} { + Validate(hidden_states, w1, w2, w1_scale, w2_scale, gating_output, + topk_weights, topk_ids, expert_map, global_scale1, global_scale2, + g_idx1, g_idx2, sort_indices1, sort_indices2, w1_zeros, w2_zeros, + workspace, out); + } + + virtual void operator()( + const Tensor hidden_states, const Tensor w1, const Tensor w2, + const Tensor w1_scale, const Tensor w2_scale, const Tensor gating_output, + const Tensor topk_weights, const Tensor topk_ids, + const int64_t quant_type_id, const bool apply_router_weight_on_input, + const int64_t global_num_experts, std::optional expert_map, + std::optional global_scale1, std::optional global_scale2, + std::optional g_idx1, std::optional g_idx2, + std::optional sort_indices1, std::optional sort_indices2, + std::optional w1_zeros, std::optional w2_zeros, + std::optional workspace, const bool is_k_full, const bool inplace, + Tensor out) const = 0; + + protected: + void ValidateCallMetadata( + const Tensor hidden_states, const Tensor w1, const Tensor w2, + const Tensor w1_scale, const Tensor w2_scale, const Tensor gating_output, + const Tensor topk_weights, const Tensor topk_ids, + const int64_t quant_type_id, const bool apply_router_weight_on_input, + const int64_t global_num_experts, std::optional expert_map, + std::optional global_scale1, std::optional global_scale2, + std::optional g_idx1, std::optional g_idx2, + std::optional sort_indices1, std::optional sort_indices2, + std::optional w1_zeros, std::optional w2_zeros, + std::optional workspace, const bool is_k_full, const bool inplace, + Tensor out) const { + assert(quant_type_id == quant_type_id_ && + apply_router_weight_on_input == apply_router_weight_on_input_ && + global_num_experts == global_num_experts_ && + is_k_full == is_k_full_ && inplace == inplace_ && + "`FusedMarlinMoe` 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(hidden_states_metadata_, hidden_states) && + same_metadata(w1_metadata_, w1) && same_metadata(w2_metadata_, w2) && + same_metadata(w1_scale_metadata_, w1_scale) && + same_metadata(w2_scale_metadata_, w2_scale) && + same_metadata(gating_output_metadata_, gating_output) && + same_metadata(topk_weights_metadata_, topk_weights) && + same_metadata(topk_ids_metadata_, topk_ids) && + optional_matches(expert_map_metadata_, expert_map) && + optional_matches(global_scale1_metadata_, global_scale1) && + optional_matches(global_scale2_metadata_, global_scale2) && + optional_matches(g_idx1_metadata_, g_idx1) && + optional_matches(g_idx2_metadata_, g_idx2) && + optional_matches(sort_indices1_metadata_, sort_indices1) && + optional_matches(sort_indices2_metadata_, sort_indices2) && + optional_matches(w1_zeros_metadata_, w1_zeros) && + optional_matches(w2_zeros_metadata_, w2_zeros) && + optional_matches(workspace_metadata_, workspace) && + same_metadata(out_metadata_, out); + assert(matches && "`FusedMarlinMoe` call metadata must match descriptor"); + + const auto aliases = hidden_states.data() == out.data(); + assert(aliases == inplace && + "`FusedMarlinMoe` `out` alias does not match `inplace`"); + } + + private: + void Validate(const Tensor hidden_states, const Tensor w1, const Tensor w2, + const Tensor w1_scale, const Tensor w2_scale, + const Tensor gating_output, const Tensor topk_weights, + const Tensor topk_ids, std::optional expert_map, + std::optional global_scale1, + std::optional global_scale2, + std::optional g_idx1, std::optional g_idx2, + std::optional sort_indices1, + std::optional sort_indices2, + std::optional w1_zeros, std::optional w2_zeros, + std::optional workspace, const Tensor out) const { + assert(hidden_states.ndim() == 2 && hidden_states.size(0) > 0 && + hidden_states.size(1) > 0 && hidden_states.IsContiguous() && + (hidden_states.dtype() == DataType::kFloat16 || + hidden_states.dtype() == DataType::kBFloat16) && + "`FusedMarlinMoe` requires non-empty contiguous float16 or " + "bfloat16 `hidden_states`"); + + constexpr int64_t kUint4 = 1125899906843648; + constexpr int64_t kUint4B8 = 1125899907892224; + [[maybe_unused]] constexpr int64_t kUint8B128 = 1125899923621888; + [[maybe_unused]] constexpr int64_t kFloat8E4M3Fn = 2814749767172868; + constexpr int64_t kFloat4E2M1F = 562949953487106; + const auto num_bits = quant_type_id_ == kUint4 || + quant_type_id_ == kUint4B8 || + quant_type_id_ == kFloat4E2M1F + ? 4 + : 8; + assert((quant_type_id_ == kUint4 || quant_type_id_ == kUint4B8 || + quant_type_id_ == kUint8B128 || quant_type_id_ == kFloat8E4M3Fn || + quant_type_id_ == kFloat4E2M1F) && + "`FusedMarlinMoe` received an unsupported `quant_type_id`"); + + const auto num_tokens = hidden_states.size(0); + const auto hidden_size = hidden_states.size(1); + const auto num_experts = w1.ndim() == 3 ? w1.size(0) : 0; + const auto topk = topk_ids.ndim() == 2 ? topk_ids.size(1) : 0; + const auto intermediate_size = w2.ndim() == 3 ? w2.size(1) * 16 : 0; + [[maybe_unused]] const auto pack_factor = 32 / num_bits; + + assert(num_experts > 0 && num_experts < 1024 && topk > 0 && + hidden_size % 64 == 0 && intermediate_size > 0 && + intermediate_size % 32 == 0 && + intermediate_size <= std::numeric_limits::max() / 32 && + "`FusedMarlinMoe` received unsupported dimensions"); + assert(w1.ndim() == 3 && w2.ndim() == 3 && w1.size(1) * 16 == hidden_size && + w1.size(2) == intermediate_size * 2 * 16 / pack_factor && + w2.size(0) == num_experts && + w2.size(2) == hidden_size * 16 / pack_factor && + w1.dtype() == DataType::kInt32 && w2.dtype() == DataType::kInt32 && + w1.IsContiguous() && w2.IsContiguous() && + "`FusedMarlinMoe` received invalid packed weights"); + + assert(w1_scale.ndim() == 3 && w2_scale.ndim() == 3 && + w1_scale.size(0) == num_experts && w1_scale.size(1) > 0 && + w1_scale.size(2) == intermediate_size * 2 && + hidden_size % w1_scale.size(1) == 0 && + w2_scale.size(0) == num_experts && w2_scale.size(1) > 0 && + w2_scale.size(2) == hidden_size && + intermediate_size % w2_scale.size(1) == 0 && + w1_scale.dtype() == hidden_states.dtype() && + w2_scale.dtype() == hidden_states.dtype() && + w1_scale.IsContiguous() && w2_scale.IsContiguous() && + "`FusedMarlinMoe` received invalid weight scales"); + + assert(gating_output.ndim() > 0 && gating_output.size(0) == num_tokens && + "`FusedMarlinMoe` `gating_output` token count must match " + "`hidden_states`"); + assert(topk_ids.shape() == topk_weights.shape() && + topk_ids.size(0) == num_tokens && + topk_ids.dtype() == DataType::kInt32 && + topk_weights.dtype() == DataType::kFloat32 && + topk_ids.IsContiguous() && topk_weights.IsContiguous() && + "`FusedMarlinMoe` received invalid top-k routing tensors"); + + [[maybe_unused]] const auto routed_num_experts = + global_num_experts_ == -1 ? num_experts : global_num_experts_; + assert(routed_num_experts > 0 && routed_num_experts < 1024 && + "`FusedMarlinMoe` requires `global_num_experts` in `[1, 1023]` " + "or `-1`"); + assert((expert_map || routed_num_experts == num_experts) && + "`FusedMarlinMoe` requires `global_num_experts` to match local " + "experts when `expert_map` is absent"); + if (expert_map) { + assert(expert_map->ndim() == 1 && + expert_map->numel() == routed_num_experts && + expert_map->dtype() == DataType::kInt32 && + expert_map->IsContiguous() && + "`FusedMarlinMoe` received an invalid `expert_map`"); + } + + [[maybe_unused]] const auto optional_pair = + [](const std::optional& first, + const std::optional& second) { + return first.has_value() == second.has_value(); + }; + assert(optional_pair(g_idx1, sort_indices1) && + optional_pair(g_idx2, sort_indices2) && + "`FusedMarlinMoe` requires each `g_idx` with its " + "`sort_indices`"); + + [[maybe_unused]] const auto has_zero_points = quant_type_id_ == kUint4; + const auto validate_layer_metadata = + [&](const std::optional& zeros, + const std::optional& g_idx, + const std::optional& sort_indices, const Tensor scales, + const Tensor::Size size_k, const Tensor::Size size_n) { + assert(zeros.has_value() == has_zero_points && + "`FusedMarlinMoe` zero points do not match the " + "quantization type"); + if (zeros) { + assert(zeros->ndim() == 3 && zeros->size(0) == num_experts && + zeros->size(1) == scales.size(1) && + zeros->size(2) == size_n / pack_factor && + zeros->dtype() == DataType::kInt32 && + "`FusedMarlinMoe` received invalid zero points"); + } + if (g_idx && sort_indices) { + const auto both_empty = + g_idx->numel() == 0 && sort_indices->numel() == 0; + const auto valid_nonempty = + g_idx->ndim() == 2 && g_idx->size(0) == num_experts && + g_idx->size(1) == size_k && + sort_indices->shape() == g_idx->shape() && + (!is_k_full_ || scales.size(1) > 1); + assert(g_idx->dtype() == DataType::kInt32 && + sort_indices->dtype() == DataType::kInt32 && + (both_empty || valid_nonempty) && + "`FusedMarlinMoe` received invalid activation-order " + "metadata"); + } + }; + validate_layer_metadata(w1_zeros, g_idx1, sort_indices1, w1_scale, + hidden_size, intermediate_size * 2); + validate_layer_metadata(w2_zeros, g_idx2, sort_indices2, w2_scale, + intermediate_size, hidden_size); + + const auto same_device = [&](const Tensor tensor) { + return tensor.device().type() == hidden_states.device().type() && + tensor.device().index() == hidden_states.device().index(); + }; + [[maybe_unused]] const auto valid_optional = + [&](const std::optional& tensor) { + return !tensor || (same_device(*tensor) && tensor->IsContiguous()); + }; + assert(same_device(w1) && same_device(w2) && same_device(w1_scale) && + same_device(w2_scale) && same_device(gating_output) && + same_device(topk_weights) && same_device(topk_ids) && + valid_optional(expert_map) && valid_optional(global_scale1) && + valid_optional(global_scale2) && valid_optional(g_idx1) && + valid_optional(g_idx2) && valid_optional(sort_indices1) && + valid_optional(sort_indices2) && valid_optional(w1_zeros) && + valid_optional(w2_zeros) && valid_optional(workspace) && + same_device(out) && + "`FusedMarlinMoe` requires all tensors on the input device"); + if (workspace) { + assert(workspace->ndim() == 1 && workspace->numel() > 0 && + workspace->dtype() == DataType::kInt32 && + "`FusedMarlinMoe` requires a non-empty int32 `workspace`"); + } + + assert(out.shape() == hidden_states.shape() && + out.dtype() == hidden_states.dtype() && out.IsContiguous() && + "`FusedMarlinMoe` output metadata must match `hidden_states`"); + const auto aliases = hidden_states.data() == out.data(); + assert(aliases == inplace_ && + "`FusedMarlinMoe` `out` must alias `hidden_states` exactly when " + "`inplace` is true"); + + assert(num_tokens <= std::numeric_limits::max() / topk && + hidden_size <= std::numeric_limits::max() / 16 && + "`FusedMarlinMoe` dimensions overflow"); + [[maybe_unused]] const auto route_count = num_tokens * topk; + assert(route_count <= std::numeric_limits::max() && + routed_num_experts <= + (std::numeric_limits::max() - route_count) / 63 && + "`FusedMarlinMoe` routing dimensions overflow int32 indices"); + } + + Tensor hidden_states_metadata_; + + Tensor w1_metadata_; + + Tensor w2_metadata_; + + Tensor w1_scale_metadata_; + + Tensor w2_scale_metadata_; + + Tensor gating_output_metadata_; + + Tensor topk_weights_metadata_; + + Tensor topk_ids_metadata_; + + std::optional expert_map_metadata_; + + std::optional global_scale1_metadata_; + + std::optional global_scale2_metadata_; + + std::optional g_idx1_metadata_; + + std::optional g_idx2_metadata_; + + std::optional sort_indices1_metadata_; + + std::optional sort_indices2_metadata_; + + std::optional w1_zeros_metadata_; + + std::optional w2_zeros_metadata_; + + std::optional workspace_metadata_; + + Tensor out_metadata_; + + int64_t quant_type_id_{0}; + + bool apply_router_weight_on_input_{false}; + + int64_t global_num_experts_{-1}; + + bool is_k_full_{true}; + + bool inplace_{false}; + + protected: + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_FUSED_MARLIN_MOE_H_ diff --git a/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.cc b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.cc new file mode 100644 index 000000000..ec28297ca --- /dev/null +++ b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.cc @@ -0,0 +1,105 @@ +#include "linked/torch/nvidia/ops/fused_marlin_moe/vllm.h" + +#include +#include +#include + +#include + +namespace infini::ops::linked::torch::nvidia { + +void VllmFusedMarlinMoe::Validate(const int64_t quant_type_id, + const bool has_global_scale, + const bool has_w1_zeros, + const bool has_w2_zeros) { + VllmMoeWna16MarlinGemm::Validate(quant_type_id, has_global_scale); + + constexpr int64_t kUint4 = 1125899906843648; + constexpr int64_t kUint4B8 = 1125899907892224; + constexpr int64_t kUint8B128 = 1125899923621888; + constexpr int64_t kFloat8E4M3Fn = 2814749767172868; + TORCH_CHECK(quant_type_id == kUint4 || quant_type_id == kUint4B8 || + quant_type_id == kUint8B128 || quant_type_id == kFloat8E4M3Fn, + "Linked `fused_marlin_moe` received an unsupported " + "`quant_type_id`."); + + const auto expects_zeros = quant_type_id == kUint4; + TORCH_CHECK(has_w1_zeros == expects_zeros && has_w2_zeros == expects_zeros, + "Linked `fused_marlin_moe` requires zero points for both weight " + "tensors exactly when `quant_type_id` is `uint4`."); +} + +int64_t VllmFusedMarlinMoe::WorkspaceSize(const int device_index) { + cudaDeviceProp properties{}; + const auto status = cudaGetDeviceProperties(&properties, device_index); + TORCH_CHECK(status == cudaSuccess, + "Failed to query the CUDA device for linked " + "`fused_marlin_moe`."); + + return static_cast(properties.multiProcessorCount) * 4; +} + +bool VllmFusedMarlinMoe::UseAtomicAdd(const at::ScalarType dtype, + const int device_index) { + cudaDeviceProp properties{}; + const auto status = cudaGetDeviceProperties(&properties, device_index); + TORCH_CHECK(status == cudaSuccess, + "Failed to query the CUDA device for linked " + "`fused_marlin_moe`."); + + return dtype == at::kHalf || properties.major >= 9; +} + +void VllmFusedMarlinMoe::CallAlign(at::Tensor topk_ids, + const int64_t num_experts, + const int64_t block_size, + at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded) { + static const auto op = c10::Dispatcher::singleton().findSchemaOrThrow( + "_moe_C::moe_align_block_size", ""); + c10::Stack stack; + stack.reserve(6); + stack.emplace_back(std::move(topk_ids)); + stack.emplace_back(num_experts); + stack.emplace_back(block_size); + stack.emplace_back(std::move(sorted_token_ids)); + stack.emplace_back(std::move(expert_ids)); + stack.emplace_back(std::move(num_tokens_post_padded)); + op.callBoxed(&stack); + + TORCH_CHECK(stack.empty(), + "Linked `moe_align_block_size` returned an unexpected value."); +} + +void VllmFusedMarlinMoe::CallMarlin( + at::Tensor a, at::Tensor out, at::Tensor b_q_weight, at::Tensor b_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, + const int64_t moe_block_size, const int64_t top_k, + const bool mul_topk_weights, const bool is_ep, const int64_t b_q_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) { + VllmMoeWna16MarlinGemm::Call( + std::move(a), std::move(out), std::move(b_q_weight), std::move(b_scales), + std::move(global_scale), std::move(b_zeros_or_none), + std::move(g_idx_or_none), std::move(perm_or_none), std::move(workspace), + std::move(sorted_token_ids), std::move(expert_ids), + std::move(num_tokens_past_padded), std::move(topk_weights), + moe_block_size, top_k, mul_topk_weights, is_ep, b_q_type_id, size_m, + size_n, size_k, is_full_k, use_atomic_add, use_fp32_reduce, is_zp_float); +} + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +template class TorchFusedMarlinMoe< + ::infini::ops::linked::torch::nvidia::VllmFusedMarlinMoe>; + +} // namespace infini::ops::linked::torch diff --git a/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.h b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.h new file mode 100644 index 000000000..d8d7b0361 --- /dev/null +++ b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.h @@ -0,0 +1,65 @@ +#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FUSED_MARLIN_MOE_VLLM_H_ +#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FUSED_MARLIN_MOE_VLLM_H_ + +#include + +#include "linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h" +#include "linked/torch/ops/fused_marlin_moe.h" + +namespace infini::ops::linked::torch::nvidia { + +struct VllmFusedMarlinMoe : C10 { + static void Validate(int64_t quant_type_id, bool has_global_scale, + bool has_w1_zeros, bool has_w2_zeros); + + static int64_t WorkspaceSize(int device_index); + + static bool UseAtomicAdd(at::ScalarType dtype, int device_index); + + static void CallAlign(at::Tensor topk_ids, int64_t num_experts, + int64_t block_size, at::Tensor sorted_token_ids, + at::Tensor expert_ids, + at::Tensor num_tokens_post_padded); + + static void CallMarlin(at::Tensor a, at::Tensor out, at::Tensor b_q_weight, + at::Tensor b_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, bool is_ep, + int64_t b_q_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); +}; + +} // namespace infini::ops::linked::torch::nvidia + +namespace infini::ops::linked::torch { + +extern template class TorchFusedMarlinMoe< + ::infini::ops::linked::torch::nvidia::VllmFusedMarlinMoe>; + +} // namespace infini::ops::linked::torch + +namespace infini::ops { + +template <> +class Operator + : public linked::torch::TorchFusedMarlinMoe< + linked::torch::nvidia::VllmFusedMarlinMoe> { + public: + using linked::torch::TorchFusedMarlinMoe< + linked::torch::nvidia::VllmFusedMarlinMoe>::TorchFusedMarlinMoe; + + using linked::torch::TorchFusedMarlinMoe< + linked::torch::nvidia::VllmFusedMarlinMoe>::operator(); +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FUSED_MARLIN_MOE_VLLM_H_ diff --git a/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.yaml b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.yaml new file mode 100644 index 000000000..eac59e53a --- /dev/null +++ b/src/linked/torch/nvidia/ops/fused_marlin_moe/vllm.yaml @@ -0,0 +1,16 @@ +library: vllm_moe +operator_schema: + - >- + _moe_C::moe_align_block_size(Tensor topk_ids, int num_experts, + int block_size, Tensor! sorted_token_ids, Tensor! experts_ids, + Tensor! num_tokens_post_pad) -> () + - >- + _moe_C::moe_wna16_marlin_gemm(Tensor! a, Tensor? c_or_none, + Tensor! b_q_weight, Tensor! b_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, bool is_ep, int b_q_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) -> Tensor +dispatch_key: CUDA diff --git a/src/linked/torch/ops/fused_marlin_moe.h b/src/linked/torch/ops/fused_marlin_moe.h new file mode 100644 index 000000000..8fcbdcf22 --- /dev/null +++ b/src/linked/torch/ops/fused_marlin_moe.h @@ -0,0 +1,196 @@ +#ifndef INFINI_OPS_LINKED_TORCH_OPS_FUSED_MARLIN_MOE_H_ +#define INFINI_OPS_LINKED_TORCH_OPS_FUSED_MARLIN_MOE_H_ + +#include +#include +#include + +#include +#include +#include +#include + +#include "base/fused_marlin_moe.h" +#include "torch/tensor_.h" + +namespace infini::ops::linked::torch { + +template +class TorchFusedMarlinMoe : public ::infini::ops::FusedMarlinMoe { + public: + using ::infini::ops::FusedMarlinMoe::FusedMarlinMoe; + + using ::infini::ops::FusedMarlinMoe::operator(); + + void operator()( + const Tensor hidden_states, const Tensor w1, const Tensor w2, + const Tensor w1_scale, const Tensor w2_scale, const Tensor gating_output, + const Tensor topk_weights, const Tensor topk_ids, + const int64_t quant_type_id, const bool apply_router_weight_on_input, + const int64_t global_num_experts, std::optional expert_map, + std::optional global_scale1, std::optional global_scale2, + std::optional g_idx1, std::optional g_idx2, + std::optional sort_indices1, std::optional sort_indices2, + std::optional w1_zeros, std::optional w2_zeros, + std::optional workspace, const bool is_k_full, const bool inplace, + Tensor out) const override { + ValidateCallMetadata(hidden_states, w1, w2, w1_scale, w2_scale, + gating_output, topk_weights, topk_ids, quant_type_id, + apply_router_weight_on_input, global_num_experts, + expert_map, global_scale1, global_scale2, g_idx1, + g_idx2, sort_indices1, sort_indices2, w1_zeros, + w2_zeros, workspace, is_k_full, inplace, out); + + Backend::Validate(quant_type_id, + global_scale1.has_value() || global_scale2.has_value(), + w1_zeros.has_value(), w2_zeros.has_value()); + const auto aliases = hidden_states.data() == out.data(); + TORCH_CHECK(g_idx1.has_value() == sort_indices1.has_value() && + g_idx2.has_value() == sort_indices2.has_value(), + "Linked `fused_marlin_moe` requires each `g_idx` with its " + "`sort_indices`."); + TORCH_CHECK(aliases == inplace, + "Linked `fused_marlin_moe` `out` must alias `hidden_states` " + "exactly when `inplace` is true."); + TORCH_CHECK(expert_map.has_value() || global_num_experts == -1 || + global_num_experts == static_cast(w1.size(0)), + "Linked `fused_marlin_moe` requires `global_num_experts` to " + "match local experts when `expert_map` is absent."); + + const typename Backend::StreamGuard stream_guard{ + Backend::GetStreamFromExternal(stream_, device_index_)}; + + auto at_hidden_states = ToAten(hidden_states); + auto at_w1 = ToAten(w1); + auto at_w2 = ToAten(w2); + auto at_w1_scale = ToAten(w1_scale); + auto at_w2_scale = ToAten(w2_scale); + [[maybe_unused]] auto at_gating_output = ToAten(gating_output); + auto at_topk_weights = ToAten(topk_weights); + auto at_topk_ids = ToAten(topk_ids); + auto at_expert_map = ToOptionalAten(expert_map); + auto at_global_scale1 = ToOptionalAten(global_scale1); + auto at_global_scale2 = ToOptionalAten(global_scale2); + auto at_g_idx1 = ToOptionalAten(g_idx1); + auto at_g_idx2 = ToOptionalAten(g_idx2); + auto at_sort_indices1 = ToOptionalAten(sort_indices1); + auto at_sort_indices2 = ToOptionalAten(sort_indices2); + auto at_w1_zeros = ToOptionalAten(w1_zeros); + auto at_w2_zeros = ToOptionalAten(w2_zeros); + auto at_workspace_input = ToOptionalAten(workspace); + auto at_out = ToAten(out); + + const auto num_tokens = at_hidden_states.size(0); + const auto hidden_size = at_hidden_states.size(1); + const auto num_experts = at_w1.size(0); + const auto topk = at_topk_ids.size(1); + const auto intermediate_size = at_w2.size(1) * 16; + const auto route_count = num_tokens * topk; + const auto block_size = SelectBlockSize(num_tokens, topk, num_experts); + const auto routed_num_experts = + global_num_experts == -1 ? num_experts : global_num_experts; + + const auto max_num_tokens_padded = + route_count + routed_num_experts * (block_size - 1); + auto sorted_token_ids = at::empty({max_num_tokens_padded}, + at_topk_ids.options().dtype(at::kInt)); + auto expert_ids = + at::empty({(max_num_tokens_padded + block_size - 1) / block_size}, + at_topk_ids.options().dtype(at::kInt)); + auto num_tokens_post_padded = + at::empty({1}, at_topk_ids.options().dtype(at::kInt)); + Backend::CallAlign(at_topk_ids, routed_num_experts, block_size, + sorted_token_ids, expert_ids, num_tokens_post_padded); + if (at_expert_map) { + expert_ids = at_expert_map->index({expert_ids.to(at::kLong)}); + } + + at::Tensor at_workspace; + if (at_workspace_input) { + at_workspace = *at_workspace_input; + TORCH_CHECK( + at_workspace.numel() >= Backend::WorkspaceSize(device_index_), + "Linked `fused_marlin_moe` requires an int32 workspace with at " + "least four entries per streaming multiprocessor."); + } else { + at_workspace = at::zeros({Backend::WorkspaceSize(device_index_)}, + at_hidden_states.options().dtype(at::kInt)); + } + + const auto cache13_size = + route_count * std::max(intermediate_size * 2, hidden_size); + auto cache13 = at::empty({cache13_size}, at_hidden_states.options()); + auto cache1 = cache13.narrow(0, 0, route_count * intermediate_size * 2) + .view({route_count, intermediate_size * 2}); + auto cache2 = + at::empty({route_count, intermediate_size}, at_hidden_states.options()); + auto cache3 = cache13.narrow(0, 0, route_count * hidden_size) + .view({route_count, hidden_size}); + + const auto is_ep = at_expert_map.has_value(); + const auto use_atomic_add = + Backend::UseAtomicAdd(at_hidden_states.scalar_type(), device_index_); + Backend::CallMarlin(at_hidden_states, cache1, at_w1, at_w1_scale, + at_global_scale1, at_w1_zeros, at_g_idx1, + at_sort_indices1, at_workspace, sorted_token_ids, + expert_ids, num_tokens_post_padded, at_topk_weights, + block_size, topk, apply_router_weight_on_input, is_ep, + quant_type_id, num_tokens, intermediate_size * 2, + hidden_size, is_k_full, use_atomic_add, true, false); + + const auto cache1_left = cache1.narrow(1, 0, intermediate_size); + const auto cache1_right = + cache1.narrow(1, intermediate_size, intermediate_size); + at::silu_out(cache2, cache1_left); + cache2.mul_(cache1_right); + + if (is_ep) { + cache3.zero_(); + } + Backend::CallMarlin( + cache2, cache3, at_w2, at_w2_scale, at_global_scale2, at_w2_zeros, + at_g_idx2, at_sort_indices2, at_workspace, sorted_token_ids, expert_ids, + num_tokens_post_padded, at_topk_weights, block_size, 1, + !apply_router_weight_on_input, is_ep, quant_type_id, route_count, + hidden_size, intermediate_size, is_k_full, use_atomic_add, true, false); + + const auto routed_output = cache3.view({num_tokens, topk, hidden_size}); + at::sum_out(at_out, routed_output, {1}, false, std::nullopt); + } + + private: + static int64_t SelectBlockSize(const int64_t num_tokens, const int64_t topk, + const int64_t num_experts) { + int64_t block_size = 64; + for (const auto candidate : std::array{8, 16, 32, 48, 64}) { + block_size = candidate; + if (static_cast(num_tokens * topk) / + static_cast(num_experts) / + static_cast(block_size) < + 0.9) { + break; + } + } + + return block_size; + } + + 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_FUSED_MARLIN_MOE_H_ diff --git a/tests/test_fused_marlin_moe.py b/tests/test_fused_marlin_moe.py new file mode 100644 index 000000000..87839e5c8 --- /dev/null +++ b/tests/test_fused_marlin_moe.py @@ -0,0 +1,257 @@ +import importlib + +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +if not hasattr(infini.ops, "FusedMarlinMoe"): + pytest.skip( + "`FusedMarlinMoe` is not available on this platform", + allow_module_level=True, + ) +if 16 not in infini.ops.FusedMarlinMoe.active_implementation_indices("nvidia"): + pytest.skip( + "the linked vLLM `FusedMarlinMoe` implementation is unavailable", + allow_module_level=True, + ) + + +try: + importlib.import_module("vllm.model_executor.layers.fused_moe.fused_marlin_moe") +except (ImportError, OSError, RuntimeError) as error: + pytest.skip( + f"vLLM `fused_marlin_moe` reference is unavailable: {error}", + allow_module_level=True, + ) + + +@pytest.mark.parametrize( + "inplace, has_expert_map, has_empty_act_order", + ( + (False, False, False), + (True, False, False), + (False, True, True), + ), +) +@pytest.mark.parametrize("device, implementation_index", (("cuda", 16),)) +def test_fused_marlin_moe( + inplace, + has_expert_map, + has_empty_act_order, + device, + implementation_index, +): + provider_case = _make_case(device, inplace, has_expert_map, has_empty_act_order) + case = _make_case(device, inplace, has_expert_map, has_empty_act_order) + expected = _call_provider(provider_case) + + assert (expected.data_ptr() == provider_case["hidden_states"].data_ptr()) is inplace + result = _call_infini( + case, + implementation_index, + get_stream(case["hidden_states"].device), + ) + + assert result is None + torch.testing.assert_close(case["out"], expected, rtol=2e-2, atol=2e-2) + + +@pytest.mark.parametrize("device, implementation_index", (("cuda", 16),)) +def test_fused_marlin_moe_non_default_stream(device, implementation_index): + provider_case = _make_case(device, False, False) + case = _make_case(device, False, False) + expected = _call_provider(provider_case) + 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) + + +@pytest.mark.parametrize( + "unsupported, message", + ( + ("float4", "float4_e2m1f"), + ("global_scale", "global_scale"), + ("short_workspace", "workspace"), + ), +) +@pytest.mark.parametrize("device, implementation_index", (("cuda", 16),)) +def test_fused_marlin_moe_unsupported( + unsupported, + message, + device, + implementation_index, +): + case = _make_case(device, False, False) + if unsupported == "float4": + case["quant_type_id"] = 562949953487106 + elif unsupported == "global_scale": + case["global_scale1"] = torch.ones(1, device=device) + else: + case["workspace"] = torch.zeros(1, dtype=torch.int32, device=device) + + with pytest.raises(RuntimeError, match=message): + _call_infini( + case, + implementation_index, + get_stream(case["hidden_states"].device), + ) + + +def _make_case(device, inplace, has_expert_map, has_empty_act_order=False): + torch.manual_seed(0) + num_tokens, hidden_size = 1, 256 + intermediate_size = 128 + num_experts = 4 + pack_factor = 8 + hidden_states = torch.randn( + (num_tokens, hidden_size), dtype=torch.float16, device=device + ) + + def make_packed_weight(shape): + return torch.randint( + -(2**31), + 2**31 - 1, + shape, + dtype=torch.int32, + device=device, + ) + + w1 = make_packed_weight( + ( + num_experts, + hidden_size // 16, + intermediate_size * 2 * 16 // pack_factor, + ) + ) + w2 = make_packed_weight( + ( + num_experts, + intermediate_size // 16, + hidden_size * 16 // pack_factor, + ) + ) + w1_scale = ( + torch.rand( + (num_experts, 1, intermediate_size * 2), + dtype=torch.float32, + device=device, + ) + * 0.02 + ).to(hidden_states.dtype) + w2_scale = ( + torch.rand( + (num_experts, 1, hidden_size), + dtype=torch.float32, + device=device, + ) + * 0.02 + ).to(hidden_states.dtype) + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device=device + ) + topk_weights = torch.tensor(((0.75, 0.25),), dtype=torch.float32, device=device) + topk_ids = torch.tensor(((0, 1),), dtype=torch.int32, device=device) + expert_map = ( + torch.arange(num_experts, dtype=torch.int32, device=device) + if has_expert_map + else None + ) + act_order = ( + tuple(torch.empty(0, dtype=torch.int32, device=device) for _ in range(4)) + if has_empty_act_order + else (None,) * 4 + ) + out = hidden_states if inplace else torch.empty_like(hidden_states) + + return { + "hidden_states": hidden_states, + "w1": w1, + "w2": w2, + "w1_scale": w1_scale, + "w2_scale": w2_scale, + "gating_output": gating_output, + "topk_weights": topk_weights, + "topk_ids": topk_ids, + "quant_type_id": 1125899907892224, + "apply_router_weight_on_input": False, + "global_num_experts": num_experts if has_expert_map else -1, + "expert_map": expert_map, + "global_scale1": None, + "global_scale2": None, + "g_idx1": act_order[0], + "g_idx2": act_order[1], + "sort_indices1": act_order[2], + "sort_indices2": act_order[3], + "w1_zeros": None, + "w2_zeros": None, + "workspace": None, + "is_k_full": True, + "inplace": inplace, + "out": out, + } + + +def _call_provider(case): + return torch.ops.vllm.fused_marlin_moe.default( + case["hidden_states"], + case["w1"], + case["w2"], + case["w1_scale"], + case["w2_scale"], + case["gating_output"], + case["topk_weights"], + case["topk_ids"], + case["quant_type_id"], + case["apply_router_weight_on_input"], + case["global_num_experts"], + case["expert_map"], + case["global_scale1"], + case["global_scale2"], + case["g_idx1"], + case["g_idx2"], + case["sort_indices1"], + case["sort_indices2"], + case["w1_zeros"], + case["w2_zeros"], + case["workspace"], + case["is_k_full"], + case["inplace"], + ) + + +def _call_infini(case, implementation_index, stream): + return infini.ops.fused_marlin_moe( + case["hidden_states"], + case["w1"], + case["w2"], + case["w1_scale"], + case["w2_scale"], + case["gating_output"], + case["topk_weights"], + case["topk_ids"], + case["quant_type_id"], + case["apply_router_weight_on_input"], + case["global_num_experts"], + case["expert_map"], + case["global_scale1"], + case["global_scale2"], + case["g_idx1"], + case["g_idx2"], + case["sort_indices1"], + case["sort_indices2"], + case["w1_zeros"], + case["w2_zeros"], + case["workspace"], + case["is_k_full"], + case["inplace"], + case["out"], + stream=stream, + implementation_index=implementation_index, + )