diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 723dc2d4f..d84ddf7d0 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -472,7 +472,7 @@ set(_infini_ops_smoke_ops set(_infini_ops_smoke_torch_ops abs clamp exp) if(WITH_NVIDIA) - list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm) + list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum) endif() if(WITH_NVIDIA AND WITH_TORCH) diff --git a/src/base/moe_sum.h b/src/base/moe_sum.h new file mode 100644 index 000000000..d5b5259de --- /dev/null +++ b/src/base/moe_sum.h @@ -0,0 +1,213 @@ +#ifndef INFINI_OPS_BASE_MOE_SUM_H_ +#define INFINI_OPS_BASE_MOE_SUM_H_ + +#include +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// Aligned with vLLM `_moe_C::moe_sum`. +class MoeSum : public Operator { + public: + MoeSum(const Tensor input, Tensor output) + : MoeSum{input, std::nullopt, std::nullopt, output} {} + + MoeSum(const Tensor input, std::optional topk_ids, + std::optional expert_map, Tensor output) + : num_tokens_{input.ndim() == 3 ? input.size(0) : 0}, + topk_{input.ndim() == 3 ? input.size(1) : 0}, + hidden_size_{input.ndim() == 3 ? input.size(2) : 0}, + input_strides_{input.strides()}, + output_strides_{output.strides()}, + dtype_{input.dtype()}, + device_type_{input.device().type()}, + has_topk_ids_{topk_ids.has_value()}, + topk_ids_dtype_{topk_ids ? topk_ids->dtype() : DataType::kInt32}, + topk_ids_token_stride_{ + topk_ids && topk_ids->ndim() == 2 ? topk_ids->stride(0) : 0}, + topk_ids_slot_stride_{ + topk_ids && topk_ids->ndim() == 2 ? topk_ids->stride(1) : 0}, + has_expert_map_{expert_map.has_value()}, + expert_map_size_{expert_map ? expert_map->numel() : 0}, + expert_map_stride_{ + expert_map && expert_map->ndim() == 1 ? expert_map->stride(0) : 0}, + device_index_{input.device().index()} { + assert(input.ndim() == 3 && output.ndim() == 2 && + "`MoeSum` requires `[num_tokens, topk, hidden_size]` input and " + "`[num_tokens, hidden_size]` output"); + assert(output.size(0) == num_tokens_ && output.size(1) == hidden_size_ && + "`MoeSum` output shape is incompatible with the input"); + assert(topk_ > 0 && "`MoeSum` requires at least one top-k slot"); + assert((dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat16 || + dtype_ == DataType::kBFloat16) && + "`MoeSum` supports float32, float16, and bfloat16 inputs"); + assert(output.dtype() == dtype_ && + "`MoeSum` input and output dtypes must match"); + assert(output.IsContiguous() && "`MoeSum` requires contiguous output"); + + constexpr auto kMaxSignedIndex = + static_cast(std::numeric_limits::max()); + assert(num_tokens_ <= kMaxSignedIndex && topk_ <= kMaxSignedIndex && + hidden_size_ <= kMaxSignedIndex && + "`MoeSum` dimensions must fit signed index arithmetic"); + assert(num_tokens_ <= std::numeric_limits::max() && + "`MoeSum` token count exceeds the CUDA grid limit"); + assert( + (hidden_size_ == 0 || num_tokens_ <= kMaxSignedIndex / hidden_size_) && + "`MoeSum` output size must fit signed index arithmetic"); + + const auto offsets_fit = [](const Tensor::Shape& shape, + const Tensor::Strides& strides) { + uint64_t max_offset = 0; + constexpr auto kMaxOffset = + static_cast(std::numeric_limits::max()); + + for (Tensor::Size dim = 0; dim < shape.size(); ++dim) { + if (static_cast(shape[dim]) > kMaxOffset) { + return false; + } + if (strides[dim] < 0) { + return false; + } + if (shape[dim] == 0) { + continue; + } + + const auto extent = static_cast(shape[dim] - 1); + const auto stride = static_cast(strides[dim]); + if (extent != 0 && stride > kMaxOffset / extent) { + return false; + } + + const auto term = extent * stride; + if (term > kMaxOffset - max_offset) { + return false; + } + max_offset += term; + } + + return true; + }; + assert(offsets_fit(input.shape(), input.strides()) && + "`MoeSum` input requires non-negative strides with signed offsets"); + + const auto same_device_as_input = [&](const Tensor tensor) { + return tensor.device().type() == input.device().type() && + tensor.device().index() == input.device().index(); + }; + assert(same_device_as_input(output) && + "`MoeSum` input and output must be on the same device"); + assert((!expert_map || topk_ids) && + "`MoeSum` expert_map requires topk_ids"); + + if (topk_ids) { + assert(topk_ids->ndim() == 2 && topk_ids->size(0) == num_tokens_ && + topk_ids->size(1) == topk_ && + "`MoeSum` topk_ids must have shape `[num_tokens, topk]`"); + assert((topk_ids_dtype_ == DataType::kInt32 || + topk_ids_dtype_ == DataType::kInt64) && + "`MoeSum` topk_ids must have int32 or int64 dtype"); + assert(same_device_as_input(*topk_ids) && + "`MoeSum` topk_ids must be on the input device"); + assert(offsets_fit(topk_ids->shape(), topk_ids->strides()) && + "`MoeSum` topk_ids requires non-negative strides with signed " + "offsets"); + } + + if (expert_map) { + assert(expert_map->ndim() == 1 && + expert_map->dtype() == DataType::kInt32 && + "`MoeSum` expert_map must be a 1D int32 tensor"); + assert(same_device_as_input(*expert_map) && + "`MoeSum` expert_map must be on the input device"); + assert(offsets_fit(expert_map->shape(), expert_map->strides()) && + "`MoeSum` expert_map requires non-negative strides with signed " + "offsets"); + } + } + + void operator()(const Tensor input, Tensor output) const { + (*this)(input, std::nullopt, std::nullopt, output); + } + + virtual void operator()(const Tensor input, std::optional topk_ids, + std::optional expert_map, + Tensor output) const = 0; + + protected: + void ValidateCallMetadata(const Tensor input, std::optional topk_ids, + std::optional expert_map, + const Tensor output) const { + const auto same_device_as_descriptor = [&](const Tensor tensor) { + return tensor.device().type() == device_type_ && + tensor.device().index() == device_index_; + }; + auto matches = + input.ndim() == 3 && input.size(0) == num_tokens_ && + input.size(1) == topk_ && input.size(2) == hidden_size_ && + input.strides() == input_strides_ && input.dtype() == dtype_ && + same_device_as_descriptor(input) && output.ndim() == 2 && + output.size(0) == num_tokens_ && output.size(1) == hidden_size_ && + output.strides() == output_strides_ && output.dtype() == dtype_ && + same_device_as_descriptor(output) && + topk_ids.has_value() == has_topk_ids_ && + expert_map.has_value() == has_expert_map_; + + if (matches && topk_ids) { + matches = topk_ids->ndim() == 2 && topk_ids->size(0) == num_tokens_ && + topk_ids->size(1) == topk_ && + topk_ids->stride(0) == topk_ids_token_stride_ && + topk_ids->stride(1) == topk_ids_slot_stride_ && + topk_ids->dtype() == topk_ids_dtype_ && + same_device_as_descriptor(*topk_ids); + } + + if (matches && expert_map) { + matches = expert_map->ndim() == 1 && + expert_map->numel() == expert_map_size_ && + expert_map->stride(0) == expert_map_stride_ && + expert_map->dtype() == DataType::kInt32 && + same_device_as_descriptor(*expert_map); + } + + assert(matches && "`MoeSum` call metadata must match descriptor"); + } + + Tensor::Size num_tokens_{0}; + + Tensor::Size topk_{0}; + + Tensor::Size hidden_size_{0}; + + Tensor::Strides input_strides_; + + Tensor::Strides output_strides_; + + DataType dtype_; + + Device::Type device_type_; + + bool has_topk_ids_{false}; + + DataType topk_ids_dtype_; + + Tensor::Stride topk_ids_token_stride_{0}; + + Tensor::Stride topk_ids_slot_stride_{0}; + + bool has_expert_map_{false}; + + Tensor::Size expert_map_size_{0}; + + Tensor::Stride expert_map_stride_{0}; + + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_MOE_SUM_H_ diff --git a/src/native/cuda/nvidia/ops/moe_sum/kernel.cu b/src/native/cuda/nvidia/ops/moe_sum/kernel.cu new file mode 100644 index 000000000..82d6a8482 --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_sum/kernel.cu @@ -0,0 +1,86 @@ +#include +#include + +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/ops/moe_sum/kernel.h" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/moe_sum/kernel.cuh" + +namespace infini::ops { +namespace { + +class DeviceGuard { + public: + explicit DeviceGuard(int device_index) { + auto status = cudaGetDevice(&previous_device_); + assert(status == cudaSuccess && + "`MoeSum` failed to query the current CUDA device"); + + if (previous_device_ != device_index) { + status = cudaSetDevice(device_index); + assert(status == cudaSuccess && + "`MoeSum` failed to select the input CUDA device"); + restore_ = true; + } + } + + ~DeviceGuard() { + if (restore_) { + const auto status = cudaSetDevice(previous_device_); + assert(status == cudaSuccess && + "`MoeSum` failed to restore the CUDA device"); + } + } + + private: + int previous_device_{0}; + + bool restore_{false}; +}; + +} // namespace + +void Operator::operator()( + const Tensor input, std::optional topk_ids, + std::optional expert_map, Tensor output) const { + ValidateCallMetadata(input, topk_ids, expert_map, output); + + if (num_tokens_ == 0 || hidden_size_ == 0) { + return; + } + + DeviceGuard device_guard{device_index_}; + auto stream = static_cast(stream_ ? stream_ : 0); + constexpr unsigned int kBlockSize = 256; + using DataTypes = ConcatType, ReducedFloatTypes>; + using IndexTypes = List; + + DispatchFunc( + {static_cast(dtype_), static_cast(topk_ids_dtype_)}, + [&](auto list_tag) { + using Data = TypeMapType(list_tag)>; + using Index = TypeMapType(list_tag)>; + + MoeSumKernel + <<(num_tokens_), kBlockSize, 0, stream>>>( + reinterpret_cast(input.data()), + topk_ids ? reinterpret_cast(topk_ids->data()) + : nullptr, + expert_map + ? reinterpret_cast(expert_map->data()) + : nullptr, + reinterpret_cast(output.data()), + static_cast(topk_), static_cast(hidden_size_), + input_strides_[0], input_strides_[1], input_strides_[2], + topk_ids_token_stride_, topk_ids_slot_stride_, + static_cast(expert_map_size_), expert_map_stride_); + }, + "Operator::operator()"); + + const auto status = cudaGetLastError(); + assert(status == cudaSuccess && "`MoeSum` CUDA kernel launch failed"); +} + +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/moe_sum/kernel.h b/src/native/cuda/nvidia/ops/moe_sum/kernel.h new file mode 100644 index 000000000..9cfd91b39 --- /dev/null +++ b/src/native/cuda/nvidia/ops/moe_sum/kernel.h @@ -0,0 +1,24 @@ +#ifndef INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_ + +#include + +#include "base/moe_sum.h" + +namespace infini::ops { + +template <> +class Operator : public MoeSum { + public: + using MoeSum::MoeSum; + + using MoeSum::operator(); + + void operator()(const Tensor input, std::optional topk_ids, + std::optional expert_map, + Tensor output) const override; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_MOE_SUM_KERNEL_H_ diff --git a/src/native/cuda/ops/moe_sum/kernel.cuh b/src/native/cuda/ops/moe_sum/kernel.cuh new file mode 100644 index 000000000..c6318fe0c --- /dev/null +++ b/src/native/cuda/ops/moe_sum/kernel.cuh @@ -0,0 +1,52 @@ +#ifndef INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_ + +#include + +#include "native/cuda/caster.cuh" + +namespace infini::ops { + +template +__global__ void MoeSumKernel( + const Data* __restrict__ input, const Index* __restrict__ topk_ids, + const int32_t* __restrict__ expert_map, Data* __restrict__ output, + int64_t topk, int64_t hidden_size, int64_t input_token_stride, + int64_t input_slot_stride, int64_t input_hidden_stride, + int64_t topk_ids_token_stride, int64_t topk_ids_slot_stride, + int64_t expert_map_size, int64_t expert_map_stride) { + const int64_t token = static_cast(blockIdx.x); + + for (int64_t hidden = static_cast(threadIdx.x); hidden < hidden_size; + hidden += static_cast(blockDim.x)) { + float sum = 0.0f; + + for (int64_t slot = 0; slot < topk; ++slot) { + if (topk_ids != nullptr) { + const auto expert_id = + static_cast(topk_ids[token * topk_ids_token_stride + + slot * topk_ids_slot_stride]); + if (expert_id < 0) { + continue; + } + if (expert_map != nullptr && + (expert_id >= expert_map_size || + expert_map[expert_id * expert_map_stride] < 0)) { + continue; + } + } + + const auto offset = token * input_token_stride + + slot * input_slot_stride + + hidden * input_hidden_stride; + sum += Caster::template Cast(input[offset]); + } + + output[token * hidden_size + hidden] = + Caster::template Cast(sum); + } +} + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_MOE_SUM_KERNEL_CUH_ diff --git a/tests/conftest.py b/tests/conftest.py index 5b825bb71..df2393cd7 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -196,6 +196,7 @@ def _is_smoke_item(item): "tests/test_gemm.py": _is_smoke_gemm_case, "tests/test_linear.py": _is_smoke_linear_case, "tests/test_matmul.py": _is_smoke_matmul_case, + "tests/test_moe_sum.py": _is_smoke_moe_sum_case, "tests/test_mul.py": _is_smoke_mul_case, "tests/test_rms_norm.py": _is_smoke_rms_norm_case, "tests/test_cutlass_scaled_mm.py": _is_smoke_cutlass_scaled_mm_case, @@ -385,6 +386,14 @@ def _is_smoke_rms_norm_case(params): ) +def _is_smoke_moe_sum_case(params): + return ( + params.get("topk") == 2 + and params.get("hidden_size") == 64 + and params.get("dtype") == torch.float16 + ) + + def _is_smoke_cutlass_scaled_mm_case(params): return ( params.get("m") == 16 diff --git a/tests/test_moe_sum.py b/tests/test_moe_sum.py new file mode 100644 index 000000000..8a342d62a --- /dev/null +++ b/tests/test_moe_sum.py @@ -0,0 +1,279 @@ +import subprocess +import sys +import textwrap + +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +if not hasattr(infini.ops, "MoeSum"): + pytest.skip("`MoeSum` is not available on this platform", allow_module_level=True) + + +@pytest.mark.parametrize("topk, hidden_size", ((1, 7), (2, 64), (4, 129))) +@pytest.mark.parametrize( + "dtype, rtol, atol", + ( + (torch.float32, 1e-6, 1e-6), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 1e-2, 1e-2), + ), +) +def test_moe_sum( + topk, + hidden_size, + dtype, + device, + implementation_index, + rtol, + atol, +): + if device != "cuda": + pytest.skip("moe_sum requires the NVIDIA backend") + + input = torch.randn((5, topk, hidden_size), dtype=dtype, device=device) + out = torch.full( + (input.size(0), hidden_size), torch.nan, dtype=dtype, device=device + ) + + result = infini.ops.moe_sum( + input, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + assert result is None + expected = input.float().sum(dim=1).to(dtype) + torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("topk_ids_dtype", (torch.int32, torch.int64)) +@pytest.mark.parametrize("has_expert_map", (False, True)) +@pytest.mark.parametrize( + "dtype, rtol, atol", + ( + (torch.float32, 1e-6, 1e-6), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 1e-2, 1e-2), + ), +) +def test_moe_sum_routing( + topk_ids_dtype, + has_expert_map, + dtype, + device, + implementation_index, + rtol, + atol, +): + if device != "cuda": + pytest.skip("moe_sum requires the NVIDIA backend") + + input = torch.randn((3, 4, 17), dtype=dtype, device=device) + topk_ids = torch.tensor( + ((0, 1, -1, 2), (2, -1, 1, 0), (1, 2, 0, -1)), + dtype=topk_ids_dtype, + device=device, + ) + expert_map = torch.tensor((0, -1, 1), dtype=torch.int32, device=device) + out = torch.full((3, 17), torch.nan, dtype=dtype, device=device) + + result = infini.ops.moe_sum( + input, + topk_ids, + expert_map if has_expert_map else None, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + assert result is None + valid_ids = topk_ids.clamp_min(0).to(torch.int64) + active = topk_ids >= 0 + if has_expert_map: + active &= expert_map[valid_ids] >= 0 + expected = (input.float() * active.unsqueeze(-1)).sum(dim=1).to(dtype) + torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) + + +def test_moe_sum_non_contiguous_positive_strides(device, implementation_index): + if device != "cuda": + pytest.skip("moe_sum requires the NVIDIA backend") + + input = torch.empty_strided( + (3, 4, 17), (167, 40, 2), dtype=torch.float32, device=device + ) + input.normal_() + topk_ids = torch.empty_strided((3, 4), (11, 2), dtype=torch.int64, device=device) + topk_ids.copy_( + torch.tensor( + ((0, 1, -1, 2), (2, -1, 1, 0), (1, 2, 0, -1)), + dtype=topk_ids.dtype, + device=device, + ) + ) + output = torch.full((3, 17), torch.nan, dtype=input.dtype, device=device) + + infini.ops.moe_sum( + input, + topk_ids, + None, + output, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + expected = (input * (topk_ids >= 0).unsqueeze(-1)).sum(dim=1) + torch.testing.assert_close(output, expected) + + +def test_moe_sum_descriptor_reuses_metadata_with_fresh_tensors( + device, implementation_index +): + if device != "cuda": + pytest.skip("moe_sum requires the NVIDIA backend") + + input = torch.randn((3, 2, 17), dtype=torch.float32, device=device) + output = torch.empty((3, 17), dtype=input.dtype, device=device) + op = infini.ops.MoeSum(input, output) + fresh_input = torch.randn_like(input) + fresh_output = torch.full_like(output, torch.nan) + + result = op(fresh_input, fresh_output) + + assert result is None + torch.testing.assert_close(fresh_output, fresh_input.sum(dim=1)) + + +@pytest.mark.parametrize( + "mismatch", + ( + "shape", + "stride", + "dtype", + "device", + "optional_presence", + "optional_metadata", + ), +) +def test_moe_sum_descriptor_rejects_metadata_mismatch(mismatch, device): + if device != "cuda": + pytest.skip("moe_sum requires the NVIDIA backend") + + mismatch_setup = { + "shape": """ +call_input = torch.randn((3, 2, 18), dtype=torch.float32, device="cuda") +call_output = torch.empty((3, 18), dtype=torch.float32, device="cuda") +op(call_input, call_output) +""", + "stride": """ +call_input = torch.empty_strided( + (3, 2, 17), (50, 22, 1), dtype=torch.float32, device="cuda" +) +call_input.normal_() +op(call_input, torch.empty_like(output)) +""", + "dtype": """ +call_input = input.to(torch.float16) +call_output = output.to(torch.float16) +op(call_input, call_output) +""", + "device": """ +op(torch.randn_like(input, device="cpu"), torch.empty_like(output, device="cpu")) +""", + "optional_presence": """ +topk_ids = torch.zeros((3, 2), dtype=torch.int32, device="cuda") +op_with_routing = infini.ops.MoeSum(input, topk_ids, None, output) +op_with_routing(torch.randn_like(input), torch.empty_like(output)) +""", + "optional_metadata": """ +topk_ids = torch.zeros((3, 2), dtype=torch.int32, device="cuda") +op_with_routing = infini.ops.MoeSum(input, topk_ids, None, output) +call_topk_ids = topk_ids.to(torch.int64) +op_with_routing( + torch.randn_like(input), call_topk_ids, None, torch.empty_like(output) +) +""", + }[mismatch] + code = ( + textwrap.dedent( + """ + import infini.ops + import torch + + input = torch.randn((3, 2, 17), dtype=torch.float32, device="cuda") + output = torch.empty((3, 17), dtype=torch.float32, device="cuda") + op = infini.ops.MoeSum(input, output) + """ + ) + + textwrap.dedent(mismatch_setup) + + "\ntorch.cuda.synchronize()\n" + ) + + completed = subprocess.run( + (sys.executable, "-c", code), + check=False, + capture_output=True, + text=True, + ) + + assert completed.returncode != 0 + assert "`MoeSum` call metadata must match descriptor" in completed.stderr + + +def test_moe_sum_non_default_stream(device, implementation_index): + if device != "cuda": + pytest.skip("non-default CUDA streams require the NVIDIA backend") + + input = torch.randn((5, 2, 64), dtype=torch.float16, device=device) + out = torch.full((5, 64), torch.nan, dtype=input.dtype, device=device) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + + infini.ops.moe_sum( + input, + None, + None, + out, + stream=stream.cuda_stream, + implementation_index=implementation_index, + ) + + stream.synchronize() + expected = input.float().sum(dim=1).to(input.dtype) + torch.testing.assert_close(out, expected, rtol=1e-3, atol=1e-3) + + +def test_moe_sum_multi_gpu_device_guard(device, implementation_index): + if device != "cuda" or torch.cuda.device_count() < 2: + pytest.skip("multi-GPU device guard test requires two NVIDIA GPUs") + + original_device = torch.cuda.current_device() + + try: + torch.cuda.set_device(0) + target_device = torch.device("cuda:1") + input = torch.randn((5, 4, 64), dtype=torch.bfloat16, device=target_device) + out = torch.full((5, 64), torch.nan, dtype=input.dtype, device=target_device) + stream = torch.cuda.Stream(device=target_device) + stream.wait_stream(torch.cuda.current_stream(target_device)) + + infini.ops.moe_sum( + input, + None, + None, + out, + stream=stream.cuda_stream, + implementation_index=implementation_index, + ) + + assert torch.cuda.current_device() == 0 + stream.synchronize() + expected = input.float().sum(dim=1).to(input.dtype) + torch.testing.assert_close(out, expected, rtol=1e-3, atol=1e-3) + finally: + torch.cuda.set_device(original_device)