diff --git a/include/infinicore/ops.hpp b/include/infinicore/ops.hpp index 5e93e1457..bd42acc01 100644 --- a/include/infinicore/ops.hpp +++ b/include/infinicore/ops.hpp @@ -73,6 +73,7 @@ #include "ops/recurrent_gated_delta_rule.hpp" #include "ops/relu.hpp" #include "ops/rms_norm.hpp" +#include "ops/rms_norm_rope.hpp" #include "ops/rope.hpp" #include "ops/rot.hpp" #include "ops/rotary_embedding.hpp" diff --git a/include/infinicore/ops/rms_norm_rope.hpp b/include/infinicore/ops/rms_norm_rope.hpp new file mode 100644 index 000000000..f2f28f602 --- /dev/null +++ b/include/infinicore/ops/rms_norm_rope.hpp @@ -0,0 +1,23 @@ +#pragma once + +#include "../device.hpp" +#include "../graph/graph.hpp" +#include "../nn/rope.hpp" +#include "../tensor.hpp" +#include "common/op.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_CLASS(RMSNormRoPE, Tensor, const Tensor &, const Tensor &, const Tensor &, const Tensor &, float, infinicore::nn::RoPE::Algo); + +// Internal: fused per-head RMSNorm + RoPE (full rotary), in-place on x +// x: [num_tokens, num_heads, head_dim] +void rms_norm_rope_(Tensor x, + const Tensor &weight, + const Tensor &pos_ids, + const Tensor &sin_table, + const Tensor &cos_table, + float epsilon, + infinicore::nn::RoPE::Algo algo); + +} // namespace infinicore::op diff --git a/include/infiniop.h b/include/infiniop.h index 9f632e27f..1b9610945 100644 --- a/include/infiniop.h +++ b/include/infiniop.h @@ -128,6 +128,7 @@ #include "infiniop/ops/recurrent_gated_delta_rule.h" #include "infiniop/ops/relu.h" #include "infiniop/ops/rms_norm.h" +#include "infiniop/ops/rms_norm_rope.h" #include "infiniop/ops/rope.h" #include "infiniop/ops/rot.h" #include "infiniop/ops/rotg.h" diff --git a/include/infiniop/ops/rms_norm_rope.h b/include/infiniop/ops/rms_norm_rope.h new file mode 100644 index 000000000..775c85736 --- /dev/null +++ b/include/infiniop/ops/rms_norm_rope.h @@ -0,0 +1,35 @@ +#ifndef __INFINIOP_RMS_NORM_ROPE_API_H__ +#define __INFINIOP_RMS_NORM_ROPE_API_H__ + +#include "../operator_descriptor.h" +#include "rope.h" + +typedef struct InfiniopDescriptor *infiniopRMSNormRoPEDescriptor_t; + +__INFINI_C __export infiniStatus_t infiniopCreateRMSNormRoPEDescriptor( + infiniopHandle_t handle, + infiniopRMSNormRoPEDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t pos_ids_desc, + infiniopTensorDescriptor_t sin_table_desc, + infiniopTensorDescriptor_t cos_table_desc, + float epsilon, + infiniopRoPEAlgo_t algo); + +__INFINI_C __export infiniStatus_t infiniopGetRMSNormRoPEWorkspaceSize(infiniopRMSNormRoPEDescriptor_t desc, size_t *size); + +__INFINI_C __export infiniStatus_t infiniopRMSNormRoPE( + infiniopRMSNormRoPEDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void const *weight, + void const *pos_ids, + void const *sin_table, + void const *cos_table, + void *stream); + +__INFINI_C __export infiniStatus_t infiniopDestroyRMSNormRoPEDescriptor(infiniopRMSNormRoPEDescriptor_t desc); + +#endif diff --git a/src/infinicore/ops/rms_norm_rope/rms_norm_rope.cc b/src/infinicore/ops/rms_norm_rope/rms_norm_rope.cc new file mode 100644 index 000000000..c10976f8c --- /dev/null +++ b/src/infinicore/ops/rms_norm_rope/rms_norm_rope.cc @@ -0,0 +1,39 @@ +#include "infinicore/ops/rms_norm_rope.hpp" +#include "../../utils.hpp" + +namespace infinicore::op { + +INFINICORE_GRAPH_OP_DISPATCHERS_IMPL(RMSNormRoPE); + +RMSNormRoPE::RMSNormRoPE(Tensor x, + const Tensor &weight, + const Tensor &pos_ids, + const Tensor &sin_table, + const Tensor &cos_table, + float epsilon, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_ASSERT_TENSORS_SAME_DEVICE(x, weight, pos_ids, sin_table, cos_table); + INFINICORE_GRAPH_OP_DISPATCH(x->device().getType(), x, weight, pos_ids, sin_table, cos_table, epsilon, algo); +} + +void RMSNormRoPE::execute(Tensor x, + const Tensor &weight, + const Tensor &pos_ids, + const Tensor &sin_table, + const Tensor &cos_table, + float epsilon, + infinicore::nn::RoPE::Algo algo) { + INFINICORE_GRAPH_OP_RECORD_OR_RUN(RMSNormRoPE, x, weight, pos_ids, sin_table, cos_table, epsilon, algo); +} + +void rms_norm_rope_(Tensor x, + const Tensor &weight, + const Tensor &pos_ids, + const Tensor &sin_table, + const Tensor &cos_table, + float epsilon, + infinicore::nn::RoPE::Algo algo) { + RMSNormRoPE::execute(x, weight, pos_ids, sin_table, cos_table, epsilon, algo); +} + +} // namespace infinicore::op diff --git a/src/infinicore/ops/rms_norm_rope/rms_norm_rope_infiniop.cc b/src/infinicore/ops/rms_norm_rope/rms_norm_rope_infiniop.cc new file mode 100644 index 000000000..dbef47a22 --- /dev/null +++ b/src/infinicore/ops/rms_norm_rope/rms_norm_rope_infiniop.cc @@ -0,0 +1,85 @@ +#include "infinicore/ops/rms_norm_rope.hpp" + +#include "../infiniop_impl.hpp" + +#include + +namespace infinicore::op::rms_norm_rope_impl::infiniop { + +INFINIOP_CACHABLE_DESCRIPTOR(Descriptor, RMSNormRoPE, 100); + +struct PlannedMeta { + std::shared_ptr descriptor; + graph::GraphTensor workspace; + graph::GraphTensor x; + graph::GraphTensor weight; + graph::GraphTensor pos; + graph::GraphTensor sin; + graph::GraphTensor cos; +}; + +static infiniopRoPEAlgo_t to_infiniop_algo(infinicore::nn::RoPE::Algo algo) { + switch (algo) { + case infinicore::nn::RoPE::Algo::GPT_J: + return INFINIOP_ROPE_ALGO_GPT_J; + case infinicore::nn::RoPE::Algo::GPT_NEOX: + return INFINIOP_ROPE_ALGO_GPT_NEOX; + default: + throw std::runtime_error("Unsupported RoPE algorithm"); + } +} + +void *plan(Tensor x, + const Tensor &weight, + const Tensor &pos, + const Tensor &sin, + const Tensor &cos, + float epsilon, + infinicore::nn::RoPE::Algo algo) { + auto infiniop_algo = to_infiniop_algo(algo); + size_t key = hash_combine(x, weight, pos, sin, cos, epsilon, static_cast(infiniop_algo)); + + INFINIOP_CACHABLE_DESCRIPTOR_GET_OR_CREATE( + Descriptor, descriptor, RMSNormRoPE, key, x->desc(), + weight->desc(), + pos->desc(), + sin->desc(), + cos->desc(), + epsilon, + infiniop_algo); + + INFINIOP_WORKSPACE_TENSOR(workspace, RMSNormRoPE, descriptor); + return new PlannedMeta{ + descriptor, + graph::GraphTensor(workspace), + graph::GraphTensor(x), + graph::GraphTensor(weight), + graph::GraphTensor(pos), + graph::GraphTensor(sin), + graph::GraphTensor(cos)}; +} + +void run(void *planned_meta) { + auto *p = reinterpret_cast(planned_meta); + + INFINICORE_CHECK_ERROR( + infiniopRMSNormRoPE( + p->descriptor->desc, + p->workspace->data(), + p->workspace->numel(), + p->x->data(), + p->weight->data(), + p->pos->data(), + p->sin->data(), + p->cos->data(), + context::getStream())); +} + +void cleanup(void **planned_meta_ptr) { + delete *reinterpret_cast(planned_meta_ptr); + *planned_meta_ptr = nullptr; +} + +INFINICORE_GRAPH_OP_REGISTER_ALLDEVICE(RMSNormRoPE, &plan, &run, &cleanup); + +} // namespace infinicore::op::rms_norm_rope_impl::infiniop diff --git a/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.cc b/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.cc new file mode 100644 index 000000000..6a4bf0779 --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.cc @@ -0,0 +1,135 @@ +#include "rms_norm_rope_cpu.h" +#include "../../../devices/cpu/common_cpu.h" + +namespace op::rms_norm_rope::cpu { + +Descriptor::~Descriptor() {} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + infiniopTensorDescriptor_t pos_desc, + infiniopTensorDescriptor_t sin_desc, + infiniopTensorDescriptor_t cos_desc, + float epsilon, + infiniopRoPEAlgo_t algo) { + auto result = RMSNormRoPEInfo::create(x_desc, w_desc, pos_desc, sin_desc, cos_desc, epsilon, algo); + CHECK_RESULT(result); + *desc_ptr = new Descriptor(nullptr, result.take(), 0, handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +template +infiniStatus_t rms_norm_rope(const RMSNormRoPEInfo *info, + Tdata *x, const Tweight *w, const Tindex *pos_ids, + const Tdata *sin_table, const Tdata *cos_table) { + const size_t num_tokens = info->num_tokens; + const size_t nhead = info->num_heads; + const size_t table_dim = info->table_dim; + const size_t dim = 2 * table_dim; + const ptrdiff_t total_blocks = static_cast(num_tokens * nhead); + +#pragma omp parallel for + for (ptrdiff_t block_idx = 0; block_idx < total_blocks; ++block_idx) { + const size_t tok = block_idx / nhead; // token index + const size_t h = block_idx % nhead; // head index + + Tdata *x_ptr = x + tok * info->x_stride_token + h * info->x_stride_head; + + const size_t pos_id = size_t(pos_ids[tok * info->pos_stride]); + const Tdata *sin_ptr = sin_table + pos_id * table_dim; + const Tdata *cos_ptr = cos_table + pos_id * table_dim; + + // [Reduce] sum of x^2 on the row + float ss = 0.f; + for (size_t k = 0; k < dim; k++) { + float v = utils::cast(x_ptr[k]); + ss += v * v; + } + + // 1 / (sqrt(sum/dim + eps)) + float rms = 1.f / std::sqrt(ss / (float)(dim) + info->epsilon); + + for (size_t i = 0; i < table_dim; i++) { + // Calculate positions based on algorithm + size_t pos0, pos1; + if (info->algo == infiniopRoPEAlgo_t::INFINIOP_ROPE_ALGO_GPT_J) { + // GPT-J style: interleaved pairs + pos0 = 2 * i; + pos1 = 2 * i + 1; + } else { + // GPT-NeoX style: first half and second half + pos0 = i; + pos1 = i + table_dim; + } + + // Round normalized values to the storage dtype first + // (mimicking the two-kernel pipeline), then rotate in fp32 + // with the same conversion and rounding order as the rope cpu kernel + float x0 = utils::cast(utils::cast(utils::cast(x_ptr[pos0]) * utils::cast(w[pos0]) * rms)); + float x1 = utils::cast(utils::cast(utils::cast(x_ptr[pos1]) * utils::cast(w[pos1]) * rms)); + float sin__ = utils::cast(sin_ptr[i]); + float cos__ = utils::cast(cos_ptr[i]); + + x_ptr[pos0] = utils::cast(x0 * cos__ - x1 * sin__); + x_ptr[pos1] = utils::cast(x0 * sin__ + x1 * cos__); + } + } + + return INFINI_STATUS_SUCCESS; +} + +#define CALCULATE(Tdata, Tweight, Tindex) \ + rms_norm_rope(&_info, (Tdata *)x, (const Tweight *)w, (const Tindex *)pos_ids, (const Tdata *)sin_table, (const Tdata *)cos_table) + +#define DISPATCH_POS(Tdata, Tweight) \ + if (_info.pos_type == INFINI_DTYPE_I32) { \ + return CALCULATE(Tdata, Tweight, int32_t); \ + } else if (_info.pos_type == INFINI_DTYPE_I64) { \ + return CALCULATE(Tdata, Tweight, int64_t); \ + } else { \ + return INFINI_STATUS_BAD_TENSOR_DTYPE; \ + } + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *x, const void *w, + const void *pos_ids, + const void *sin_table, + const void *cos_table, + void *stream) const { + if (_info.atype == INFINI_DTYPE_F16) { + if (_info.wtype == INFINI_DTYPE_F16) { + DISPATCH_POS(fp16_t, fp16_t); + } else if (_info.wtype == INFINI_DTYPE_BF16) { + DISPATCH_POS(fp16_t, bf16_t); + } else if (_info.wtype == INFINI_DTYPE_F32) { + DISPATCH_POS(fp16_t, float); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + } else if (_info.atype == INFINI_DTYPE_BF16) { + if (_info.wtype == INFINI_DTYPE_BF16) { + DISPATCH_POS(bf16_t, bf16_t); + } else if (_info.wtype == INFINI_DTYPE_F16) { + DISPATCH_POS(bf16_t, fp16_t); + } else if (_info.wtype == INFINI_DTYPE_F32) { + DISPATCH_POS(bf16_t, float); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + } else if (_info.atype == INFINI_DTYPE_F32) { + DISPATCH_POS(float, float); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + + return INFINI_STATUS_SUCCESS; +} + +#undef DISPATCH_POS +#undef CALCULATE + +} // namespace op::rms_norm_rope::cpu diff --git a/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.h b/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.h new file mode 100644 index 000000000..2c08d95e5 --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/cpu/rms_norm_rope_cpu.h @@ -0,0 +1,7 @@ +#ifndef __RMS_NORM_ROPE_CPU_H__ +#define __RMS_NORM_ROPE_CPU_H__ +#include "../rms_norm_rope.h" + +DESCRIPTOR(cpu) + +#endif diff --git a/src/infiniop/ops/rms_norm_rope/cuda/kernel.cuh b/src/infiniop/ops/rms_norm_rope/cuda/kernel.cuh new file mode 100644 index 000000000..7e63336c0 --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/cuda/kernel.cuh @@ -0,0 +1,122 @@ +#ifndef __RMS_NORM_ROPE_CUDA_KERNEL_CUH__ +#define __RMS_NORM_ROPE_CUDA_KERNEL_CUH__ + +#include + +// Fused per-head RMSNorm + RoPE, in-place on x. +// Each block takes care of one head in one token. +// Each thread deals with every BLOCK_SIZE-th rotation pair in the row; +// pair i covers elements (2i, 2i+1) for GPT-J and (i, i + table_dim) for GPT-NeoX, +// so the pairs partition the whole row for both algorithms. +template +__device__ void rmsNormRopeBlock( + Tdata *__restrict__ x, + ptrdiff_t stride_x_token, + ptrdiff_t stride_x_head, + const Tweight *__restrict__ w, + const Tindex *__restrict__ pos_ids, + ptrdiff_t pos_stride, + const Tangle *__restrict__ sin_table, + const Tangle *__restrict__ cos_table, + size_t table_dim, + float epsilon) { + + size_t token_idx = blockIdx.x; + size_t head_idx = blockIdx.y; + + auto x_ptr = x + token_idx * stride_x_token + head_idx * stride_x_head; + + size_t pos_id = size_t(pos_ids[token_idx * pos_stride]); + auto sin_ptr = sin_table + pos_id * table_dim; + auto cos_ptr = cos_table + pos_id * table_dim; + + // [Reduce] sum of x^2 on the row (fp32 accumulate, same as the rms_norm kernel) + float sum_squared = 0; + for (size_t i = threadIdx.x; i < table_dim; i += BLOCK_SIZE) { + size_t pos0 = IsGPTJ ? 2 * i : i; + size_t pos1 = IsGPTJ ? 2 * i + 1 : i + table_dim; + float x0 = float(x_ptr[pos0]); + float x1 = float(x_ptr[pos1]); + sum_squared += x0 * x0 + x1 * x1; + } + + // Block-reduce sum of squares + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + sum_squared = BlockReduce(temp_storage).Sum(sum_squared); + + // Thread_0 computes RMS=1/sqrt(ss/dim+epsilon) and stores in shared memory + __shared__ float rms; + if (threadIdx.x == 0) { + rms = rsqrtf(sum_squared / float(2 * table_dim) + epsilon); + } + __syncthreads(); + + // Normalize and round to the storage dtype first: n is bit-identical to the + // rms_norm output, i.e. exactly what the rope kernel would read in the + // two-kernel pipeline. Then rotate with the same per-branch arithmetic and + // rounding points as the rope kernel (Tangle == Tdata, so e.g. the bf16 + // operators round back to bf16 after every operation). + // Each pair is owned by a single thread, so reloading x here still reads + // the original values. + for (size_t i = threadIdx.x; i < table_dim; i += BLOCK_SIZE) { + size_t pos0 = IsGPTJ ? 2 * i : i; + size_t pos1 = IsGPTJ ? 2 * i + 1 : i + table_dim; + + Tdata n0 = Tdata(float(x_ptr[pos0]) * float(w[pos0]) * rms); + Tdata n1 = Tdata(float(x_ptr[pos1]) * float(w[pos1]) * rms); + Tangle sin__ = sin_ptr[i], + cos__ = cos_ptr[i]; + + if constexpr (IsGPTJ) { + if constexpr (std::is_same::value) { + // Same as the half2 path of the rope kernel: packed pair, + // per-op rounding in half arithmetic + half2 x_pair = half2(n0, n1); + Tangle y0 = x_pair.x * cos__ - x_pair.y * sin__, + y1 = x_pair.x * sin__ + x_pair.y * cos__; + half2 y_pair = half2(y0, y1); + x_ptr[pos0] = y_pair.x; + x_ptr[pos1] = y_pair.y; + } else if constexpr (std::is_same::value) { + // Same as the bfloat162 path of the rope kernel + cuda_bfloat162 x_pair = cuda_bfloat162(n0, n1); + Tangle x0 = __low2bfloat16(x_pair); + Tangle x1 = __high2bfloat16(x_pair); + Tangle y0 = x0 * cos__ - x1 * sin__; + Tangle y1 = x0 * sin__ + x1 * cos__; + cuda_bfloat162 y_pair = __floats2bfloat162_rn(y0, y1); + x_ptr[pos0] = __low2bfloat16(y_pair); + x_ptr[pos1] = __high2bfloat16(y_pair); + } else { + Tangle x0 = Tangle(n0), + x1 = Tangle(n1); + x_ptr[pos0] = Tdata(x0 * cos__ - x1 * sin__); + x_ptr[pos1] = Tdata(x0 * sin__ + x1 * cos__); + } + } else { + if constexpr (std::is_same::value) { + Tangle x0 = __half2float(n0); + Tangle x1 = __half2float(n1); + Tangle y0 = x0 * cos__ - x1 * sin__; + Tangle y1 = x0 * sin__ + x1 * cos__; + x_ptr[pos0] = __float2half(y0); + x_ptr[pos1] = __float2half(y1); + } else if constexpr (std::is_same::value) { + Tangle x0 = __bfloat162float(n0); + Tangle x1 = __bfloat162float(n1); + Tangle y0 = x0 * cos__ - x1 * sin__; + Tangle y1 = x0 * sin__ + x1 * cos__; + x_ptr[pos0] = __float2bfloat16(y0); + x_ptr[pos1] = __float2bfloat16(y1); + } else { + Tangle x0 = Tangle(n0), + x1 = Tangle(n1); + x_ptr[pos0] = x0 * cos__ - x1 * sin__; + x_ptr[pos1] = x0 * sin__ + x1 * cos__; + } + } + } +} + +#endif diff --git a/src/infiniop/ops/rms_norm_rope/info.h b/src/infiniop/ops/rms_norm_rope/info.h new file mode 100644 index 000000000..8c8e07eaf --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/info.h @@ -0,0 +1,110 @@ +#ifndef __RMS_NORM_ROPE_INFO_H__ +#define __RMS_NORM_ROPE_INFO_H__ + +#include "../../../utils.h" +#include "../../tensor.h" +#include "infiniop/ops/rope.h" + +namespace op::rms_norm_rope { + +class RMSNormRoPEInfo { + RMSNormRoPEInfo() = default; + +public: + infiniDtype_t atype; + infiniDtype_t wtype; + infiniDtype_t pos_type; + float epsilon; + size_t num_tokens, num_heads, head_dim, table_len, table_dim; + ptrdiff_t x_stride_token; + ptrdiff_t x_stride_head; + ptrdiff_t pos_stride; + infiniopRoPEAlgo_t algo; + + static utils::Result create( + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + infiniopTensorDescriptor_t pos_desc, + infiniopTensorDescriptor_t sin_desc, + infiniopTensorDescriptor_t cos_desc, + float epsilon, + infiniopRoPEAlgo_t algo) { + + CHECK_OR_RETURN( + x_desc != nullptr && w_desc != nullptr && pos_desc != nullptr && sin_desc != nullptr && cos_desc != nullptr, + INFINI_STATUS_NULL_POINTER); + CHECK_OR_RETURN(algo < infiniopRoPEAlgo_t::INFINIOP_ROPE_ALGO_COUNT, INFINI_STATUS_BAD_PARAM); + + const infiniDtype_t atype = x_desc->dtype(); + const infiniDtype_t wtype = w_desc->dtype(); + const infiniDtype_t pos_type = pos_desc->dtype(); + + if (atype == INFINI_DTYPE_F16 || atype == INFINI_DTYPE_BF16) { + // For half-precision types (FP16/BF16), weights can be the same half-precision type or FP32 + if (wtype != atype && wtype != INFINI_DTYPE_F32 && wtype != INFINI_DTYPE_BF16 && wtype != INFINI_DTYPE_F16) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + } else if (atype == INFINI_DTYPE_F32) { + // For FP32, activations and weights must be of the same type + if (wtype != INFINI_DTYPE_F32) { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + // Position IDs must be 32/64-bit integers + CHECK_DTYPE(pos_type, INFINI_DTYPE_I32, INFINI_DTYPE_I64); + // sin/cos tables must have the same dtype as x (same constraint as the rope op) + CHECK_OR_RETURN(atype == sin_desc->dtype() && atype == cos_desc->dtype(), + INFINI_STATUS_BAD_TENSOR_DTYPE); + + // x: [num_tokens, num_heads, head_dim], weight: [head_dim], pos_ids: [num_tokens] + CHECK_OR_RETURN(x_desc->ndim() == 3 && w_desc->ndim() == 1 && pos_desc->ndim() == 1, + INFINI_STATUS_BAD_TENSOR_SHAPE); + CHECK_OR_RETURN(sin_desc->ndim() == 2 && cos_desc->ndim() == 2, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + const size_t num_tokens = x_desc->dim(0); + const size_t num_heads = x_desc->dim(1); + const size_t head_dim = x_desc->dim(2); + + CHECK_OR_RETURN(w_desc->dim(0) == head_dim && pos_desc->dim(0) == num_tokens, + INFINI_STATUS_BAD_TENSOR_SHAPE); + + const auto table_len = sin_desc->dim(0); + const auto table_dim = sin_desc->dim(1); + CHECK_OR_RETURN(table_len == cos_desc->dim(0) && table_dim == cos_desc->dim(1), + INFINI_STATUS_BAD_TENSOR_SHAPE); + + // v1 only supports full rotary (rotary_dim == head_dim) + CHECK_OR_RETURN(head_dim == table_dim * 2, INFINI_STATUS_BAD_TENSOR_SHAPE); + + // Last dimension of x must be contiguous, weight must be contiguous + CHECK_OR_RETURN(x_desc->stride(2) == 1 && w_desc->stride(0) == 1, + INFINI_STATUS_BAD_TENSOR_STRIDES); + + // sin table and cos table must be totally contiguous + CHECK_OR_RETURN(sin_desc->isContiguous() && cos_desc->isContiguous(), + INFINI_STATUS_BAD_TENSOR_STRIDES); + + return utils::Result(RMSNormRoPEInfo{ + atype, + wtype, + pos_type, + epsilon, + num_tokens, + num_heads, + head_dim, + table_len, + table_dim, + x_desc->stride(0), + x_desc->stride(1), + pos_desc->stride(0), + algo, + }); + } +}; + +} // namespace op::rms_norm_rope + +#endif // __RMS_NORM_ROPE_INFO_H__ diff --git a/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cu b/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cu new file mode 100644 index 000000000..be58d498e --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cu @@ -0,0 +1,151 @@ +#include "../../../devices/nvidia/nvidia_common.cuh" +#include "rms_norm_rope_nvidia.cuh" + +#include "../../../devices/nvidia/nvidia_kernel_common.cuh" +#include + +#include "../cuda/kernel.cuh" + +template +INFINIOP_CUDA_KERNEL rmsNormRopeKernel( + Tdata *__restrict__ x, + ptrdiff_t stride_x_token, + ptrdiff_t stride_x_head, + const Tweight *__restrict__ w, + const Tindex *__restrict__ pos_ids, + ptrdiff_t pos_stride, + const Tangle *__restrict__ sin_table, + const Tangle *__restrict__ cos_table, + size_t table_dim, + float epsilon) { + rmsNormRopeBlock(x, stride_x_token, stride_x_head, w, pos_ids, pos_stride, sin_table, cos_table, table_dim, epsilon); +} + +namespace op::rms_norm_rope::nvidia { + +struct Descriptor::Opaque { + std::shared_ptr internal; +}; + +Descriptor::~Descriptor() { + delete _opaque; +} + +infiniStatus_t Descriptor::create( + infiniopHandle_t handle, + Descriptor **desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t w_desc, + infiniopTensorDescriptor_t pos_desc, + infiniopTensorDescriptor_t sin_desc, + infiniopTensorDescriptor_t cos_desc, + float epsilon, + infiniopRoPEAlgo_t algo) { + auto result = RMSNormRoPEInfo::create(x_desc, w_desc, pos_desc, sin_desc, cos_desc, epsilon, algo); + CHECK_RESULT(result); + auto info = result.take(); + + *desc_ptr = new Descriptor( + new Opaque{reinterpret_cast(handle)->internal()}, + std::move(info), + 0, + handle->device, handle->device_id); + return INFINI_STATUS_SUCCESS; +} + +// launch kernel with different data types +// sin/cos tables always have the same dtype as x (checked in RMSNormRoPEInfo), +// so they are interpreted as Tdata, same as the rope kernel launch +template +infiniStatus_t launchKernel( + uint32_t num_tokens, size_t nhead, size_t table_dim, + void *x, infiniDtype_t atype, ptrdiff_t stride_x_token, ptrdiff_t stride_x_head, + const void *w, infiniDtype_t wtype, + const void *pos_ids, infiniDtype_t pos_type, ptrdiff_t pos_stride, + const void *sin_table, const void *cos_table, + float epsilon, infiniopRoPEAlgo_t algo, + cudaStream_t cuda_stream) { + + dim3 grid_dim(num_tokens, uint32_t(nhead)); + +#define LAUNCH_KERNEL(Tdata, Tweight, Tindex, IsGPTJ) \ + rmsNormRopeKernel \ + <<>>( \ + reinterpret_cast(x), \ + stride_x_token, \ + stride_x_head, \ + reinterpret_cast(w), \ + reinterpret_cast(pos_ids), \ + pos_stride, \ + reinterpret_cast(sin_table), \ + reinterpret_cast(cos_table), \ + table_dim, \ + epsilon) + +#define DISPATCH_ALGO(Tdata, Tweight, Tindex) \ + if (algo == INFINIOP_ROPE_ALGO_GPT_J) { \ + LAUNCH_KERNEL(Tdata, Tweight, Tindex, true); \ + } else { \ + LAUNCH_KERNEL(Tdata, Tweight, Tindex, false); \ + } \ + return INFINI_STATUS_SUCCESS + +#define DISPATCH_POS(Tdata, Tweight) \ + if (pos_type == INFINI_DTYPE_I32) { \ + DISPATCH_ALGO(Tdata, Tweight, int32_t); \ + } else if (pos_type == INFINI_DTYPE_I64) { \ + DISPATCH_ALGO(Tdata, Tweight, int64_t); \ + } else { \ + return INFINI_STATUS_BAD_TENSOR_DTYPE; \ + } + + if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) { + DISPATCH_POS(half, half); + } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) { + DISPATCH_POS(half, __nv_bfloat16); + } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) { + DISPATCH_POS(half, float); + } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) { + DISPATCH_POS(__nv_bfloat16, __nv_bfloat16); + } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) { + DISPATCH_POS(__nv_bfloat16, half); + } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) { + DISPATCH_POS(__nv_bfloat16, float); + } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) { + DISPATCH_POS(float, float); + } else { + return INFINI_STATUS_BAD_TENSOR_DTYPE; + } + +#undef DISPATCH_POS +#undef DISPATCH_ALGO +#undef LAUNCH_KERNEL + + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t Descriptor::calculate( + void *workspace, size_t workspace_size, + void *x, const void *w, + const void *pos_ids, + const void *sin_table, + const void *cos_table, + void *stream) const { + + if (workspace_size < _workspace_size) { + return INFINI_STATUS_INSUFFICIENT_WORKSPACE; + } + + auto cuda_stream = reinterpret_cast(stream); + + // launch kernel with block size matching the number of rotation pairs + if (_info.table_dim <= 64) { + CHECK_STATUS(launchKernel<64>(uint32_t(_info.num_tokens), _info.num_heads, _info.table_dim, x, _info.atype, _info.x_stride_token, _info.x_stride_head, w, _info.wtype, pos_ids, _info.pos_type, _info.pos_stride, sin_table, cos_table, _info.epsilon, _info.algo, cuda_stream)); + } else if (_info.table_dim <= 128) { + CHECK_STATUS(launchKernel<128>(uint32_t(_info.num_tokens), _info.num_heads, _info.table_dim, x, _info.atype, _info.x_stride_token, _info.x_stride_head, w, _info.wtype, pos_ids, _info.pos_type, _info.pos_stride, sin_table, cos_table, _info.epsilon, _info.algo, cuda_stream)); + } else { + CHECK_STATUS(launchKernel<256>(uint32_t(_info.num_tokens), _info.num_heads, _info.table_dim, x, _info.atype, _info.x_stride_token, _info.x_stride_head, w, _info.wtype, pos_ids, _info.pos_type, _info.pos_stride, sin_table, cos_table, _info.epsilon, _info.algo, cuda_stream)); + } + return INFINI_STATUS_SUCCESS; +} +} // namespace op::rms_norm_rope::nvidia diff --git a/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cuh b/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cuh new file mode 100644 index 000000000..fe0351332 --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/nvidia/rms_norm_rope_nvidia.cuh @@ -0,0 +1,8 @@ +#ifndef __RMS_NORM_ROPE_NVIDIA_CUH__ +#define __RMS_NORM_ROPE_NVIDIA_CUH__ + +#include "../rms_norm_rope.h" + +DESCRIPTOR(nvidia) + +#endif // __RMS_NORM_ROPE_NVIDIA_CUH__ diff --git a/src/infiniop/ops/rms_norm_rope/operator.cc b/src/infiniop/ops/rms_norm_rope/operator.cc new file mode 100644 index 000000000..843991f01 --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/operator.cc @@ -0,0 +1,171 @@ +#include "../../operator.h" +#include "../../handle.h" +#include "infiniop/ops/rms_norm_rope.h" + +#ifdef ENABLE_CPU_API +#include "cpu/rms_norm_rope_cpu.h" +#endif +#if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ILUVATAR_API) || defined(ENABLE_QY_API) || defined(ENABLE_HYGON_API) || defined(ENABLE_ALI_API) +#include "nvidia/rms_norm_rope_nvidia.cuh" +#endif + +__INFINI_C infiniStatus_t infiniopCreateRMSNormRoPEDescriptor( + infiniopHandle_t handle, + infiniopRMSNormRoPEDescriptor_t *desc_ptr, + infiniopTensorDescriptor_t x_desc, + infiniopTensorDescriptor_t weight_desc, + infiniopTensorDescriptor_t pos_ids_desc, + infiniopTensorDescriptor_t sin_table_desc, + infiniopTensorDescriptor_t cos_table_desc, + float epsilon, + infiniopRoPEAlgo_t algo) { + +#define CREATE(CASE, NAMESPACE) \ + case CASE: \ + return op::rms_norm_rope::NAMESPACE::Descriptor::create( \ + handle, \ + reinterpret_cast(desc_ptr), \ + x_desc, \ + weight_desc, \ + pos_ids_desc, \ + sin_table_desc, \ + cos_table_desc, \ + epsilon, \ + algo) + + 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_QY_API + CREATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CREATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } + +#undef CREATE +} + +__INFINI_C infiniStatus_t infiniopGetRMSNormRoPEWorkspaceSize(infiniopRMSNormRoPEDescriptor_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_QY_API + GET(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + GET(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef GET + + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +__INFINI_C infiniStatus_t infiniopRMSNormRoPE( + infiniopRMSNormRoPEDescriptor_t desc, + void *workspace, + size_t workspace_size, + void *x, + void const *weight, + void const *pos_ids, + void const *sin_table, + void const *cos_table, + void *stream) { + +#define CALCULATE(CASE, NAMESPACE) \ + case CASE: \ + return reinterpret_cast(desc) \ + ->calculate(workspace, workspace_size, x, weight, pos_ids, sin_table, cos_table, 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_QY_API + CALCULATE(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + CALCULATE(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef CALCULATE +} + +__INFINI_C infiniStatus_t infiniopDestroyRMSNormRoPEDescriptor(infiniopRMSNormRoPEDescriptor_t desc) { + if (desc == nullptr) { + return INFINI_STATUS_SUCCESS; + } + +#define DESTROY(CASE, NAMESPACE) \ + case CASE: \ + delete reinterpret_cast(desc); \ + return INFINI_STATUS_SUCCESS + + switch (desc->device_type) { +#ifdef ENABLE_CPU_API + DESTROY(INFINI_DEVICE_CPU, cpu); +#endif +#ifdef ENABLE_NVIDIA_API + DESTROY(INFINI_DEVICE_NVIDIA, nvidia); +#endif +#ifdef ENABLE_ILUVATAR_API + DESTROY(INFINI_DEVICE_ILUVATAR, nvidia); +#endif +#ifdef ENABLE_ALI_API + DESTROY(INFINI_DEVICE_ALI, nvidia); +#endif +#ifdef ENABLE_QY_API + DESTROY(INFINI_DEVICE_QY, nvidia); +#endif +#ifdef ENABLE_HYGON_API + DESTROY(INFINI_DEVICE_HYGON, nvidia); +#endif + default: + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + } +#undef DESTROY +} diff --git a/src/infiniop/ops/rms_norm_rope/rms_norm_rope.h b/src/infiniop/ops/rms_norm_rope/rms_norm_rope.h new file mode 100644 index 000000000..0bd8b75bd --- /dev/null +++ b/src/infiniop/ops/rms_norm_rope/rms_norm_rope.h @@ -0,0 +1,54 @@ +#ifndef RMS_NORM_ROPE_H +#define RMS_NORM_ROPE_H + +#include "../../operator.h" +#include "info.h" + +#define DESCRIPTOR(NAMESPACE) \ + \ + namespace op::rms_norm_rope::NAMESPACE { \ + class Descriptor final : public InfiniopDescriptor { \ + struct Opaque; \ + Opaque *_opaque; \ + RMSNormRoPEInfo _info; \ + size_t _workspace_size; \ + \ + Descriptor( \ + Opaque *opaque, \ + RMSNormRoPEInfo info, \ + size_t workspace_size, \ + 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 x_desc, \ + infiniopTensorDescriptor_t w_desc, \ + infiniopTensorDescriptor_t pos_desc, \ + infiniopTensorDescriptor_t sin_desc, \ + infiniopTensorDescriptor_t cos_desc, \ + float epsilon, \ + infiniopRoPEAlgo_t algo); \ + \ + infiniStatus_t calculate( \ + void *workspace, size_t workspace_size, \ + void *x, \ + const void *w, \ + const void *pos_ids, \ + const void *sin_table, \ + const void *cos_table, \ + void *stream) const; \ + }; \ + } + +#endif // RMS_NORM_ROPE_H