From bd80a6fbffae4094bbb8e98d8d4ee1e636e6cee5 Mon Sep 17 00:00:00 2001 From: rubik Date: Fri, 18 Sep 2026 15:55:59 +0800 Subject: [PATCH] feat(metax): add moe_fused_dense operator for fused MoE expert forward Implement fused dense forward on Metax backend, combining gate-up projection (w1/w3), SiLU/GELU activation, down projection (w2), and token shuffle into a single dispatch. Supports both fp16 and bf16 precision for end-to-end MoE layer inference. --- src/base/moe_fused_dense.h | 99 +++++++ .../cuda/metax/ops/moe_fused_dense/kernel.h | 21 ++ .../cuda/ops/moe_fused_dense/kernel.cuh | 130 ++++++++ src/native/cuda/ops/moe_fused_dense/kernel.h | 279 ++++++++++++++++++ tests/test_moe_fused_dense.py | 272 +++++++++++++++++ 5 files changed, 801 insertions(+) create mode 100644 src/base/moe_fused_dense.h create mode 100644 src/native/cuda/metax/ops/moe_fused_dense/kernel.h create mode 100644 src/native/cuda/ops/moe_fused_dense/kernel.cuh create mode 100644 src/native/cuda/ops/moe_fused_dense/kernel.h create mode 100644 tests/test_moe_fused_dense.py diff --git a/src/base/moe_fused_dense.h b/src/base/moe_fused_dense.h new file mode 100644 index 000000000..5450c5448 --- /dev/null +++ b/src/base/moe_fused_dense.h @@ -0,0 +1,99 @@ +#ifndef INFINI_OPS_BASE_MOE_FUSED_DENSE_H_ +#define INFINI_OPS_BASE_MOE_FUSED_DENSE_H_ + +#include + +#include "operator.h" + +namespace infini::ops { + +class MoeFusedDense : public Operator { + public: + MoeFusedDense(Tensor output, Tensor hidden_states, Tensor w13, Tensor w2, + Tensor topk_weights, Tensor topk_ids, Tensor sorted_token_ids, + Tensor expert_ids, Tensor num_tokens_post_padded) + : output_shape_{output.shape()}, + hidden_states_shape_{hidden_states.shape()}, + w13_shape_{w13.shape()}, + w2_shape_{w2.shape()}, + topk_weights_shape_{topk_weights.shape()}, + topk_ids_shape_{topk_ids.shape()}, + sorted_token_ids_shape_{sorted_token_ids.shape()}, + expert_ids_shape_{expert_ids.shape()}, + num_tokens_post_padded_shape_{num_tokens_post_padded.shape()}, + num_tokens_{hidden_states.size(0)}, + hidden_size_{hidden_states.size(1)}, + num_experts_{w13.size(0)}, + intermediate_size_{w2.size(2)}, + topk_{static_cast(topk_ids.size(1))}, + max_num_tokens_padded_{sorted_token_ids.numel()}, + max_num_blocks_{expert_ids.numel()}, + dtype_{output.dtype()} { + assert(output.ndim() == 2 && hidden_states.ndim() == 2 && + "`MoeFusedDense` output and hidden_states must be 2D tensors"); + assert(w13.ndim() == 3 && w2.ndim() == 3 && + "`MoeFusedDense` w13 and w2 must be 3D tensors"); + assert(topk_weights.ndim() == 2 && topk_ids.ndim() == 2 && + "`MoeFusedDense` topk_weights and topk_ids must be 2D tensors"); + assert( + sorted_token_ids.ndim() == 1 && expert_ids.ndim() == 1 && + "`MoeFusedDense` sorted_token_ids and expert_ids must be 1D tensors"); + assert(num_tokens_post_padded.ndim() == 1 && + num_tokens_post_padded.numel() == 1 && + "`MoeFusedDense` num_tokens_post_padded must be a scalar tensor"); + assert( + output.dtype() == hidden_states.dtype() && + output.dtype() == w13.dtype() && output.dtype() == w2.dtype() && + "`MoeFusedDense` all weight/output tensors must have the same dtype"); + assert(topk_weights.dtype() == DataType::kFloat32 && + "`MoeFusedDense` topk_weights must be float32"); + assert(topk_ids.dtype() == DataType::kInt32 && + sorted_token_ids.dtype() == DataType::kInt32 && + expert_ids.dtype() == DataType::kInt32 && + num_tokens_post_padded.dtype() == DataType::kInt32 && + "`MoeFusedDense` index tensors must be int32"); + assert(output.size(0) == num_tokens_ && output.size(1) == hidden_size_ && + "`MoeFusedDense` output shape must be (num_tokens, hidden_size)"); + assert(w13.size(2) == hidden_size_ && + "`MoeFusedDense` w13 must have shape (num_experts, w13_rows, " + "hidden_size)"); + assert(w2.size(0) == num_experts_ && w2.size(1) == hidden_size_ && + "`MoeFusedDense` w2 must have shape (num_experts, hidden_size, " + "intermediate_size)"); + assert(topk_weights.size(0) == num_tokens_ && + topk_weights.size(1) == static_cast(topk_) && + topk_ids.size(0) == num_tokens_ && + "`MoeFusedDense` topk shapes must match"); + assert(max_num_tokens_padded_ >= num_tokens_ * topk_ && + max_num_blocks_ > 0 && + "`MoeFusedDense` invalid sorted_token_ids or expert_ids sizes"); + } + + virtual void operator()(Tensor output, Tensor hidden_states, Tensor w13, + Tensor w2, Tensor topk_weights, Tensor topk_ids, + Tensor sorted_token_ids, Tensor expert_ids, + Tensor num_tokens_post_padded) const = 0; + + protected: + Tensor::Shape output_shape_; + Tensor::Shape hidden_states_shape_; + Tensor::Shape w13_shape_; + Tensor::Shape w2_shape_; + Tensor::Shape topk_weights_shape_; + Tensor::Shape topk_ids_shape_; + Tensor::Shape sorted_token_ids_shape_; + Tensor::Shape expert_ids_shape_; + Tensor::Shape num_tokens_post_padded_shape_; + Tensor::Size num_tokens_{0}; + Tensor::Size hidden_size_{0}; + Tensor::Size num_experts_{0}; + Tensor::Size intermediate_size_{0}; + int64_t topk_{0}; + Tensor::Size max_num_tokens_padded_{0}; + Tensor::Size max_num_blocks_{0}; + DataType dtype_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/moe_fused_dense/kernel.h b/src/native/cuda/metax/ops/moe_fused_dense/kernel.h new file mode 100644 index 000000000..0920e3c77 --- /dev/null +++ b/src/native/cuda/metax/ops/moe_fused_dense/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_ +#define INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_ + +#include + +#include "native/cuda/metax/blas.h" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/moe_fused_dense/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaMoeFusedDense> { + public: + using CudaMoeFusedDense>::CudaMoeFusedDense; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_METAX_MOE_FUSED_DENSE_KERNEL_H_ \ No newline at end of file diff --git a/src/native/cuda/ops/moe_fused_dense/kernel.cuh b/src/native/cuda/ops/moe_fused_dense/kernel.cuh new file mode 100644 index 000000000..1ed1cd247 --- /dev/null +++ b/src/native/cuda/ops/moe_fused_dense/kernel.cuh @@ -0,0 +1,130 @@ +#ifndef INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_ +#define INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_ + +#include + +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { + +// Computes the inclusive prefix sum of `counts[0..num_experts-1]`, storing the +// exclusive prefix in `offsets[0..num_experts-1]` and the total in +// `offsets[num_experts]`. Only one thread (thread 0) is used. +__global__ void ExclusivePrefixCountsKernel(const int* counts, int* offsets, + int num_experts) { + if (threadIdx.x == 0) { + offsets[0] = 0; + int sum = 0; + for (int i = 0; i < num_experts; ++i) { + sum += counts[i]; + offsets[i + 1] = sum; + } + } +} + +// For each aligned block of `block_size` tokens, atomically adds `block_size` +// to the count of the expert that owns the block. Rows whose expert id is out +// of range (padding rows) are skipped. +// NOTE: `num_tokens_post_padded` is a host-side scalar (not a device pointer), +// because Metax does not support device-side pointer dereference inside +// kernels. +__global__ void CountAlignedExpertsKernel(const int* expert_ids, + int num_tokens_post_padded, + int* counts, int num_experts, + int block_size) { + int block = blockIdx.x * blockDim.x + threadIdx.x; + int num_blocks = (num_tokens_post_padded + block_size - 1) / block_size; + if (block >= num_blocks) { + return; + } + int expert = expert_ids[block]; + if (expert >= 0 && expert < num_experts) { + atomicAdd(counts + expert, block_size); + } +} + +// Gathers hidden states into a packed, expert-bucketed buffer. Each output row +// `row` maps to `pair = sorted_token_ids[row]`, and its hidden state is copied +// from the source token `pair / topk`. Rows pointing past the valid pair range +// (padding) are zero-filled. +template +__global__ void PackHiddenAlignedKernel(const T* hidden, + const int* sorted_token_ids, + int* output_permutation, + T* packed_hidden, int pairs, int topk, + int hidden_size, + int max_num_tokens_padded) { + int row = blockIdx.x; + int tid = threadIdx.x; + if (row >= max_num_tokens_padded) { + return; + } + int pair = sorted_token_ids[row]; + if (pair >= 0 && pair < pairs) { + if (tid == 0) { + output_permutation[pair] = row; + } + int token = pair / topk; + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = + hidden[static_cast(token) * hidden_size + h]; + } + } else { + for (int h = tid; h < hidden_size; h += blockDim.x) { + packed_hidden[static_cast(row) * hidden_size + h] = + Caster::template Cast(0.0f); + } + } +} + +// SwiGLU activation: out = up * silu(gate), where silu(x) = x / (1 + exp(-x)). +template +__global__ void SwigluKernel(const T* gate_up, T* activated, int rows, + int intermediate_size) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + int total = rows * intermediate_size; + if (idx >= total) { + return; + } + int row = idx / intermediate_size; + int col = idx - row * intermediate_size; + const T* base = gate_up + static_cast(row) * intermediate_size * 2; + float gate = Caster::template Cast(base[col]); + float up = Caster::template Cast(base[intermediate_size + col]); + float silu = gate / (1.0f + expf(-gate)); + activated[idx] = Caster::template Cast(up * silu); +} + +// Scatters the weighted expert outputs back to the original token rows. For +// each token `token`, sums over its `topk` pairs, gathering rows through +// `output_permutation` and weighting by `topk_weights`. +template +__global__ void ApplyShuffleMulSumKernel( + const T* __restrict__ expert_out, T* __restrict__ out, + const int* __restrict__ output_permutation, + const float* __restrict__ topk_weights, int num_tokens, int topk, + int hidden_size) { + int token = blockIdx.x; + if (token >= num_tokens) { + return; + } + + for (int h = threadIdx.x; h < hidden_size; h += blockDim.x) { + float sum = 0.0f; + for (int k = 0; k < topk; ++k) { + int pair = token * topk + k; + int src_row = output_permutation[pair]; + if (src_row >= 0) { + sum += Caster::template Cast( + expert_out[static_cast(src_row) * hidden_size + h]) * + topk_weights[pair]; + } + } + out[static_cast(token) * hidden_size + h] = + Caster::template Cast(sum); + } +} + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_CUH_ \ No newline at end of file diff --git a/src/native/cuda/ops/moe_fused_dense/kernel.h b/src/native/cuda/ops/moe_fused_dense/kernel.h new file mode 100644 index 000000000..48c00c192 --- /dev/null +++ b/src/native/cuda/ops/moe_fused_dense/kernel.h @@ -0,0 +1,279 @@ +#ifndef INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_H_ +#define INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/moe_fused_dense.h" +#include "common/generic_utils.h" +#include "data_type.h" +#include "native/cuda/ops/moe_fused_dense/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +// CUTLASS-compatible MoE fused-dense reference. Metax has no grouped GEMM, so +// the prefill and decode paths share one implementation: sorted tokens are +// packed into per-expert contiguous buckets, one GEMM runs per expert, SwiGLU + +// the down-projection run, and the weighted results are scattered back to the +// original tokens. +// +// The `Backend` type must expose the CUDA-like runtime members (`Malloc`, +// `Free`, `Memcpy`, `Memset`, `DeviceSynchronize`, `Stream`) and the BLAS +// members used here (`BlasHandle`, `BlasCreate`, `BlasSetStream`, +// `BlasGemmStridedBatchedEx`) plus `BLAS_OP_*` / dtype / algorithm constants. +template +class CudaMoeFusedDense : public MoeFusedDense { + public: + CudaMoeFusedDense(Tensor output, Tensor hidden_states, Tensor w13, Tensor w2, + Tensor topk_weights, Tensor topk_ids, + Tensor sorted_token_ids, Tensor expert_ids, + Tensor num_tokens_post_padded) + : MoeFusedDense{output, + hidden_states, + w13, + w2, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded} {} + + ~CudaMoeFusedDense() override = default; + + std::size_t workspace_size_in_bytes() const override { + const std::size_t dtype_size = kDataTypeToSize.at(dtype_); + const std::size_t pairs = num_tokens_ * topk_; + std::size_t bytes = 0; + bytes += AlignUp((num_experts_ + 1) * sizeof(int)); + bytes += AlignUp((num_experts_ + 1) * sizeof(int)); + bytes += AlignUp(pairs * sizeof(int)); + bytes += AlignUp(max_num_tokens_padded_ * hidden_size_ * dtype_size); + bytes += + AlignUp(max_num_tokens_padded_ * intermediate_size_ * 2 * dtype_size); + bytes += AlignUp(max_num_tokens_padded_ * intermediate_size_ * dtype_size); + bytes += AlignUp(max_num_tokens_padded_ * hidden_size_ * dtype_size); + return bytes; + } + + void operator()(Tensor output, Tensor hidden_states, Tensor w13, Tensor w2, + Tensor topk_weights, Tensor topk_ids, Tensor sorted_token_ids, + Tensor expert_ids, + Tensor num_tokens_post_padded) const override { + const int num_tokens = static_cast(num_tokens_); + const int hidden_size = static_cast(hidden_size_); + const int num_experts = static_cast(num_experts_); + const int intermediate_size = static_cast(intermediate_size_); + const int topk = static_cast(topk_); + const int num_tokens_padded = static_cast(max_num_tokens_padded_); + const int max_num_blocks = static_cast(max_num_blocks_); + const int pairs = num_tokens * topk; + const int block_size = static_cast( + (max_num_tokens_padded_ + max_num_blocks_ - 1) / max_num_blocks_); + const std::size_t dtype_size = kDataTypeToSize.at(dtype_); + + auto stream = static_cast(stream_ ? stream_ : 0); + + // ------------------------------------------------------------------ + // If the caller did not provide a workspace (Python binding currently + // leaves Handle::workspace empty), allocate a temporary one ourselves. + // ------------------------------------------------------------------ + const std::size_t ws_bytes = workspace_size_in_bytes(); + bool own_workspace = false; + void* workspace_ptr = workspace_; + if (workspace_ptr == nullptr) { + Backend::Malloc(&workspace_ptr, ws_bytes); + own_workspace = true; + } + + // ------------------------------------------------------------------ + // Workspace layout: counts, offsets, permutation, then data buffers. + // ------------------------------------------------------------------ + std::uint8_t* ptr = reinterpret_cast(workspace_ptr); + int* counts = Advance(ptr, num_experts + 1); + int* offsets = Advance(ptr, num_experts + 1); + int* output_permutation = Advance(ptr, pairs); + void* packed_hidden = + AdvanceBytes(ptr, static_cast(num_tokens_padded) * + hidden_size * dtype_size); + void* gate_up = + AdvanceBytes(ptr, static_cast(num_tokens_padded) * + intermediate_size * 2 * dtype_size); + void* activated = + AdvanceBytes(ptr, static_cast(num_tokens_padded) * + intermediate_size * dtype_size); + void* expert_out = + AdvanceBytes(ptr, static_cast(num_tokens_padded) * + hidden_size * dtype_size); + + auto& blas_handle = GetHandle(); + Backend::BlasSetStream(blas_handle, stream); + + Backend::Memset(output_permutation, 0xFF, pairs * sizeof(int)); + Backend::Memset(counts, 0, (num_experts + 1) * sizeof(int)); + Backend::Memset( + expert_out, 0, + static_cast(num_tokens_padded) * hidden_size * dtype_size); + + // Read num_tokens_post_padded on the host before launch – Metax cannot + // dereference a device pointer inside a kernel. + int num_tokens_post_padded_host = 0; + Backend::Memcpy(&num_tokens_post_padded_host, num_tokens_post_padded.data(), + sizeof(int), Backend::kMemcpyDeviceToHost); + int num_blocks_aligned = + (num_tokens_post_padded_host + block_size - 1) / block_size; + + CountAlignedExpertsKernel<<<(num_blocks_aligned + 255) / 256, 256, 0, + stream>>>( + reinterpret_cast(expert_ids.data()), + num_tokens_post_padded_host, counts, num_experts, block_size); + ExclusivePrefixCountsKernel<<<1, 1, 0, stream>>>(counts, offsets, + num_experts); + + DispatchFunc( + dtype_, + [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + PackHiddenAlignedKernel + <<>>( + reinterpret_cast(hidden_states.data()), + reinterpret_cast(sorted_token_ids.data()), + output_permutation, reinterpret_cast(packed_hidden), + pairs, topk, hidden_size, num_tokens_padded); + }, + "CudaMoeFusedDense::PackHiddenAligned"); + + // Per-expert GEMMs launch from the host, so copy counts/offsets back and + // wait for the counting / packing kernels as well as the D2H copies. + std::vector host_counts(num_experts + 1); + std::vector host_offsets(num_experts + 1); + Backend::Memcpy(host_counts.data(), counts, (num_experts + 1) * sizeof(int), + Backend::kMemcpyDeviceToHost); + Backend::Memcpy(host_offsets.data(), offsets, + (num_experts + 1) * sizeof(int), + Backend::kMemcpyDeviceToHost); + Backend::DeviceSynchronize(); + + const auto blas_dtype = + (dtype_ == DataType::kFloat16) ? Backend::R_16F : Backend::R_16BF; + + // GEMM1: gate_up (m x 2*I) = w13^T (2*I x H) @ packed_hidden^T. + // w13 is col-major, dimensions (H, 2*I) -> OP_T with lda = H. + { + const float alpha = 1.0f; + const float beta = 0.0f; + for (int e = 0; e < num_experts; ++e) { + const int m = host_counts[e]; + if (m <= 0) { + continue; + } + const int off = host_offsets[e]; + Backend::BlasGemmStridedBatchedEx( + blas_handle, Backend::BLAS_OP_T, Backend::BLAS_OP_N, + intermediate_size * 2, m, hidden_size, &alpha, + reinterpret_cast(w13.data()) + + static_cast(e) * intermediate_size * 2 * + hidden_size * dtype_size, + blas_dtype, hidden_size, 0, + reinterpret_cast(packed_hidden) + + static_cast(off) * hidden_size * dtype_size, + blas_dtype, hidden_size, 0, &beta, + reinterpret_cast(gate_up) + + static_cast(off) * intermediate_size * 2 * + dtype_size, + blas_dtype, intermediate_size * 2, 0, /*batch_count=*/1, + Backend::BLAS_COMPUTE_32F, Backend::BLAS_GEMM_DEFAULT); + } + } + + DispatchFunc( + dtype_, + [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + const int total = num_tokens_padded * intermediate_size; + SwigluKernel + <<<(total + 255) / 256, 256, 0, stream>>>( + reinterpret_cast(gate_up), + reinterpret_cast(activated), num_tokens_padded, + intermediate_size); + }, + "CudaMoeFusedDense::Swiglu"); + + // GEMM2: expert_out (m x H) = w2^T (H x I) @ activated^T. + // w2 is col-major, dimensions (I, H) -> OP_T with lda = I. + { + const float alpha = 1.0f; + const float beta = 0.0f; + for (int e = 0; e < num_experts; ++e) { + const int m = host_counts[e]; + if (m <= 0) { + continue; + } + const int off = host_offsets[e]; + Backend::BlasGemmStridedBatchedEx( + blas_handle, Backend::BLAS_OP_T, Backend::BLAS_OP_N, hidden_size, m, + intermediate_size, &alpha, + reinterpret_cast(w2.data()) + + static_cast(e) * hidden_size * intermediate_size * + dtype_size, + blas_dtype, intermediate_size, 0, + reinterpret_cast(activated) + + static_cast(off) * intermediate_size * dtype_size, + blas_dtype, intermediate_size, 0, &beta, + reinterpret_cast(expert_out) + + static_cast(off) * hidden_size * dtype_size, + blas_dtype, hidden_size, 0, /*batch_count=*/1, + Backend::BLAS_COMPUTE_32F, Backend::BLAS_GEMM_DEFAULT); + } + } + + DispatchFunc( + dtype_, + [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + ApplyShuffleMulSumKernel + <<>>( + reinterpret_cast(expert_out), + reinterpret_cast(output.data()), output_permutation, + reinterpret_cast(topk_weights.data()), + num_tokens, topk, hidden_size); + }, + "CudaMoeFusedDense::ApplyShuffleMulSum"); + if (own_workspace) { + Backend::Free(workspace_ptr); + } + } + + protected: + static std::size_t AlignUp(std::size_t value, std::size_t alignment = 16) { + return (value + alignment - 1) / alignment * alignment; + } + + template + static T* Advance(std::uint8_t*& ptr, std::size_t count) { + T* out = reinterpret_cast(ptr); + ptr += AlignUp(count * sizeof(T)); + return out; + } + + static void* AdvanceBytes(std::uint8_t*& ptr, std::size_t bytes) { + void* out = reinterpret_cast(ptr); + ptr += AlignUp(bytes); + return out; + } + + static typename Backend::BlasHandle& GetHandle() { + thread_local typename Backend::BlasHandle handle = []() { + typename Backend::BlasHandle h; + Backend::BlasCreate(&h); + return h; + }(); + return handle; + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_MOE_FUSED_DENSE_KERNEL_H_ \ No newline at end of file diff --git a/tests/test_moe_fused_dense.py b/tests/test_moe_fused_dense.py new file mode 100644 index 000000000..62fb66f4a --- /dev/null +++ b/tests/test_moe_fused_dense.py @@ -0,0 +1,272 @@ +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import infini.ops +import pytest + +import torch +from tests.utils import empty_strided, get_stream + + +def _ref_moe_align(topk_ids, num_experts, block_size, pad_sorted_token_ids): + """Reference: pure-Python moe_align (re-implemented here for independence).""" + numel = topk_ids.numel() + + # Count tokens per expert + counts = [0] * (num_experts + 1) + flat_ids = topk_ids.view(-1) + for i in range(flat_ids.numel()): + expert_id = int(flat_ids[i].item()) + # No expert_map in fused_dense test path + expert_id += 1 # Shift by 1: index 0 is reserved + if 1 <= expert_id <= num_experts: + counts[expert_id] += 1 + + # Padded cumulative sums + prefix = [0] * (num_experts + 1) + total_tokens_post_pad = 0 + for i in range(1, num_experts + 1): + padded = (counts[i] + block_size - 1) // block_size * block_size + prefix[i] = total_tokens_post_pad + padded + total_tokens_post_pad = prefix[i] + + # Buffer size must equal the actual padded output length + max_num_tokens_padded = total_tokens_post_pad + max_num_blocks = (max_num_tokens_padded + block_size - 1) // block_size + + # Fill sorted_token_ids + sorted_token_ids = torch.full( + (max_num_tokens_padded,), fill_value=numel, dtype=torch.int32 + ) + positions = [0] * (num_experts + 1) + for i in range(flat_ids.numel()): + expert_id = int(flat_ids[i].item()) + expert_id += 1 + if 1 <= expert_id <= num_experts: + rank = prefix[expert_id - 1] + positions[expert_id] + sorted_token_ids[rank] = i + positions[expert_id] += 1 + + # expert_ids per block + expert_ids = torch.full((max_num_blocks,), -1, dtype=torch.int32) + num_blocks = total_tokens_post_pad // block_size + for i in range(num_blocks): + block_start = i * block_size + expert = 0 + for e in range(1, num_experts + 1): + if prefix[e] > block_start: + expert = e - 1 + break + expert_ids[i] = expert + + num_tokens_post_padded = torch.tensor( + total_tokens_post_pad, dtype=torch.int32 + ) + return sorted_token_ids, expert_ids, num_tokens_post_padded + + +def _ref_moe_fused_dense( + hidden_states, w13, w2, topk_weights, sorted_token_ids, expert_ids, num_tokens_post_padded +): + """Reference implementation using PyTorch.""" + num_tokens_post_pad = int(num_tokens_post_padded.item()) + num_tokens = hidden_states.size(0) + hidden_size = hidden_states.size(1) + num_experts, twice_inter, _ = w13.shape + inter_size = twice_inter // 2 + topk = topk_weights.size(1) + + hidden_fp32 = hidden_states.float() + w13_fp32 = w13.float() + w2_fp32 = w2.float() + + max_num_tokens_padded = sorted_token_ids.numel() + max_num_blocks = expert_ids.numel() + blk_size = (max_num_tokens_padded + max_num_blocks - 1) // max_num_blocks + + numel = num_tokens * topk + + # Gather hidden states into packed buffer. + # Padding slots (pair_idx == numel) stay zero-filled. + packed_hidden = torch.zeros( + num_tokens_post_pad, hidden_size, dtype=torch.float32 + ) + for i in range(num_tokens_post_pad): + pair_idx = int(sorted_token_ids[i].item()) + if pair_idx < numel: + token = pair_idx // topk + packed_hidden[i] = hidden_fp32[token] + + # Determine expert positions from blocks (includes padding slots, + # because the GPU GEMM runs on the padded count). + expert_positions = [[] for _ in range(num_experts)] + if num_tokens_post_pad > 0: + num_blocks = (num_tokens_post_pad + blk_size - 1) // blk_size + for b in range(num_blocks): + expert = int(expert_ids[b].item()) + if expert < 0 or expert >= num_experts: + continue + start = b * blk_size + end = min(start + blk_size, num_tokens_post_pad) + for pos in range(start, end): + expert_positions[expert].append(pos) + + output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32) + + for e in range(num_experts): + positions = expert_positions[e] + if len(positions) == 0: + continue + + expert_input = packed_hidden[positions] # (m, hidden_size) + + # GEMM1: (m, hidden_size) @ (hidden_size, 2*inter_size) + gate_up = expert_input @ w13_fp32[e].t() # (m, 2*inter_size) + + # SwiGLU + gate = gate_up[:, :inter_size] + up = gate_up[:, inter_size:] + silu = gate / (1 + torch.exp(-gate)) + activated = silu * up # (m, inter_size) + + # GEMM2: (m, inter_size) @ (inter_size, hidden_size) + out = activated @ w2_fp32[e].t() # (m, hidden_size) + + # Scatter with weight (skip padding slots, matching the GPU kernel). + for i_pos, pos in enumerate(positions): + pair = int(sorted_token_ids[pos].item()) + if pair >= numel: + continue + token = pair // topk + weight = topk_weights.view(-1)[pair].item() + output[token] += out[i_pos] * weight + + return output + + +@pytest.mark.parametrize( + "num_tokens,num_experts,topk,block_size,hidden_size,inter_size", + ( + (4, 4, 2, 4, 16, 32), + (8, 8, 2, 4, 16, 32), + (8, 4, 4, 4, 32, 64), + (1, 2, 1, 4, 16, 32), + (16, 8, 2, 8, 32, 64), + ), +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float16, 5e-2, 5e-2), + (torch.bfloat16, 1e-1, 1e-0), + ), +) +def test_moe_fused_dense( + num_tokens, + num_experts, + topk, + block_size, + hidden_size, + inter_size, + dtype, + device, + rtol, + atol, +): + # --- Generate inputs --- + hidden_states = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device) + gating_output = torch.randn(num_tokens, num_experts, dtype=torch.float32, device=device) + + # Top-k softmax + topk_weights = empty_strided( + (num_tokens, topk), None, dtype=torch.float32, device=device + ) + topk_ids = empty_strided( + (num_tokens, topk), None, dtype=torch.int32, device=device + ) + infini.ops.moe_topk_softmax( + topk_weights, + topk_ids, + gating_output, + torch.empty(0, dtype=torch.float32, device=device), + True, # renormalize + 0.0, # moe_softcapping + stream=get_stream(device), + ) + + # Align -- first run reference on CPU to learn the exact output sizes, + # then allocate GPU buffers of that size so the C++ layer computes the + # same block_size that the reference uses. + ref_sorted_cpu, ref_expert_ids_cpu, _ref_num_post_cpu = _ref_moe_align( + topk_ids.cpu(), num_experts, block_size, True + ) + max_num_tokens_padded = ref_sorted_cpu.numel() + max_num_blocks = ref_expert_ids_cpu.numel() + + sorted_token_ids = empty_strided( + (max_num_tokens_padded,), None, dtype=torch.int32, device=device + ) + expert_ids = empty_strided( + (max_num_blocks,), None, dtype=torch.int32, device=device + ) + num_tokens_post_padded = empty_strided( + (1,), None, dtype=torch.int32, device=device + ) + expert_map = torch.empty(0, dtype=torch.int32, device=device) + + infini.ops.moe_align( + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + topk_ids, + expert_map, + num_experts, + block_size, + True, # pad_sorted_token_ids + stream=get_stream(device), + ) + + # Weights + w13 = torch.randn( + num_experts, 2 * inter_size, hidden_size, dtype=dtype, device=device + ) + w2 = torch.randn( + num_experts, hidden_size, inter_size, dtype=dtype, device=device + ) + + # Output + output = empty_strided( + (num_tokens, hidden_size), None, dtype=dtype, device=device + ) + + # Call operator + infini.ops.moe_fused_dense( + output, + hidden_states, + w13, + w2, + topk_weights, + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + stream=get_stream(device), + ) + + # Reference + ref_sorted, ref_expert_ids, ref_num_post = _ref_moe_align( + topk_ids.cpu(), num_experts, block_size, True + ) + ref_out = _ref_moe_fused_dense( + hidden_states.cpu(), + w13.cpu(), + w2.cpu(), + topk_weights.cpu(), + ref_sorted, + ref_expert_ids, + ref_num_post, + ) + + torch.testing.assert_close(output.float().cpu(), ref_out, rtol=rtol, atol=atol)