From dff8cf3cef689f448b9cbd2165cb37b4a55afdc7 Mon Sep 17 00:00:00 2001 From: sunteng Date: Sun, 12 Jul 2026 19:19:00 +0800 Subject: [PATCH] feat(grouped_gemm): add grouped_gemm op (CPU/NVIDIA/MetaX) for Qwen3-MoE Signed-off-by: sunteng --- include/infinicore/ops.hpp | 3 + include/infinicore/ops/grouped_gemm.hpp | 51 +++++ include/infiniop.h | 1 + include/infiniop/ops/grouped_gemm.h | 67 ++++++ .../ops/grouped_gemm/grouped_gemm.cc | 53 +++++ .../ops/grouped_gemm/grouped_gemm_infiniop.cc | 63 ++++++ src/infinicore/tensor/tensor.cc | 10 +- src/infiniop/devices/metax/metax_ht2mc.h | 1 + .../ops/grouped_gemm/cpu/grouped_gemm_cpu.cc | 138 ++++++++++++ .../ops/grouped_gemm/cpu/grouped_gemm_cpu.h | 8 + src/infiniop/ops/grouped_gemm/grouped_gemm.h | 54 +++++ src/infiniop/ops/grouped_gemm/info.h | 91 ++++++++ .../grouped_gemm/metax/grouped_gemm_metax.cc | 172 +++++++++++++++ .../grouped_gemm/metax/grouped_gemm_metax.h | 8 + .../nvidia/grouped_gemm_nvidia.cu | 199 ++++++++++++++++++ .../nvidia/grouped_gemm_nvidia.cuh | 8 + src/infiniop/ops/grouped_gemm/operator.cc | 180 ++++++++++++++++ test/infiniop/grouped_gemm.py | 176 ++++++++++++++++ 18 files changed, 1282 insertions(+), 1 deletion(-) create mode 100644 include/infinicore/ops/grouped_gemm.hpp create mode 100644 include/infiniop/ops/grouped_gemm.h create mode 100644 src/infinicore/ops/grouped_gemm/grouped_gemm.cc create mode 100644 src/infinicore/ops/grouped_gemm/grouped_gemm_infiniop.cc create mode 100644 src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.cc create mode 100644 src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.h create mode 100644 src/infiniop/ops/grouped_gemm/grouped_gemm.h create mode 100644 src/infiniop/ops/grouped_gemm/info.h create mode 100644 src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.cc create mode 100644 src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.h create mode 100644 src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cu create mode 100644 src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cuh create mode 100644 src/infiniop/ops/grouped_gemm/operator.cc create mode 100644 test/infiniop/grouped_gemm.py diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index b5c4ff18f..57ffee10a 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -32,7 +32,9 @@ #include "ops/fused_gated_delta_net_gating.hpp" #include "ops/gelu.hpp" #include "ops/gelutanh.hpp" +#include "ops/grouped_gemm.hpp" #include "ops/hardswish.hpp" +#include "ops/index_add.hpp" #include "ops/hardtanh.hpp" #include "ops/kv_caching.hpp" #include "ops/layer_norm.hpp" @@ -45,6 +47,7 @@ #include "ops/moe_sum.hpp" #include "ops/moe_topk_sigmoid.hpp" #include "ops/moe_topk_softmax.hpp" +#include "ops/mul.hpp" #include "ops/nrm2.hpp" #include "ops/ones.hpp" #include "ops/paged_attention.hpp" diff --git a/include/infinicore/ops/grouped_gemm.hpp b/include/infinicore/ops/grouped_gemm.hpp new file mode 100644 index 000000000..398d1eee2 --- /dev/null +++ b/include/infinicore/ops/grouped_gemm.hpp @@ -0,0 +1,51 @@ +#pragma once + +#include + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(GroupedGemm, + Tensor, + const Tensor &, + const Tensor &, + const Tensor &, + float, + float, + const int32_t *); + +// Variable-batched matmul shared across `num_groups` weight slabs. +// +// Shapes (all row-major): +// a : [M_total, K] +// b : [num_groups, N, K] // out, in -- matches torch linear +// c : [M_total, N] +// group_sizes : [num_groups], int32 // sum == M_total +// +// For each group g: +// c[off_g : off_g + group_sizes[g]] = alpha * a[off_g : off_g + group_sizes[g]] @ b[g].T +// + beta * c[...] +// where `off_g = sum(group_sizes[0..g])`. +// `group_sizes_host` (optional): host-side copy of `group_sizes`. When given, +// device backends use it directly and skip the per-call device->host sync of +// the sizes array. Must stay valid until the call returns; pass nullptr (the +// default) under graph capture or when no host copy is available. +Tensor grouped_gemm(const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha = 1.0f, + float beta = 0.0f, + const int32_t *group_sizes_host = nullptr); + +void grouped_gemm_(Tensor c, + const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha, + float beta, + const int32_t *group_sizes_host = nullptr); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 95bf75a0d..1f37b1824 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -62,6 +62,7 @@ #include "infiniop/ops/gemm.h" #include "infiniop/ops/gptq_marlin_gemm.h" #include "infiniop/ops/gptq_qyblas_gemm.h" +#include "infiniop/ops/grouped_gemm.h" #include "infiniop/ops/hardswish.h" #include "infiniop/ops/hardtanh.h" #include "infiniop/ops/hinge_embedding_loss.h" diff --git a/include/infiniop/ops/grouped_gemm.h b/include/infiniop/ops/grouped_gemm.h new file mode 100644 index 000000000..8f2d5cb21 --- /dev/null +++ b/include/infiniop/ops/grouped_gemm.h @@ -0,0 +1,67 @@ +#ifndef __INFINIOP_GROUPED_GEMM_API_H__ +#define __INFINIOP_GROUPED_GEMM_API_H__ + +#include "../operator_descriptor.h" + +/** + * Variable-batched (a.k.a. "grouped" or "segment") GEMM. + * + * For each group `g` in `[0, num_groups)`: + * rows_g = group_sizes[g] + * off_g = sum(group_sizes[0..g]) + * A_g = a[off_g : off_g + rows_g, :] // [rows_g, K] + * B_g = b[g, :, :] // [N, K] + * c[off_g : off_g + rows_g, :] = alpha * A_g @ B_g^T + beta * c_g + * + * Shapes + * a : [M_total, K] + * b : [num_groups, N, K] -- matches PyTorch / HF linear weight layout + * c : [M_total, N] + * group_sizes : [num_groups], int32, sum == M_total + * + * Constraints + * - M_total = sum(group_sizes); enforced at calculate-time. + * - All `*_desc` must be row-major contiguous on the leading axes. + * - dtype: F16, BF16 or F32 (same for a, b, c). + * - `group_sizes` is `INFINI_DTYPE_I32`. + * - The `group_sizes` data pointer lives on the same device as `a/b/c`. + * CPU backends read it directly; device backends sync it to host + * inside `calculate`. + * - `group_sizes_host` (optional, may be NULL): host-side copy of the same + * int32 group sizes. When provided, device backends use it directly and + * skip the per-call device->host copy + stream sync. The caller must keep + * it valid and consistent with `group_sizes` until the call returns. NULL + * preserves the legacy device-sync behavior. Unsafe under graph capture + * (pass NULL there). + */ +typedef struct InfiniopDescriptor *infiniopGroupedGemmDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateGroupedGemmDescriptor( + infiniopHandle_t handle, + infiniopGroupedGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc); + +__INFINI_C __export infiniStatus_t infiniopGetGroupedGemmWorkspaceSize( + infiniopGroupedGemmDescriptor_t desc, + size_t *size); + +__INFINI_C __export infiniStatus_t infiniopGroupedGemm( + infiniopGroupedGemmDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + const void *group_sizes, + const void *group_sizes_host, + float alpha, + float beta, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyGroupedGemmDescriptor( + infiniopGroupedGemmDescriptor_t desc); + +#endif // __INFINIOP_GROUPED_GEMM_API_H__ diff --git a/src/infinicore/ops/grouped_gemm/grouped_gemm.cc b/src/infinicore/ops/grouped_gemm/grouped_gemm.cc new file mode 100644 index 000000000..c1645215e --- /dev/null +++ b/src/infinicore/ops/grouped_gemm/grouped_gemm.cc @@ -0,0 +1,53 @@ +#include "infinicore/ops/grouped_gemm.hpp" + +#include "../../utils.hpp" + +namespace infinicore::op { +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(GroupedGemm); + +GroupedGemm::GroupedGemm(Tensor c, + const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha, + float beta, + const int32_t *group_sizes_host) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(c, a, b, group_sizes); + INFINICORE_GRAPH_OP_DISPATCH(c->device().getType(), c, a, b, group_sizes, alpha, beta, group_sizes_host); +} + +void GroupedGemm::execute(Tensor c, + const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha, + float beta, + const int32_t *group_sizes_host) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(GroupedGemm, c, a, b, group_sizes, alpha, beta, group_sizes_host); +} + +Tensor grouped_gemm(const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha, + float beta, + const int32_t *group_sizes_host) { + // a: [M_total, K], b: [num_groups, N, K] -> c: [M_total, N]. + Shape shape = a->shape(); + shape[shape.size() - 1] = b->size(1); + auto c = Tensor::empty(shape, a->dtype(), a->device()); + grouped_gemm_(c, a, b, group_sizes, alpha, beta, group_sizes_host); + return c; +} + +void grouped_gemm_(Tensor c, + const Tensor &a, + const Tensor &b, + const Tensor &group_sizes, + float alpha, + float beta, + const int32_t *group_sizes_host) { + GroupedGemm::execute(c, a, b, group_sizes, alpha, beta, group_sizes_host); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/grouped_gemm/grouped_gemm_infiniop.cc b/src/infinicore/ops/grouped_gemm/grouped_gemm_infiniop.cc new file mode 100644 index 000000000..7ce7fc165 --- /dev/null +++ b/src/infinicore/ops/grouped_gemm/grouped_gemm_infiniop.cc @@ -0,0 +1,63 @@ +#include "../infiniop_impl.hpp" +#include "infinicore/ops/grouped_gemm.hpp" + +namespace infinicore::op::grouped_gemm_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, GroupedGemm, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace, c, a, b, group_sizes; + float alpha, beta; + // Optional host-side group sizes; nullptr falls back to the device sync. + // Only valid for immediate (non-recorded) execution -- the pointer is not + // owned, so graph replay must pass nullptr. + const int32_t *group_sizes_host; +}; + +void *plan(Tensor c, const Tensor &a, const Tensor &b, const Tensor &group_sizes, float alpha, float beta, const int32_t *group_sizes_host) { + size_t seed = hash_combine(c, a, b, group_sizes); + + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, GroupedGemm, + seed, c->desc(), a->desc(), b->desc(), group_sizes->desc()); + + INFINIOP_WORKSPACE_TENSOR(workspace, GroupedGemm, descriptor); + + auto planned = new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(c), + graph::GraphTensor(a), + graph::GraphTensor(b), + graph::GraphTensor(group_sizes), + alpha, beta, + group_sizes_host}; + + return planned; +} + +void run(void *planned_meta) { + auto planned = reinterpret_cast(planned_meta); + + INFINICORE_CHECK_ERROR(infiniopGroupedGemm( + planned->descriptor->desc, + planned->workspace->data(), planned->workspace->numel(), + planned->c->data(), + planned->a->data(), + planned->b->data(), + planned->group_sizes->data(), + planned->group_sizes_host, + planned->alpha, + planned->beta, + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(GroupedGemm, &plan, &run, &cleanup); + +} // namespace infinicore::op::grouped_gemm_impl::infiniop diff --git a/src/infinicore/tensor/tensor.cc b/src/infinicore/tensor/tensor.cc index e7e54aa59..b08acbe5d 100644 --- a/src/infinicore/tensor/tensor.cc +++ b/src/infinicore/tensor/tensor.cc @@ -4,6 +4,7 @@ #include "infinicore/context/context.hpp" #include "infinicore/dtype.hpp" +#include #include namespace { @@ -242,7 +243,14 @@ std::shared_ptr TensorImpl::zeros(const Shape &shape, const DataType &dtype, const Device &device, bool pin_memory) { - + // MetaX device-side integer zeros is unreliable (does not actually clear), + // so construct zeros on CPU and copy over. Other backends use the faster + // device-side byte memset. + if (device.getType() == Device::Type::METAX) { + auto cpu = empty(shape, dtype, Device(Device::Type::CPU), false); + std::memset(cpu->data(), 0, cpu->nbytes()); + return cpu->to(device).impl_; + } auto result = empty(shape, dtype, device, pin_memory); context::setDeviceMemoryAsync(result->data(), 0, result->nbytes(), context::getStream()); return result; diff --git a/src/infiniop/devices/metax/metax_ht2mc.h b/src/infiniop/devices/metax/metax_ht2mc.h index 9e0da85ef..64b5f67c1 100644 --- a/src/infiniop/devices/metax/metax_ht2mc.h +++ b/src/infiniop/devices/metax/metax_ht2mc.h @@ -116,6 +116,7 @@ #define hcclAllReduce mcclAllReduce #define hcblasSetStream mcblasSetStream #define hcblasHandle_t mcblasHandle_t +#define hcblasGemmBatchedEx mcblasGemmBatchedEx #define hcblasGemmStridedBatchedEx mcblasGemmStridedBatchedEx #define hcblasGemmEx mcblasGemmEx #define hcblasCreate mcblasCreate diff --git a/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.cc b/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.cc new file mode 100644 index 000000000..52b9c303c --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.cc @@ -0,0 +1,138 @@ +#include "grouped_gemm_cpu.h" +#include "../../../devices/cpu/common_cpu.h" + +#include + +namespace op::grouped_gemm::cpu { + +Descriptor::~Descriptor() = default; + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc) { + + auto handle = reinterpret_cast(handle_); + + CHECK_DTYPE(c_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + + auto result = GroupedGemmInfo::create(c_desc, a_desc, b_desc, group_sizes_desc); + CHECK_RESULT(result); + + *desc_ptr = new Descriptor( + result.take(), 0, + nullptr, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +static void calculate_typed( + const GroupedGemmInfo &info, + void *c_, + const void *a_, + const void *b_, + const int32_t *group_sizes, + float alpha, + float beta) { + + auto a = reinterpret_cast(a_); + auto b = reinterpret_cast(b_); + auto c = reinterpret_cast(c_); + + // Walk each group and dispatch a plain triple-loop matmul on its rows. + // Parallelism is across groups + rows; the tight K loop stays sequential. + ptrdiff_t row_offset = 0; + for (size_t g = 0; g < info.num_groups; ++g) { + ptrdiff_t rows = group_sizes[g]; + if (rows <= 0) { + // Skip degenerate / empty groups; routing is allowed to produce them. + continue; + } + const Tdata *b_g = b + g * info.b_group_stride; + +#pragma omp parallel for collapse(2) + for (ptrdiff_t r = 0; r < rows; ++r) { + for (ptrdiff_t n = 0; n < ptrdiff_t(info.n); ++n) { + const Tdata *a_row = a + (row_offset + r) * info.a_row_stride; + const Tdata *b_row = b_g + n * info.b_row_stride; + Tdata *c_pos = c + (row_offset + r) * info.c_row_stride + n * info.c_col_stride; + + float sum = 0.f; + for (ptrdiff_t k = 0; k < ptrdiff_t(info.k); ++k) { + if constexpr (std::is_same::value || std::is_same::value) { + sum += utils::cast(a_row[k * info.a_col_stride]) + * utils::cast(b_row[k * info.b_col_stride]); + } else { + sum += a_row[k * info.a_col_stride] * b_row[k * info.b_col_stride]; + } + } + + if constexpr (std::is_same::value || std::is_same::value) { + if (beta == 0.f) { + *c_pos = utils::cast(alpha * sum); + } else { + *c_pos = utils::cast(beta * utils::cast(*c_pos) + alpha * sum); + } + } else { + if (beta == 0.f) { + *c_pos = alpha * sum; + } else { + *c_pos = beta * (*c_pos) + alpha * sum; + } + } + } + } + row_offset += rows; + } +} + +infiniStatus_t Descriptor::calculate( + void * /*workspace*/, size_t /*workspace_size*/, + void *c, + const void *a, + const void *b, + const void *group_sizes, + const void *group_sizes_host, + float alpha, + float beta, + void * /*stream*/) const { + + // On CPU `group_sizes` is already host memory, so the optional host copy is + // redundant; prefer it when supplied to mirror the device backends' contract. + auto sizes = reinterpret_cast( + group_sizes_host != nullptr ? group_sizes_host : group_sizes); + + // Validate that group_sizes sums to m_total to keep `row_offset + r` in range. + { + size_t total = 0; + for (size_t g = 0; g < _info.num_groups; ++g) { + if (sizes[g] < 0) { + return INFINI_STATUS_BAD_PARAM; + } + total += size_t(sizes[g]); + } + if (total != _info.m_total) { + return INFINI_STATUS_BAD_PARAM; + } + } + + switch (_info.dtype) { + case INFINI_DTYPE_F16: + calculate_typed(_info, c, a, b, sizes, alpha, beta); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + calculate_typed(_info, c, a, b, sizes, alpha, beta); + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + calculate_typed(_info, c, a, b, sizes, alpha, beta); + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +} // namespace op::grouped_gemm::cpu diff --git a/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.h b/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.h new file mode 100644 index 000000000..7c267963a --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/cpu/grouped_gemm_cpu.h @@ -0,0 +1,8 @@ +#ifndef __GROUPED_GEMM_CPU_H__ +#define __GROUPED_GEMM_CPU_H__ + +#include "../grouped_gemm.h" + +DESCRIPTOR(cpu) + +#endif // __GROUPED_GEMM_CPU_H__ diff --git a/src/infiniop/ops/grouped_gemm/grouped_gemm.h b/src/infiniop/ops/grouped_gemm/grouped_gemm.h new file mode 100644 index 000000000..5a36fdb87 --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/grouped_gemm.h @@ -0,0 +1,54 @@ +#ifndef __GROUPED_GEMM_H__ +#define __GROUPED_GEMM_H__ + +#include "../../operator.h" +#include "info.h" + +// See `ops/gemm/gemm.h` for the rationale of the PImpl-style DESCRIPTOR macro. +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::grouped_gemm::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + GroupedGemmInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + GroupedGemmInfo info, \ + size_t workspace_size_, \ + Opaque *opaque, \ + infiniDevice_t device_type, \ + int device_id) \ + : InfiniopDescriptor{device_type, device_id}, \ + _opaque(opaque), \ + _info(info), \ + _workspace_size(workspace_size_) {} \ + \ + public: \ + ~Descriptor(); \ + \ + size_t workspaceSize() const { return _workspace_size; } \ + \ + static infiniStatus_t create( \ + infiniopHandle_t handle, \ + Descriptor **desc_ptr, \ + infiniopTensorDescriptor_t c_desc, \ + infiniopTensorDescriptor_t a_desc, \ + infiniopTensorDescriptor_t b_desc, \ + infiniopTensorDescriptor_t group_sizes_desc); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *c, \ + const void *a, \ + const void *b, \ + const void *group_sizes, \ + const void *group_sizes_host, \ + float alpha, \ + float beta, \ + void *stream) const; \ + }; \ + } + +#endif // __GROUPED_GEMM_H__ diff --git a/src/infiniop/ops/grouped_gemm/info.h b/src/infiniop/ops/grouped_gemm/info.h new file mode 100644 index 000000000..dc42bc36e --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/info.h @@ -0,0 +1,91 @@ +#ifndef __GROUPED_GEMM_INFO_H__ +#define __GROUPED_GEMM_INFO_H__ + +#include "../../../utils.h" +#include "../../operator.h" +#include "../../tensor.h" + +namespace op::grouped_gemm { + +// Shape/stride metadata captured at descriptor-creation time. +// +// We require the leading-axis layouts that match `torch.matmul` on row-major +// tensors so we can address each group's slab as `base + offset * row_stride`. +struct GroupedGemmInfo { + infiniDtype_t dtype; + size_t num_groups; + size_t m_total; + size_t n; + size_t k; + + // Strides in elements (not bytes) on the row/col axes. + ptrdiff_t a_row_stride; // stride between A rows + ptrdiff_t a_col_stride; // == 1 expected + ptrdiff_t c_row_stride; // stride between C rows + ptrdiff_t c_col_stride; // == 1 expected + ptrdiff_t b_group_stride; // stride between expert slabs of B + ptrdiff_t b_row_stride; // stride between rows of B[g] + ptrdiff_t b_col_stride; // == 1 expected + + static utils::Result create( + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc) { + + if (a_desc->ndim() != 2 || c_desc->ndim() != 2 || b_desc->ndim() != 3) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + if (group_sizes_desc->ndim() != 1) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + if (group_sizes_desc->dtype() != INFINI_DTYPE_I32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + auto dtype = c_desc->dtype(); + if (dtype != a_desc->dtype() || dtype != b_desc->dtype()) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + size_t m_total = a_desc->dim(0); + size_t k_a = a_desc->dim(1); + size_t num_groups = b_desc->dim(0); + size_t n_b = b_desc->dim(1); + size_t k_b = b_desc->dim(2); + size_t m_c = c_desc->dim(0); + size_t n_c = c_desc->dim(1); + + if (k_a != k_b || m_total != m_c || n_b != n_c) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + if (group_sizes_desc->dim(0) != num_groups) { + return INFINI_STATUS_BAD_TENSOR_SHAPE; + } + + // We need the inner (K, N) axes to be unit-strided so each slab is a + // dense matrix that cuBLAS / our cpu loops can consume directly. + if (a_desc->stride(1) != 1 || b_desc->stride(2) != 1 || c_desc->stride(1) != 1) { + return INFINI_STATUS_BAD_TENSOR_STRIDES; + } + + GroupedGemmInfo info; + info.dtype = dtype; + info.num_groups = num_groups; + info.m_total = m_total; + info.n = n_b; + info.k = k_a; + info.a_row_stride = a_desc->stride(0); + info.a_col_stride = a_desc->stride(1); + info.c_row_stride = c_desc->stride(0); + info.c_col_stride = c_desc->stride(1); + info.b_group_stride = b_desc->stride(0); + info.b_row_stride = b_desc->stride(1); + info.b_col_stride = b_desc->stride(2); + return utils::Result(info); + } +}; + +} // namespace op::grouped_gemm + +#endif // __GROUPED_GEMM_INFO_H__ diff --git a/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.cc b/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.cc new file mode 100644 index 000000000..3b1b1c8af --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.cc @@ -0,0 +1,172 @@ +#include "grouped_gemm_metax.h" +#include "../../../devices/metax/metax_common.h" +#include "../../../devices/metax/metax_handle.h" + +#define CHECK_METAX(API) CHECK_INTERNAL(API, hcSuccess) + +#include +#include + +namespace op::grouped_gemm::metax { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc) { + + auto handle = reinterpret_cast(handle_); + + CHECK_DTYPE(c_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + + auto result = GroupedGemmInfo::create(c_desc, a_desc, b_desc, group_sizes_desc); + CHECK_RESULT(result); + + auto info = result.take(); + + *desc_ptr = new Descriptor( + info, 0, + new Opaque{handle->internal()}, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +namespace { + +struct McblasDtypes { + hpccDataType ab_type; + hpccDataType c_type; + hcblasComputeType_t compute_type; +}; + +infiniStatus_t resolveDtypes(infiniDtype_t dtype, McblasDtypes &out) { + switch (dtype) { + case INFINI_DTYPE_F16: + out.ab_type = out.c_type = HPCC_R_16F; + out.compute_type = HCBLAS_COMPUTE_32F; + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + out.ab_type = out.c_type = HPCC_R_16BF; + out.compute_type = HCBLAS_COMPUTE_32F; + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + out.ab_type = out.c_type = HPCC_R_32F; + out.compute_type = HCBLAS_COMPUTE_32F_FAST_TF32; + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +inline size_t element_size(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + case INFINI_DTYPE_BF16: + return 2; + case INFINI_DTYPE_F32: + return 4; + default: + return 0; + } +} + +} // namespace + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + const void *group_sizes, + const void *group_sizes_host, + float alpha, + float beta, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + McblasDtypes dtypes; + CHECK_STATUS(resolveDtypes(_info.dtype, dtypes)); + + auto hc_stream = static_cast(stream); + + // Resolve the host-side group sizes. When the caller already holds them on + // the host (MoE routing is bucketed on CPU) it passes `group_sizes_host`, + // letting us skip the per-call device->host copy + stream sync. On MetaX a + // pageable hcMemcpyAsync is effectively blocking, so eliminating this + // round-trip removes a hard sync from every grouped GEMM on the decode path. + std::vector sizes_host_vec; + const int32_t *sizes_host; + if (group_sizes_host != nullptr) { + sizes_host = reinterpret_cast(group_sizes_host); + } else { + sizes_host_vec.resize(_info.num_groups); + CHECK_METAX(hcMemcpyAsync( + sizes_host_vec.data(), group_sizes, + _info.num_groups * sizeof(int32_t), + hcMemcpyDeviceToHost, hc_stream)); + CHECK_METAX(hcStreamSynchronize(hc_stream)); + sizes_host = sizes_host_vec.data(); + } + + const size_t elem_size = element_size(_info.dtype); + auto a_bytes = reinterpret_cast(a); + auto b_bytes = reinterpret_cast(b); + auto c_bytes = reinterpret_cast(c); + + // Row-major C[rows_g, N] = alpha * A[rows_g, K] @ B[g][N, K]^T + beta * C. + // We compute the transpose `C^T = B @ A^T` in column-major space; the + // memory is identical. See NVIDIA backend for the full derivation. + CHECK_STATUS(_opaque->internal->useMcblas( + hc_stream, + [&](hcblasHandle_t handle) { + ptrdiff_t row_offset = 0; + for (size_t g = 0; g < _info.num_groups; ++g) { + int32_t rows = sizes_host[g]; + if (rows <= 0) { + continue; + } + const void *a_g = a_bytes + size_t(row_offset) * size_t(_info.a_row_stride) * elem_size; + const void *b_g = b_bytes + g * size_t(_info.b_group_stride) * elem_size; + void *c_g = c_bytes + size_t(row_offset) * size_t(_info.c_row_stride) * elem_size; + + CHECK_MCBLAS(hcblasGemmEx( + handle, + HCBLAS_OP_T, // op on B[g] view -> N x K + HCBLAS_OP_N, // op on A view -> K x rows_g + static_cast(_info.n), // m_cb + static_cast(rows), // n_cb + static_cast(_info.k), // k_cb + &alpha, + b_g, dtypes.ab_type, + static_cast(_info.b_row_stride), + a_g, dtypes.ab_type, + static_cast(_info.a_row_stride), + &beta, + c_g, dtypes.c_type, + static_cast(_info.c_row_stride), + dtypes.compute_type, + HCBLAS_GEMM_DEFAULT_TENSOR_OP)); + + row_offset += rows; + } + return INFINI_STATUS_SUCCESS; + })); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::grouped_gemm::metax diff --git a/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.h b/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.h new file mode 100644 index 000000000..f14ccaf0c --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/metax/grouped_gemm_metax.h @@ -0,0 +1,8 @@ +#ifndef __GROUPED_GEMM_METAX_H__ +#define __GROUPED_GEMM_METAX_H__ + +#include "../grouped_gemm.h" + +DESCRIPTOR(metax) + +#endif // __GROUPED_GEMM_METAX_H__ diff --git a/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cu b/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cu new file mode 100644 index 000000000..729b852b6 --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cu @@ -0,0 +1,199 @@ +#include "../../../devices/nvidia/nvidia_handle.cuh" +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include "grouped_gemm_nvidia.cuh" + +#include +#include + +namespace op::grouped_gemm::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle_, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc) { + + auto handle = reinterpret_cast(handle_); + + CHECK_DTYPE(c_desc->dtype(), INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_BF16); + + auto result = GroupedGemmInfo::create(c_desc, a_desc, b_desc, group_sizes_desc); + CHECK_RESULT(result); + + auto info = result.take(); + // Workspace holds the host-side copy of `group_sizes` so calculate() can + // walk each group's slab without round-tripping through the user's pointer. + size_t workspace_size = info.num_groups * sizeof(int32_t); + + *desc_ptr = new Descriptor( + info, workspace_size, + new Opaque{handle->internal()}, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +namespace { + +struct CublasDtypes { + cudaDataType ab_type; + cudaDataType c_type; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + cudaDataType compute_type; +#else + cublasComputeType_t compute_type; +#endif +}; + +infiniStatus_t resolveDtypes(infiniDtype_t dtype, CublasDtypes &out) { + switch (dtype) { + case INFINI_DTYPE_F16: + out.ab_type = out.c_type = CUDA_R_16F; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + out.compute_type = CUDA_R_32F; +#else + out.compute_type = CUBLAS_COMPUTE_32F; +#endif + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_BF16: + out.ab_type = out.c_type = CUDA_R_16BF; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + out.compute_type = CUDA_R_32F; +#else + out.compute_type = CUBLAS_COMPUTE_32F; +#endif + return INFINI_STATUS_SUCCESS; + case INFINI_DTYPE_F32: + out.ab_type = out.c_type = CUDA_R_32F; +#if defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) + out.compute_type = CUDA_R_32F; +#else + out.compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; +#endif + return INFINI_STATUS_SUCCESS; + default: + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } +} + +inline size_t element_size(infiniDtype_t dtype) { + switch (dtype) { + case INFINI_DTYPE_F16: + case INFINI_DTYPE_BF16: + return 2; + case INFINI_DTYPE_F32: + return 4; + default: + return 0; + } +} + +} // namespace + +infiniStatus_t Descriptor::calculate( + void *workspace, + size_t workspace_size, + void *c, + const void *a, + const void *b, + const void *group_sizes, + const void *group_sizes_host, + float alpha, + float beta, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + CublasDtypes dtypes; + CHECK_STATUS(resolveDtypes(_info.dtype, dtypes)); + + auto cuda_stream = static_cast(stream); + + // Resolve the host-side group sizes. When the caller already holds them on + // the host (e.g. MoE routing computed on CPU) it passes `group_sizes_host`, + // letting us skip the per-call device->host copy + stream sync -- the only + // host/device round-trip on this hot path. Otherwise we fall back to syncing + // the small (num_groups * 4 bytes) array down through the workspace. + // (The true batched path, cublasGemmGroupedBatched, needs CUDA 12.4+ and is + // left for a future revision.) + const int32_t *sizes_host; + if (group_sizes_host != nullptr) { + sizes_host = reinterpret_cast(group_sizes_host); + } else { + auto sizes_ws = reinterpret_cast(workspace); + CHECK_CUDA(cudaMemcpyAsync( + sizes_ws, group_sizes, + _info.num_groups * sizeof(int32_t), + cudaMemcpyDeviceToHost, cuda_stream)); + CHECK_CUDA(cudaStreamSynchronize(cuda_stream)); + sizes_host = sizes_ws; + } + + const size_t elem_size = element_size(_info.dtype); + auto a_bytes = reinterpret_cast(a); + auto b_bytes = reinterpret_cast(b); + auto c_bytes = reinterpret_cast(c); + + // C[rows_g, N] = alpha * A[rows_g, K] @ B[g][N, K]^T + beta * C[rows_g, N] + // (all row-major). cuBLAS is column-major, so we ask it to compute the + // transpose `C^T[N, rows_g] = B[g] @ A^T` instead — same memory. + // + // With cuBLAS conventions, when we view row-major data as column-major it + // appears mathematically transposed. So for `op(A_cb) @ op(B_cb)`: + // cuBLAS A_cb = B[g] memory, op = OP_T -> recovers the N x K math matrix + // cuBLAS B_cb = A memory, op = OP_N -> recovers the K x rows_g matrix + // cuBLAS C_cb = C memory, no transpose -> N x rows_g output + // The leading dimensions equal the row stride of the row-major tensor + // because that is the stride between successive columns once viewed as + // column-major. + CHECK_STATUS(_opaque->internal->useCublas( + cuda_stream, + [&](cublasHandle_t handle) { + ptrdiff_t row_offset = 0; + for (size_t g = 0; g < _info.num_groups; ++g) { + int32_t rows = sizes_host[g]; + if (rows <= 0) { + continue; + } + const void *a_g = a_bytes + size_t(row_offset) * size_t(_info.a_row_stride) * elem_size; + const void *b_g = b_bytes + g * size_t(_info.b_group_stride) * elem_size; + void *c_g = c_bytes + size_t(row_offset) * size_t(_info.c_row_stride) * elem_size; + + CHECK_CUBLAS(cublasGemmEx( + handle, + CUBLAS_OP_T, // op on B[g] view -> N x K + CUBLAS_OP_N, // op on A view -> K x rows_g + static_cast(_info.n), // m_cb (rows of C^T) + static_cast(rows), // n_cb (cols of C^T) + static_cast(_info.k), // k_cb + &alpha, + b_g, dtypes.ab_type, + static_cast(_info.b_row_stride), // lda = stride between cols of B view + a_g, dtypes.ab_type, + static_cast(_info.a_row_stride), // ldb = stride between cols of A view + &beta, + c_g, dtypes.c_type, + static_cast(_info.c_row_stride), // ldc = stride between cols of C view + dtypes.compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + + row_offset += rows; + } + return INFINI_STATUS_SUCCESS; + })); + + return INFINI_STATUS_SUCCESS; +} + +} // namespace op::grouped_gemm::nvidia diff --git a/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cuh b/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cuh new file mode 100644 index 000000000..17c3d115c --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/nvidia/grouped_gemm_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __GROUPED_GEMM_NVIDIA_CUH__ +#define __GROUPED_GEMM_NVIDIA_CUH__ + +#include "../grouped_gemm.h" + +DESCRIPTOR(nvidia) + +#endif // __GROUPED_GEMM_NVIDIA_CUH__ diff --git a/src/infiniop/ops/grouped_gemm/operator.cc b/src/infiniop/ops/grouped_gemm/operator.cc new file mode 100644 index 000000000..61c30470d --- /dev/null +++ b/src/infiniop/ops/grouped_gemm/operator.cc @@ -0,0 +1,180 @@ +#include "../../handle.h" +#include "../../operator.h" +#include "infiniop/ops/grouped_gemm.h" + +#ifdef ENABLE_CPU_API +#include "cpu/grouped_gemm_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/grouped_gemm_nvidia.cuh" +#endif +#ifdef ENABLE_METAX_API +#include "metax/grouped_gemm_metax.h" +#endif + +__INFINI_C infiniStatus_t infiniopCreateGroupedGemmDescriptor( + infiniopHandle_t handle, + infiniopGroupedGemmDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t c_desc, + infiniopTensorDescriptor_t a_desc, + infiniopTensorDescriptor_t b_desc, + infiniopTensorDescriptor_t group_sizes_desc) { + +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::grouped_gemm::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + c_desc, \ + a_desc, \ + b_desc, \ + group_sizes_desc) + + switch (handle->device) { + +#ifdef ENABLE_CPU_API + CREATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CREATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CREATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_ALI_API + CREATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CREATE(INFINI_DEVICE_METAX, metax); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetGroupedGemmWorkspaceSize( + infiniopGroupedGemmDescriptor_t desc, + size_t *size) { + +#define GET(CASE, NAMESPACE) \ + case CASE: \ + *size = reinterpret_cast(desc)->workspaceSize(); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { + +#ifdef ENABLE_CPU_API + GET(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + GET(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + GET(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_ALI_API + GET(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + GET(INFINI_DEVICE_METAX, metax); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef GET +} + +__INFINI_C infiniStatus_t infiniopGroupedGemm( + infiniopGroupedGemmDescriptor_t desc, + void *workspace, size_t workspace_size, + void *c, + const void *a, + const void *b, + const void *group_sizes, + const void *group_sizes_host, + float alpha, + float beta, + void *stream) { + +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, \ + c, a, b, group_sizes, group_sizes_host, \ + alpha, beta, \ + stream) + + switch (desc->device_type) { + +#ifdef ENABLE_CPU_API + CALCULATE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + CALCULATE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + CALCULATE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_ALI_API + CALCULATE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + CALCULATE(INFINI_DEVICE_METAX, metax); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyGroupedGemmDescriptor( + infiniopGroupedGemmDescriptor_t desc) { + +#define DELETE(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS; + + switch (desc->device_type) { + +#ifdef ENABLE_CPU_API + DELETE(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + DELETE(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DELETE(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_ALI_API + DELETE(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DELETE(INFINI_DEVICE_HYGON, nvidia); +#endif +#ifdef ENABLE_METAX_API + DELETE(INFINI_DEVICE_METAX, metax); +#endif + + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef DELETE +} diff --git a/test/infiniop/grouped_gemm.py b/test/infiniop/grouped_gemm.py new file mode 100644 index 000000000..fff1ded25 --- /dev/null +++ b/test/infiniop/grouped_gemm.py @@ -0,0 +1,176 @@ +import torch +import ctypes +from ctypes import c_uint64 +from libinfiniop import ( + LIBINFINIOP, + TestTensor, + get_test_devices, + check_error, + test_operator, + get_args, + debug, + get_tolerance, + profile_operation, + TestWorkspace, + InfiniDtype, + InfiniDtypeNames, + InfiniDeviceNames, + infiniopOperatorDescriptor_t, +) + +# ============================================================================== +# Configuration (internal use only) +# ============================================================================== +# (alpha, beta, group_sizes, K, N) +# a: [sum(group_sizes), K] b: [num_groups, N, K] c: [sum(group_sizes), N] +# +# The MoE shapes we care about for Qwen3-30B-A3B are around K=2048 N=768 E=128 +# with very skewed group_sizes; we keep a small smoke set and one realistic +# shape so CI doesn't blow up. +_TEST_CASES = [ + (1.0, 0.0, [1, 1, 1, 1], 16, 8), + (1.0, 0.0, [4, 0, 2, 7], 32, 24), + (1.0, 1.0, [3, 5], 64, 16), + (0.5, 0.0, [8, 8, 8, 8, 8, 8, 8, 8], 128, 64), + # Qwen3-MoE one-layer ish: 128 experts, top-8 routing, ~256 tokens -> ~16 tokens/expert avg + (1.0, 0.0, [16] * 128, 2048, 768), +] + +_TENSOR_DTYPES = [InfiniDtype.F16, InfiniDtype.BF16, InfiniDtype.F32] + +_TOLERANCE_MAP = { + InfiniDtype.F16: {"atol": 0, "rtol": 1e-2}, + InfiniDtype.F32: {"atol": 0, "rtol": 1e-3}, + InfiniDtype.BF16: {"atol": 0, "rtol": 5e-2}, +} + +DEBUG = False +PROFILE = False +NUM_PRERUN = 10 +NUM_ITERATIONS = 100 + + +def torch_grouped_gemm(c_inout, a, b, group_sizes, alpha, beta): + """Reference: per-group `linear` against the matching B slab.""" + out = torch.zeros_like(c_inout) + offset = 0 + for g, rows in enumerate(group_sizes): + if rows == 0: + continue + a_g = a[offset : offset + rows] + b_g = b[g] # [N, K] + out[offset : offset + rows] = a_g @ b_g.transpose(0, 1) + offset += rows + return alpha * out + beta * c_inout + + +def test( + handle, + device, + alpha, + beta, + group_sizes, + k, + n, + dtype=InfiniDtype.F16, + sync=None, +): + num_groups = len(group_sizes) + m_total = sum(group_sizes) + + print( + f"Testing GroupedGemm on {InfiniDeviceNames[device]} with" + f" alpha:{alpha}, beta:{beta}, groups:{num_groups}," + f" M_total:{m_total}, K:{k}, N:{n}, dtype:{InfiniDtypeNames[dtype]}" + ) + + a = TestTensor((m_total, k), None, dtype, device) + b = TestTensor((num_groups, n, k), None, dtype, device) + c = TestTensor((m_total, n), None, dtype, device, mode="ones") + ans_t = torch_grouped_gemm( + c.torch_tensor(), a.torch_tensor(), b.torch_tensor(), group_sizes, alpha, beta + ) + + # Pass the set tensor's own stride: TestTensor's "manual" mode asserts + # torch_strides == set_tensor.stride(), and strides=None makes torch_strides + # None (mirrors topksoftmax.py, which passes data.stride()). + group_sizes_data = torch.tensor(group_sizes, dtype=torch.int32) + group_sizes_tensor = TestTensor( + (num_groups,), + group_sizes_data.stride(), + InfiniDtype.I32, + device, + mode="manual", + set_tensor=group_sizes_data, + ) + + if sync is not None: + sync() + + descriptor = infiniopOperatorDescriptor_t() + check_error( + LIBINFINIOP.infiniopCreateGroupedGemmDescriptor( + handle, + ctypes.byref(descriptor), + c.descriptor, + a.descriptor, + b.descriptor, + group_sizes_tensor.descriptor, + ) + ) + + # Invalidate shapes/strides so the kernel can't peek at them through the desc. + for tensor in [a, b, c, group_sizes_tensor]: + tensor.destroy_desc() + + workspace_size = c_uint64(0) + check_error( + LIBINFINIOP.infiniopGetGroupedGemmWorkspaceSize( + descriptor, ctypes.byref(workspace_size) + ) + ) + workspace = TestWorkspace(workspace_size.value, device) + + def lib_run(): + check_error( + LIBINFINIOP.infiniopGroupedGemm( + descriptor, + workspace.data(), + workspace_size.value, + c.data(), + a.data(), + b.data(), + group_sizes_tensor.data(), + None, # group_sizes_host: exercise the device-sync fallback path + ctypes.c_float(alpha), + ctypes.c_float(beta), + None, + ) + ) + + lib_run() + + atol, rtol = get_tolerance(_TOLERANCE_MAP, dtype) + if DEBUG: + debug(c.actual_tensor(), ans_t, atol=atol, rtol=rtol) + assert torch.allclose(c.actual_tensor(), ans_t, atol=atol, rtol=rtol) + + if PROFILE: + # fmt: off + profile_operation(" lib", lambda: lib_run(), device, NUM_PRERUN, NUM_ITERATIONS) + # fmt: on + check_error(LIBINFINIOP.infiniopDestroyGroupedGemmDescriptor(descriptor)) + + +if __name__ == "__main__": + args = get_args() + + DEBUG = args.debug + PROFILE = args.profile + NUM_PRERUN = args.num_prerun + NUM_ITERATIONS = args.num_iterations + + for device in get_test_devices(args): + test_operator(device, test, _TEST_CASES, _TENSOR_DTYPES) + + print("\033[92mTest passed!\033[0m")