diff --git a/scripts/torch_ops.yaml b/scripts/torch_ops.yaml index e0b7ece5f..3f79960a0 100644 --- a/scripts/torch_ops.yaml +++ b/scripts/torch_ops.yaml @@ -5,13 +5,11 @@ # generator emits one InfiniOps wrapper per overload, so this # file lists 500+ base names and produces 500+ wrappers. # -# To exclude an op, comment out its line. Ops whose hand-written -# `src/base/.h` signature does not match the ATen-derived one -# (currently `add`, `linear`, `matmul`, `mul` — they pre-date this -# codegen and use a different parameter shape) must stay excluded: -# the generator skips emitting their base, but would still emit a -# torch backend declaring `operator()` with the ATen signature, and -# that override would not compile against the hand-written base. +# To exclude an op, comment out its line. Existing base operators may remain +# on this list when their public overloads match the ATen schema; the generator +# then emits their Torch backend. `cat`, `linear`, and `matmul` use this path. +# Operators with handwritten Torch backends, such as `add`, stay excluded to +# avoid generating a duplicate backend. - abs - absolute @@ -68,6 +66,7 @@ - bitwise_xor - bmm - bucketize +- cat - ceil - chain_matmul - cholesky @@ -235,6 +234,7 @@ - linalg_tensorsolve - linalg_vecdot - linalg_vector_norm +- linear - linspace - log - log10 @@ -259,6 +259,7 @@ - lu_solve - lu_unpack - masked_select +- matmul - matrix_power - max - max_pool2d_with_indices diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..ea512b3de 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -419,7 +419,8 @@ set(INFINI_OPS_OPS "" CACHE STRING set(INFINI_OPS_SMOKE_BUILD OFF CACHE BOOL "Build only the smoke-test operator subset") set(_infini_ops_smoke_ops - add mul cast cat gemm matmul linear rms_norm swiglu causal_softmax abs clamp exp) + add mul cast cat gemm matmul linear rms_norm swiglu silu_and_mul + causal_softmax causal_softmax_infinilm abs clamp exp) set(_infini_ops_smoke_torch_ops abs clamp exp) if(INFINI_OPS_SMOKE_BUILD) diff --git a/src/base/add.h b/src/base/add.h index eb4b31f4a..985e6ed03 100644 --- a/src/base/add.h +++ b/src/base/add.h @@ -1,7 +1,10 @@ #ifndef INFINI_OPS_BASE_ADD_H_ #define INFINI_OPS_BASE_ADD_H_ -#include +#include +#include +#include +#include #include "operator.h" @@ -9,20 +12,22 @@ namespace infini::ops { class Add : public Operator { public: - Add(const Tensor input, const Tensor other, Tensor out) + Add(const Tensor input, const Tensor other, const double alpha, Tensor out) : ndim_{out.ndim()}, output_size_{out.numel()}, input_type_{input.dtype()}, other_type_{other.dtype()}, out_type_{out.dtype()}, - input_shape_{input.shape()}, - other_shape_{other.shape()}, + input_shape_{out.shape()}, + other_shape_{out.shape()}, out_shape_{out.shape()}, - input_strides_{input.strides()}, - other_strides_{other.strides()}, + input_strides_{BroadcastStrides(input, out)}, + other_strides_{BroadcastStrides(other, out)}, out_strides_{out.strides()}, - is_input_contiguous_{input.IsContiguous()}, - is_other_contiguous_{other.IsContiguous()}, + is_input_contiguous_{input.shape() == out.shape() && + input.IsContiguous()}, + is_other_contiguous_{other.shape() == out.shape() && + other.IsContiguous()}, is_out_contiguous_{out.IsContiguous()} { assert(!out.HasBroadcastDim() && "the output of `Add` should NOT have broadcasted dim!"); @@ -31,18 +36,89 @@ class Add : public Operator { assert(input_type_ == other_type_ && other_type_ == out_type_ && "operator `Add` requires all input and output Tensors to have the " "same dtype"); + if (input_type_ != DataType::kFloat16 && + input_type_ != DataType::kBFloat16 && + input_type_ != DataType::kFloat32 && + input_type_ != DataType::kFloat64) { + assert(alpha == static_cast(alpha) && + "operator `Add` requires integral `alpha` for integer tensors"); + } + ValidateBroadcast(input, other, out); } + Add(const Tensor input, const Tensor other, Tensor out) + : Add{input, other, 1.0, out} {} + virtual void operator()(const Tensor input, const Tensor other, - Tensor out) const = 0; + const double alpha, Tensor out) const = 0; + + void operator()(const Tensor input, const Tensor other, Tensor out) const { + (*this)(input, other, 1.0, out); + } template static auto MakeReturnValue(const TensorLike& input, const TensorLike& other) { - return TensorLike::Empty(input.shape(), input.dtype(), input.device()); + return MakeReturnValue(input, other, 1.0); + } + + template + static auto MakeReturnValue(const TensorLike& input, const TensorLike& other, + const double /*alpha*/) { + auto input_shape = input.shape(); + auto other_shape = other.shape(); + auto ndim = std::max(input_shape.size(), other_shape.size()); + std::vector out_shape(ndim, 1); + + for (std::size_t i = 0; i < ndim; ++i) { + auto input_dim = i < ndim - input_shape.size() + ? 1 + : input_shape[i + input_shape.size() - ndim]; + auto other_dim = i < ndim - other_shape.size() + ? 1 + : other_shape[i + other_shape.size() - ndim]; + assert((input_dim == other_dim || input_dim == 1 || other_dim == 1) && + "operator `Add` requires broadcast-compatible input shapes"); + out_shape[i] = std::max(input_dim, other_dim); + } + + return TensorLike::Empty(out_shape, input.dtype(), input.device()); } protected: + static Tensor::Strides BroadcastStrides(const Tensor input, + const Tensor out) { + assert(input.ndim() <= out.ndim() && + "operator `Add` input rank must not exceed output rank"); + Tensor::Strides strides(out.ndim(), 0); + auto offset = out.ndim() - input.ndim(); + + for (Tensor::Size i = 0; i < input.ndim(); ++i) { + auto out_dim = i + offset; + assert((input.size(i) == 1 || input.size(i) == out.size(out_dim)) && + "operator `Add` input shape is not broadcast-compatible with " + "output shape"); + strides[out_dim] = input.size(i) == 1 ? 0 : input.stride(i); + } + + return strides; + } + + static void ValidateBroadcast(const Tensor input, const Tensor other, + const Tensor out) { + for (Tensor::Size i = 0; i < out.ndim(); ++i) { + auto input_dim = i < out.ndim() - input.ndim() + ? 1 + : input.size(i + input.ndim() - out.ndim()); + auto other_dim = i < out.ndim() - other.ndim() + ? 1 + : other.size(i + other.ndim() - out.ndim()); + assert(out.size(i) == std::max(input_dim, other_dim) && + "operator `Add` output shape must equal the broadcasted input " + "shape"); + } + } + Tensor::Size ndim_{0}; Tensor::Size output_size_{0}; diff --git a/src/base/add_rms_norm.h b/src/base/add_rms_norm.h index 0241f7fce..5c8e4fa6a 100644 --- a/src/base/add_rms_norm.h +++ b/src/base/add_rms_norm.h @@ -9,7 +9,9 @@ namespace infini::ops { -// Fused add + RMSNorm aligned with vLLM `fused_add_rms_norm`. +// Legacy out-of-place fused add + RMSNorm interface. +/// \deprecated Use `FusedAddRmsNorm`. This interface will be removed in a +/// future release. class AddRmsNorm : public Operator { public: AddRmsNorm(const Tensor input, const Tensor residual, const Tensor weight, diff --git a/src/base/cat.h b/src/base/cat.h index dcb0ba587..b7f2c9f4c 100644 --- a/src/base/cat.h +++ b/src/base/cat.h @@ -9,22 +9,46 @@ namespace infini::ops { class Cat : public Operator { public: - Cat(const Tensor first_input, std::vector rest_inputs, int64_t dim, - Tensor out) - : input_count_{1 + rest_inputs.size()} { - assert(input_count_ >= 2 && "`Cat` requires at least 2 input tensors"); + Cat(const std::vector tensors, const int64_t dim, Tensor out) + : out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()}, + input_count_{tensors.size()} { + assert(!tensors.empty() && "`Cat` requires a non-empty tensor list"); auto ndim = static_cast(out.ndim()); - // Normalize negative dim (e.g. -1 means last dimension). dim_ = dim < 0 ? dim + ndim : dim; assert(dim_ >= 0 && dim_ < ndim && "`Cat` dim out of range"); + + Tensor::Size cat_size = 0; + for (const auto& tensor : tensors) { + assert(tensor.ndim() == out.ndim() && + "`Cat` requires all tensors to have the output rank"); + assert(tensor.dtype() == out.dtype() && + "`Cat` requires all tensors to have the output dtype"); + + for (Tensor::Size axis = 0; axis < out.ndim(); ++axis) { + if (axis != static_cast(dim_)) { + assert(tensor.size(axis) == out.size(axis) && + "`Cat` input dimensions must match outside `dim`"); + } + } + cat_size += tensor.size(dim_); + } + assert(cat_size == out.size(dim_) && + "`Cat` output size along `dim` must equal the input sum"); } - virtual void operator()(const Tensor first_input, - std::vector rest_inputs, int64_t dim, + virtual void operator()(const std::vector tensors, const int64_t dim, Tensor out) const = 0; protected: + Tensor::Shape out_shape_; + + Tensor::Strides out_strides_; + + DataType out_type_; + int64_t dim_; size_t input_count_; diff --git a/src/base/causal_softmax.h b/src/base/causal_softmax.h index b8393d829..e03cfb2df 100644 --- a/src/base/causal_softmax.h +++ b/src/base/causal_softmax.h @@ -9,6 +9,8 @@ namespace infini::ops { +/// \deprecated Use `CausalSoftmaxInfinilm`. This interface will be removed in +/// a future release. class CausalSoftmax : public Operator { public: CausalSoftmax(const Tensor input, Tensor out) diff --git a/src/base/causal_softmax_infinilm.h b/src/base/causal_softmax_infinilm.h new file mode 100644 index 000000000..31bb206d1 --- /dev/null +++ b/src/base/causal_softmax_infinilm.h @@ -0,0 +1,52 @@ +#ifndef INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_ + +#include +#include + +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +class CausalSoftmaxInfinilm : public Operator { + public: + CausalSoftmaxInfinilm(const Tensor input, Tensor out) + : dtype_{input.dtype()}, + ndim_{out.ndim()}, + batch_size_{ndim_ == 2 ? 1 : out.size(-3)}, + seq_len_{out.size(-2)}, + total_seq_len_{out.size(-1)}, + input_strides_{input.strides()}, + out_strides_{out.strides()} { + assert(input.shape() == out.shape() && + "`CausalSoftmaxInfinilm` requires `input` and `out` same shape"); + assert(input.dtype() == out.dtype() && + "`CausalSoftmaxInfinilm` requires `input` and `out` same dtype"); + assert((ndim_ == 2 || ndim_ == 3) && + "`CausalSoftmaxInfinilm` requires 2D or 3D tensor"); + assert(seq_len_ <= total_seq_len_ && + "`CausalSoftmaxInfinilm` requires shape[-2] <= shape[-1]"); + } + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + protected: + const DataType dtype_; + + Tensor::Size ndim_{0}; + + Tensor::Size batch_size_{0}; + + Tensor::Size seq_len_{0}; + + Tensor::Size total_seq_len_{0}; + + Tensor::Strides input_strides_; + + Tensor::Strides out_strides_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/embedding.h b/src/base/embedding.h index 67126a5bd..44c9f425d 100644 --- a/src/base/embedding.h +++ b/src/base/embedding.h @@ -9,36 +9,40 @@ namespace infini::ops { -// Aligned with InfiniCore and `torch.nn.functional.embedding`. class Embedding : public Operator { public: - Embedding(const Tensor input, const Tensor weight, Tensor out) - : input_shape_{input.shape()}, + Embedding(const Tensor weight, const Tensor indices, + const int64_t padding_idx, const bool scale_grad_by_freq, + const bool sparse, Tensor out) + : indices_shape_{indices.shape()}, weight_shape_{weight.shape()}, out_shape_{out.shape()}, - input_strides_{input.strides()}, + indices_strides_{indices.strides()}, weight_strides_{weight.strides()}, out_strides_{out.strides()}, - input_dtype_{input.dtype()}, + indices_dtype_{indices.dtype()}, weight_dtype_{weight.dtype()}, out_dtype_{out.dtype()}, - num_indices_{NumIndices(input_shape_)}, + num_indices_{NumIndices(indices_shape_)}, vocab_size_{weight.size(0)}, - embedding_dim_{weight.size(1)} { + embedding_dim_{weight.size(1)}, + padding_idx_{padding_idx}, + scale_grad_by_freq_{scale_grad_by_freq}, + sparse_{sparse} { assert(weight.ndim() == 2 && "`Embedding` requires 2D `weight`"); - assert(out.ndim() == input.ndim() + 1 && - "`Embedding` output rank must be input rank + 1"); + assert(out.ndim() == indices.ndim() + 1 && + "`Embedding` output rank must be indices rank + 1"); - for (Tensor::Size i = 0; i < input.ndim(); ++i) { - assert(out.size(i) == input.size(i) && - "`Embedding` output shape must match `input` shape on non-last " + for (Tensor::Size i = 0; i < indices.ndim(); ++i) { + assert(out.size(i) == indices.size(i) && + "`Embedding` output shape must match `indices` on non-last " "dims"); } assert(out.size(-1) == embedding_dim_ && "`Embedding` output last dim must equal `weight` embedding dim"); - assert((input_dtype_ == DataType::kInt32 || - input_dtype_ == DataType::kInt64) && + assert((indices_dtype_ == DataType::kInt32 || + indices_dtype_ == DataType::kInt64) && "`Embedding` supports int32 and int64 indices only"); assert((weight_dtype_ == DataType::kFloat32 || weight_dtype_ == DataType::kFloat16 || @@ -46,11 +50,35 @@ class Embedding : public Operator { "`Embedding` supports float32, float16, and bfloat16 weights only"); assert(out_dtype_ == weight_dtype_ && "`Embedding` output dtype must match `weight` dtype"); + assert(padding_idx_ >= -static_cast(vocab_size_) && + padding_idx_ < static_cast(vocab_size_) && + "`Embedding` padding_idx must be within the weight rows"); } - virtual void operator()(const Tensor input, const Tensor weight, + Embedding(const Tensor weight, const Tensor indices, Tensor out) + : Embedding{weight, indices, -1, false, false, out} {} + + virtual void operator()(const Tensor weight, const Tensor indices, + const int64_t padding_idx, + const bool scale_grad_by_freq, const bool sparse, Tensor out) const = 0; + void operator()(const Tensor weight, const Tensor indices, Tensor out) const { + (*this)(weight, indices, -1, false, false, out); + } + + template + static auto MakeReturnValue(const TensorLike& weight, + const TensorLike& indices, + const int64_t /*padding_idx*/ = -1, + const bool /*scale_grad_by_freq*/ = false, + const bool /*sparse*/ = false) { + auto out_shape = indices.shape(); + out_shape.push_back(weight.size(1)); + + return TensorLike::Empty(out_shape, weight.dtype(), weight.device()); + } + protected: static Tensor::Size NumIndices(const Tensor::Shape& input_shape) { Tensor::Size num_indices = 1; @@ -62,19 +90,19 @@ class Embedding : public Operator { return num_indices; } - Tensor::Shape input_shape_; + Tensor::Shape indices_shape_; Tensor::Shape weight_shape_; Tensor::Shape out_shape_; - Tensor::Strides input_strides_; + Tensor::Strides indices_strides_; Tensor::Strides weight_strides_; Tensor::Strides out_strides_; - DataType input_dtype_; + DataType indices_dtype_; DataType weight_dtype_; @@ -85,6 +113,12 @@ class Embedding : public Operator { Tensor::Size vocab_size_{0}; Tensor::Size embedding_dim_{0}; + + int64_t padding_idx_{0}; + + bool scale_grad_by_freq_{false}; + + bool sparse_{false}; }; } // namespace infini::ops diff --git a/src/base/flash_attention.h b/src/base/flash_attention.h index e5952b514..9674ea789 100644 --- a/src/base/flash_attention.h +++ b/src/base/flash_attention.h @@ -9,14 +9,12 @@ namespace infini::ops { -// Fused multi-head / grouped-query attention. -// -// Interface follows vLLM v1 `AttentionImpl.forward()`: -// `vllm.v1.attention.backends.abstract.AttentionImpl` -// +// Legacy fused multi-head / grouped-query attention interface. // Layout: `query` / `key` / `value` are `[T, N, D]` (TND). // Prefill uses `cu_seqlens_q` / `cu_seqlens_kv` for variable-length packing. // Decode uses `block_table` for paged KV cache lookup. +/// \deprecated Use `ScaledDotProductAttention` for standard attention +/// semantics. This interface will be removed in a future release. class FlashAttention : public Operator { public: FlashAttention(const Tensor query, const Tensor key, const Tensor value, diff --git a/src/base/fused_add_rms_norm.h b/src/base/fused_add_rms_norm.h new file mode 100644 index 000000000..3da391f11 --- /dev/null +++ b/src/base/fused_add_rms_norm.h @@ -0,0 +1,78 @@ +#ifndef INFINI_OPS_BASE_FUSED_ADD_RMS_NORM_H_ +#define INFINI_OPS_BASE_FUSED_ADD_RMS_NORM_H_ + +#include +#include + +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +class FusedAddRmsNorm : public Operator { + public: + FusedAddRmsNorm(Tensor input, Tensor residual, + const std::optional weight, float epsilon) + : input_strides_{input.strides()}, + residual_strides_{residual.strides()}, + epsilon_{epsilon}, + dim_{input.size(-1)}, + num_tokens_{dim_ == 0 ? 0 : input.numel() / dim_} { + assert(input.ndim() >= 2 && + "`FusedAddRmsNorm` requires `input` to have at least 2 dimensions"); + assert(dim_ > 0 && + "`FusedAddRmsNorm` requires a non-empty normalized dimension"); + assert(input.shape() == residual.shape() && + "`FusedAddRmsNorm` requires `input` and `residual` to have the same " + "shape"); + assert(input.dtype() == residual.dtype() && + "`FusedAddRmsNorm` requires `input` and `residual` to have the same " + "dtype"); + assert(input.stride(-1) == 1 && + "`FusedAddRmsNorm` requires the last dimension of `input` to be " + "contiguous"); + assert(residual.stride(-1) == 1 && + "`FusedAddRmsNorm` requires the last dimension of `residual` to be " + "contiguous"); + + for (Tensor::Size i = 0; i + 2 < input.ndim(); ++i) { + assert(input.stride(i) == input.size(i + 1) * input.stride(i + 1) && + "`FusedAddRmsNorm` requires `input` rows to have a uniform " + "stride"); + assert(residual.stride(i) == + residual.size(i + 1) * residual.stride(i + 1) && + "`FusedAddRmsNorm` requires `residual` rows to have a uniform " + "stride"); + } + + if (weight.has_value()) { + assert(weight->ndim() == 1 && weight->size(0) == dim_ && + "`FusedAddRmsNorm` requires 1D `weight` with size equal to the " + "normalized dimension"); + assert(weight->dtype() == input.dtype() && + "`FusedAddRmsNorm` requires `input` and `weight` to have the " + "same dtype"); + assert(weight->stride(0) == 1 && + "`FusedAddRmsNorm` requires `weight` to be contiguous"); + } + } + + virtual void operator()(Tensor input, Tensor residual, + const std::optional weight, + float epsilon) const = 0; + + protected: + Tensor::Strides input_strides_; + + Tensor::Strides residual_strides_; + + float epsilon_{}; + + Tensor::Size dim_{0}; + + Tensor::Size num_tokens_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/linear.h b/src/base/linear.h index a5276e612..e0f6e47d9 100644 --- a/src/base/linear.h +++ b/src/base/linear.h @@ -7,55 +7,81 @@ namespace infini::ops { -// Fused linear projection: out = a @ b (+ bias). -// -// When bias is present, computes out = a @ b + bias in a single dispatch. -// When bias is absent, computes out = a @ b (equivalent to Matmul). -// `trans_a` / `trans_b`: If true, transpose the last two dims before -// multiplying. class Linear : public Operator { public: - Linear(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) - : a_shape_{a.shape()}, - b_shape_{b.shape()}, + Linear(const Tensor input, const Tensor weight, std::optional bias, + Tensor out) + : input_shape_{input.shape()}, + input_strides_{input.strides()}, + input_type_{input.dtype()}, + weight_shape_{weight.shape()}, + weight_strides_{weight.strides()}, + weight_type_{weight.dtype()}, out_shape_{out.shape()}, - a_strides_{a.strides()}, - b_strides_{b.strides()}, out_strides_{out.strides()}, - trans_a_{trans_a}, - trans_b_{trans_b}, + out_type_{out.dtype()}, + rows_{Rows(input)}, has_bias_{bias.has_value()} { - assert(a.dtype() == b.dtype() && - "operator `Linear` requires a and b to have the same dtype"); - assert(a.dtype() == out.dtype() && - "operator `Linear` requires a and out to have the same dtype"); + assert(input.ndim() >= 1 && "operator `Linear` requires non-scalar input"); + assert(weight.ndim() == 2 && "operator `Linear` requires 2D `weight`"); + assert(input.size(-1) == weight.size(-1) && + "operator `Linear` input features must match `weight`"); + assert(out.ndim() == input.ndim() && + "operator `Linear` output rank must match input rank"); + assert(out.size(-1) == weight.size(0) && + "operator `Linear` output features must match `weight`"); + + for (Tensor::Size axis = 0; axis + 1 < input.ndim(); ++axis) { + assert(input.size(axis) == out.size(axis) && + "operator `Linear` output leading dimensions must match input"); + } + + assert(input.dtype() == weight.dtype() && + "operator `Linear` requires input and weight to have the same " + "dtype"); + assert(input.dtype() == out.dtype() && + "operator `Linear` requires output to have the input dtype"); if (has_bias_) { + assert(bias->ndim() == 1 && bias->size(0) == weight.size(0) && + "operator `Linear` bias must have shape `[out_features]`"); assert(bias->dtype() == out.dtype() && - "operator `Linear` requires bias and out to have the same dtype"); + "operator `Linear` requires bias to have the output dtype"); } } - virtual void operator()(const Tensor a, const Tensor b, - std::optional bias, bool trans_a, - bool trans_b, Tensor out) const = 0; + virtual void operator()(const Tensor input, const Tensor weight, + std::optional bias, Tensor out) const = 0; protected: - Tensor::Shape a_shape_; + static Tensor::Size Rows(const Tensor input) { + Tensor::Size rows = 1; + + for (Tensor::Size axis = 0; axis + 1 < input.ndim(); ++axis) { + rows *= input.size(axis); + } - Tensor::Shape b_shape_; + return input.ndim() == 0 ? 0 : rows; + } - Tensor::Shape out_shape_; + Tensor::Shape input_shape_; - Tensor::Strides a_strides_; + Tensor::Strides input_strides_; - Tensor::Strides b_strides_; + DataType input_type_; + + Tensor::Shape weight_shape_; + + Tensor::Strides weight_strides_; + + DataType weight_type_; + + Tensor::Shape out_shape_; Tensor::Strides out_strides_; - bool trans_a_{false}; + DataType out_type_; - bool trans_b_{false}; + Tensor::Size rows_{0}; bool has_bias_{false}; }; diff --git a/src/base/matmul.h b/src/base/matmul.h index 071feaeaa..3b202a860 100644 --- a/src/base/matmul.h +++ b/src/base/matmul.h @@ -8,32 +8,43 @@ namespace infini::ops { class Matmul : public Operator { public: - // `trans_a` / `trans_b`: If true, transpose the last two dims of `a` / `b` - // before multiplying. These are constructor parameters so the `CacheKey` - // encodes the transposition and distinct descriptors are cached for each - // combination. - Matmul(const Tensor a, const Tensor b, Tensor c, bool trans_a, bool trans_b) - : a_shape_{a.shape()}, - b_shape_{b.shape()}, - c_shape_{c.shape()}, - trans_a_{trans_a}, - trans_b_{trans_b} { - assert(a.dtype() == b.dtype()); + Matmul(const Tensor input, const Tensor other, Tensor out) + : input_shape_{input.shape()}, + input_strides_{input.strides()}, + input_type_{input.dtype()}, + other_shape_{other.shape()}, + other_strides_{other.strides()}, + other_type_{other.dtype()}, + out_shape_{out.shape()}, + out_strides_{out.strides()}, + out_type_{out.dtype()} { + assert(input.dtype() == other.dtype() && + "operator `Matmul` requires inputs to have the same dtype"); + assert(input.dtype() == out.dtype() && + "operator `Matmul` requires output to have the input dtype"); } - virtual void operator()(const Tensor a, const Tensor b, Tensor c, - bool trans_a, bool trans_b) const = 0; + virtual void operator()(const Tensor input, const Tensor other, + Tensor out) const = 0; protected: - Tensor::Shape a_shape_; + Tensor::Shape input_shape_; - Tensor::Shape b_shape_; + Tensor::Strides input_strides_; - Tensor::Shape c_shape_; + DataType input_type_; - bool trans_a_{false}; + Tensor::Shape other_shape_; - bool trans_b_{false}; + Tensor::Strides other_strides_; + + DataType other_type_; + + Tensor::Shape out_shape_; + + Tensor::Strides out_strides_; + + DataType out_type_; }; } // namespace infini::ops diff --git a/src/base/reshape_and_cache.h b/src/base/reshape_and_cache.h index 4bbd5db8b..dfdb5349c 100644 --- a/src/base/reshape_and_cache.h +++ b/src/base/reshape_and_cache.h @@ -1,78 +1,161 @@ #ifndef INFINI_OPS_BASE_RESHAPE_AND_CACHE_H_ #define INFINI_OPS_BASE_RESHAPE_AND_CACHE_H_ -#include -#include +#include +#include "data_type.h" #include "operator.h" namespace infini::ops { -// Scatter `key` / `value` tokens into a paged KV cache. -// -// Interface follows vLLM's `reshape_and_cache` kernel: -// `vllm._custom_ops.reshape_and_cache_flash` -// -// `kv_cache` layout: `[2, num_blocks, block_size, num_kv_heads, head_size]`. -// `slot_mapping`: 1D `[num_tokens]`, each entry is the linear slot index -// into the cache. Padding tokens must be filtered by the caller (no -// negative indices). class ReshapeAndCache : public Operator { public: - ReshapeAndCache(const Tensor key, const Tensor value, const Tensor kv_cache, - const Tensor slot_mapping, Tensor kv_cache_out) - : num_tokens_{key.size(0)}, - num_kv_heads_{key.size(1)}, - head_size_{key.size(2)}, - block_size_{kv_cache.size(2)}, - key_shape_{key.shape()}, + ReshapeAndCache(const Tensor key, const Tensor value, Tensor key_cache, + Tensor value_cache, const Tensor slot_mapping, + const std::string kv_cache_dtype, const Tensor k_scale, + const Tensor v_scale) + : key_shape_{key.shape()}, value_shape_{value.shape()}, - kv_cache_shape_{kv_cache.shape()}, + key_cache_shape_{key_cache.shape()}, + value_cache_shape_{value_cache.shape()}, slot_mapping_shape_{slot_mapping.shape()}, + k_scale_shape_{k_scale.shape()}, + v_scale_shape_{v_scale.shape()}, key_strides_{key.strides()}, value_strides_{value.strides()}, - kv_cache_strides_{kv_cache.strides()}, + key_cache_strides_{key_cache.strides()}, + value_cache_strides_{value_cache.strides()}, slot_mapping_strides_{slot_mapping.strides()}, - kv_cache_out_strides_{kv_cache_out.strides()} { + k_scale_strides_{k_scale.strides()}, + v_scale_strides_{v_scale.strides()}, + key_type_{key.dtype()}, + key_cache_type_{key_cache.dtype()}, + value_cache_type_{value_cache.dtype()}, + k_scale_type_{k_scale.dtype()}, + v_scale_type_{v_scale.dtype()}, + kv_cache_dtype_{kv_cache_dtype}, + num_tokens_{slot_mapping.size(0)}, + num_heads_{key.size(1)}, + head_size_{key.size(2)}, + block_size_{key_cache.size(3)}, + x_{key_cache.size(4)}, + device_index_{key.device().index()} { + assert(key.ndim() == 3 && value.ndim() == 3 && + "`ReshapeAndCache` requires 3D `key` and `value`"); assert(key.shape() == value.shape() && - "`ReshapeAndCache` requires key and value same shape"); - assert(kv_cache.ndim() == 5 && - "`ReshapeAndCache` requires kv_cache to be 5D [2, num_blocks, " - "block_size, num_kv_heads, head_size]"); + "`ReshapeAndCache` requires `key` and `value` to have the same " + "shape"); + assert(key.dtype() == value.dtype() && + "`ReshapeAndCache` requires `key` and `value` to have the same " + "dtype"); + assert((key_type_ == DataType::kFloat16 || + key_type_ == DataType::kBFloat16 || + key_type_ == DataType::kFloat32) && + "`ReshapeAndCache` supports float16, bfloat16, and float32 inputs"); + assert(key_cache.ndim() == 5 && value_cache.ndim() == 4 && + "`ReshapeAndCache` requires 5D `key_cache` and 4D `value_cache`"); assert(slot_mapping.ndim() == 1 && - "`ReshapeAndCache` requires slot_mapping to be 1D"); + slot_mapping.dtype() == DataType::kInt64 && + slot_mapping.stride(0) == 1 && + "`ReshapeAndCache` requires contiguous int64 `slot_mapping`"); + assert(num_tokens_ <= key.size(0) && + "`ReshapeAndCache` requires enough key/value rows for all slots"); + assert(x_ > 0 && head_size_ % x_ == 0 && + "`ReshapeAndCache` requires `head_size` divisible by cache vector " + "width"); + assert(key_cache.size(0) == value_cache.size(0) && + key_cache.size(1) == num_heads_ && + value_cache.size(1) == num_heads_ && + key_cache.size(2) == head_size_ / x_ && + value_cache.size(2) == head_size_ && + value_cache.size(3) == block_size_ && + "`ReshapeAndCache` cache shapes do not match key/value geometry"); + assert(key.stride(2) == 1 && key.stride(1) == head_size_ && + value.stride(2) == 1 && value.stride(1) == head_size_ && + "`ReshapeAndCache` requires contiguous head dimensions"); + assert(key_cache.stride(4) == 1 && key_cache.stride(3) == x_ && + key_cache.stride(2) == block_size_ * x_ && + key_cache.stride(1) == head_size_ * block_size_ && + key_cache.stride(0) == num_heads_ * head_size_ * block_size_ && + "`ReshapeAndCache` requires contiguous vLLM key-cache layout"); + assert(value_cache.stride(3) == 1 && value_cache.stride(2) == block_size_ && + value_cache.stride(1) == head_size_ * block_size_ && + value_cache.stride(0) == num_heads_ * head_size_ * block_size_ && + "`ReshapeAndCache` requires contiguous vLLM value-cache layout"); + assert((kv_cache_dtype_ == "auto" || kv_cache_dtype_ == "fp8" || + kv_cache_dtype_ == "fp8_e4m3" || kv_cache_dtype_ == "fp8_e5m2") && + "`ReshapeAndCache` received unsupported `kv_cache_dtype`"); + + const bool quantized = kv_cache_dtype_ != "auto"; + assert((quantized ? key_cache_type_ == DataType::kUInt8 + : key_cache_type_ == key_type_) && + key_cache_type_ == value_cache_type_ && + "`ReshapeAndCache` cache storage dtype does not match " + "`kv_cache_dtype`"); + assert(k_scale.numel() == 1 && v_scale.numel() == 1 && + k_scale_type_ == DataType::kFloat32 && + v_scale_type_ == DataType::kFloat32 && + "`ReshapeAndCache` requires scalar float32 scales"); } virtual void operator()(const Tensor key, const Tensor value, - const Tensor kv_cache, const Tensor slot_mapping, - Tensor kv_cache_out) const = 0; + Tensor key_cache, Tensor value_cache, + const Tensor slot_mapping, + const std::string kv_cache_dtype, + const Tensor k_scale, const Tensor v_scale) const = 0; protected: - Tensor::Size num_tokens_{0}; - - Tensor::Size num_kv_heads_{0}; - - Tensor::Size head_size_{0}; - - Tensor::Size block_size_{0}; - Tensor::Shape key_shape_; Tensor::Shape value_shape_; - Tensor::Shape kv_cache_shape_; + Tensor::Shape key_cache_shape_; + + Tensor::Shape value_cache_shape_; Tensor::Shape slot_mapping_shape_; + Tensor::Shape k_scale_shape_; + + Tensor::Shape v_scale_shape_; + Tensor::Strides key_strides_; Tensor::Strides value_strides_; - Tensor::Strides kv_cache_strides_; + Tensor::Strides key_cache_strides_; + + Tensor::Strides value_cache_strides_; Tensor::Strides slot_mapping_strides_; - Tensor::Strides kv_cache_out_strides_; + Tensor::Strides k_scale_strides_; + + Tensor::Strides v_scale_strides_; + + DataType key_type_; + + DataType key_cache_type_; + + DataType value_cache_type_; + + DataType k_scale_type_; + + DataType v_scale_type_; + + std::string kv_cache_dtype_; + + Tensor::Size num_tokens_{0}; + + Tensor::Size num_heads_{0}; + + Tensor::Size head_size_{0}; + + Tensor::Size block_size_{0}; + + Tensor::Size x_{0}; + + int device_index_{0}; }; } // namespace infini::ops diff --git a/src/base/rotary_embedding.h b/src/base/rotary_embedding.h index 10426ee86..f405faa43 100644 --- a/src/base/rotary_embedding.h +++ b/src/base/rotary_embedding.h @@ -1,87 +1,181 @@ #ifndef INFINI_OPS_BASE_ROTARY_EMBEDDING_H_ #define INFINI_OPS_BASE_ROTARY_EMBEDDING_H_ -#include -#include +#include +#include "data_type.h" #include "operator.h" namespace infini::ops { -// Rotary position embedding (RoPE) applied in-place to Q and K. -// -// Interface follows vLLM's `RotaryEmbedding.forward_oot()`: -// `vllm.model_executor.layers.rotary_embedding.RotaryEmbedding` -// -// `positions`: `[T]` token position indices. -// `cos_sin_cache`: precomputed `[max_seq_len, rotary_dim]` table. -// `query` / `key`: `[T, N, D]` (TND layout), mutated in-place into -// `query_out` / `key_out`. class RotaryEmbedding : public Operator { public: - RotaryEmbedding(const Tensor positions, const Tensor query, const Tensor key, - const Tensor cos_sin_cache, int64_t head_size, - int64_t rotary_dim, bool is_neox_style, Tensor query_out, - Tensor key_out) - : num_tokens_{query.size(0)}, - num_heads_{static_cast(query.size(1))}, - num_kv_heads_{static_cast(key.size(1))}, - head_size_{head_size}, - rotary_dim_{rotary_dim}, - is_neox_style_{is_neox_style}, + RotaryEmbedding(const Tensor positions, Tensor query, + std::optional key, int64_t head_size, + const Tensor cos_sin_cache, bool is_neox, + int64_t rope_dim_offset = 0, bool inverse = false) + : positions_shape_{positions.shape()}, query_shape_{query.shape()}, - key_shape_{key.shape()}, + key_shape_{key.has_value() ? key->shape() : Tensor::Shape{}}, cos_sin_cache_shape_{cos_sin_cache.shape()}, - query_out_shape_{query_out.shape()}, - key_out_shape_{key_out.shape()}, + positions_strides_{positions.strides()}, query_strides_{query.strides()}, - key_strides_{key.strides()}, - query_out_strides_{query_out.strides()}, - key_out_strides_{key_out.strides()} { - assert(query.ndim() == 3 && - "`RotaryEmbedding` requires query to be 3D [T, N, D]"); - assert(key.ndim() == 3 && - "`RotaryEmbedding` requires key to be 3D [T, N_kv, D]"); - assert(rotary_dim <= head_size && - "`RotaryEmbedding` requires rotary_dim <= head_size"); + key_strides_{key.has_value() ? key->strides() : Tensor::Strides{}}, + cos_sin_cache_strides_{cos_sin_cache.strides()}, + positions_type_{positions.dtype()}, + query_type_{query.dtype()}, + cos_sin_cache_type_{cos_sin_cache.dtype()}, + num_tokens_{positions.numel()}, + positions_ndim_{positions.ndim()}, + head_size_{head_size}, + rot_dim_{static_cast(cos_sin_cache.size(1))}, + rope_dim_offset_{rope_dim_offset}, + is_neox_{is_neox}, + inverse_{inverse}, + query_hidden_size_{num_tokens_ == 0 ? 0 : query.numel() / num_tokens_}, + key_hidden_size_{key.has_value() && num_tokens_ != 0 + ? key->numel() / num_tokens_ + : 0}, + num_heads_{head_size_ == 0 ? 0 : query_hidden_size_ / head_size_}, + num_kv_heads_{head_size_ == 0 + ? 0 + : (key.has_value() ? key_hidden_size_ / head_size_ + : num_heads_)}, + query_token_stride_{query.stride(positions_ndim_ - 1)}, + key_token_stride_{key.has_value() ? key->stride(positions_ndim_ - 1) + : 0}, + query_head_stride_{query.ndim() == positions_ndim_ + 2 + ? query.stride(-2) + : head_size_}, + key_head_stride_{key.has_value() && key->ndim() == positions_ndim_ + 2 + ? key->stride(-2) + : head_size_}, + device_index_{query.device().index()} { + assert((positions_ndim_ == 1 || positions_ndim_ == 2) && + "`RotaryEmbedding` requires 1D or 2D `positions`"); + assert(positions_type_ == DataType::kInt64 && + "`RotaryEmbedding` requires int64 `positions`"); + assert(positions.stride(-1) == 1 && + (positions_ndim_ == 1 || positions.stride(0) == positions.size(1)) && + "`RotaryEmbedding` requires contiguous `positions`"); + assert((query.ndim() == positions_ndim_ + 1 || + query.ndim() == positions_ndim_ + 2) && + "`RotaryEmbedding` received unsupported `query` rank"); + assert(head_size_ > 0 && query_hidden_size_ % head_size_ == 0 && + "`RotaryEmbedding` requires query hidden size divisible by " + "`head_size`"); + assert((query_type_ == DataType::kFloat16 || + query_type_ == DataType::kBFloat16 || + query_type_ == DataType::kFloat32) && + "`RotaryEmbedding` supports float16, bfloat16, and float32 query"); + assert(cos_sin_cache.ndim() == 2 && rot_dim_ % 2 == 0 && + cos_sin_cache.stride(1) == 1 && + "`RotaryEmbedding` requires contiguous 2D `cos_sin_cache` with " + "even rotary dimension"); + assert((cos_sin_cache_type_ == DataType::kFloat16 || + cos_sin_cache_type_ == DataType::kBFloat16 || + cos_sin_cache_type_ == DataType::kFloat32) && + "`RotaryEmbedding` supports float16, bfloat16, and float32 cache"); + assert(rope_dim_offset_ >= 0 && rot_dim_ + rope_dim_offset_ <= head_size_ && + "`RotaryEmbedding` rotary dimensions exceed `head_size`"); + assert(query.stride(-1) == 1 && + "`RotaryEmbedding` requires contiguous query head dimensions"); + + if (positions_ndim_ == 1) { + assert(query.size(0) == positions.size(0) && + (!key.has_value() || key->size(0) == positions.size(0)) && + "`RotaryEmbedding` requires matching token counts"); + } else { + assert(query.size(0) == positions.size(0) && + query.size(1) == positions.size(1) && + (!key.has_value() || (key->size(0) == positions.size(0) && + key->size(1) == positions.size(1))) && + "`RotaryEmbedding` requires matching batch and sequence sizes"); + } + + if (query.ndim() == positions_ndim_ + 2) { + assert(query.size(-1) == static_cast(head_size_) && + "`RotaryEmbedding` query head dimension does not match " + "`head_size`"); + } + + if (key.has_value()) { + assert((key->ndim() == positions_ndim_ + 1 || + key->ndim() == positions_ndim_ + 2) && + key_hidden_size_ % head_size_ == 0 && + key->dtype() == query_type_ && key->stride(-1) == 1 && + "`RotaryEmbedding` key layout or dtype is incompatible with " + "query"); + assert(num_kv_heads_ > 0 && num_heads_ % num_kv_heads_ == 0 && + "`RotaryEmbedding` requires query heads divisible by key heads"); + if (key->ndim() == positions_ndim_ + 2) { + assert(key->size(-1) == static_cast(head_size_) && + "`RotaryEmbedding` key head dimension does not match " + "`head_size`"); + } + } } - virtual void operator()(const Tensor positions, const Tensor query, - const Tensor key, const Tensor cos_sin_cache, - int64_t head_size, int64_t rotary_dim, - bool is_neox_style, Tensor query_out, - Tensor key_out) const = 0; + virtual void operator()(const Tensor positions, Tensor query, + std::optional key, int64_t head_size, + const Tensor cos_sin_cache, bool is_neox, + int64_t rope_dim_offset = 0, + bool inverse = false) const = 0; protected: - Tensor::Size num_tokens_{0}; + Tensor::Shape positions_shape_; + + Tensor::Shape query_shape_; + + Tensor::Shape key_shape_; + + Tensor::Shape cos_sin_cache_shape_; + + Tensor::Strides positions_strides_; + + Tensor::Strides query_strides_; + + Tensor::Strides key_strides_; + + Tensor::Strides cos_sin_cache_strides_; - int64_t num_heads_{0}; + DataType positions_type_; - int64_t num_kv_heads_{0}; + DataType query_type_; + + DataType cos_sin_cache_type_; + + Tensor::Size num_tokens_{0}; + + Tensor::Size positions_ndim_{0}; int64_t head_size_{0}; - int64_t rotary_dim_{0}; + int64_t rot_dim_{0}; - bool is_neox_style_{true}; + int64_t rope_dim_offset_{0}; - Tensor::Shape query_shape_; + bool is_neox_{false}; - Tensor::Shape key_shape_; + bool inverse_{false}; - Tensor::Shape cos_sin_cache_shape_; + Tensor::Size query_hidden_size_{0}; - Tensor::Shape query_out_shape_; + Tensor::Size key_hidden_size_{0}; - Tensor::Shape key_out_shape_; + Tensor::Size num_heads_{0}; - Tensor::Strides query_strides_; + Tensor::Size num_kv_heads_{0}; - Tensor::Strides key_strides_; + int64_t query_token_stride_{0}; + + int64_t key_token_stride_{0}; + + int64_t query_head_stride_{0}; - Tensor::Strides query_out_strides_; + int64_t key_head_stride_{0}; - Tensor::Strides key_out_strides_; + int device_index_{0}; }; } // namespace infini::ops diff --git a/src/base/scaled_dot_product_attention.h b/src/base/scaled_dot_product_attention.h new file mode 100644 index 000000000..ace5d5cde --- /dev/null +++ b/src/base/scaled_dot_product_attention.h @@ -0,0 +1,129 @@ +#ifndef INFINI_OPS_BASE_SCALED_DOT_PRODUCT_ATTENTION_H_ +#define INFINI_OPS_BASE_SCALED_DOT_PRODUCT_ATTENTION_H_ + +#include + +#include "operator.h" + +namespace infini::ops { + +class ScaledDotProductAttention : public Operator { + public: + ScaledDotProductAttention(const Tensor query, const Tensor key, + const Tensor value, + const std::optional attn_mask, + double dropout_p, bool is_causal, + const std::optional scale, bool enable_gqa, + Tensor out) + : query_shape_{query.shape()}, + key_shape_{key.shape()}, + value_shape_{value.shape()}, + attn_mask_shape_{attn_mask.has_value() ? attn_mask->shape() + : Tensor::Shape{}}, + out_shape_{out.shape()}, + query_strides_{query.strides()}, + key_strides_{key.strides()}, + value_strides_{value.strides()}, + attn_mask_strides_{attn_mask.has_value() ? attn_mask->strides() + : Tensor::Strides{}}, + out_strides_{out.strides()}, + query_type_{query.dtype()}, + attn_mask_type_{attn_mask.has_value() ? attn_mask->dtype() + : query.dtype()}, + dropout_p_{dropout_p}, + is_causal_{is_causal}, + scale_{scale}, + enable_gqa_{enable_gqa}, + device_index_{query.device().index()} { + assert(query.ndim() >= 2 && key.ndim() >= 2 && value.ndim() >= 2 && + "`ScaledDotProductAttention` requires rank-2 or higher inputs"); + assert(query.dtype() == key.dtype() && query.dtype() == value.dtype() && + query.dtype() == out.dtype() && + "`ScaledDotProductAttention` requires matching input/output " + "dtypes"); + assert(query.size(-1) == key.size(-1) && key.size(-2) == value.size(-2) && + "`ScaledDotProductAttention` input dimensions are incompatible"); + auto expected_out_shape = query.shape(); + expected_out_shape.back() = value.size(-1); + assert(out.shape() == expected_out_shape && + "`ScaledDotProductAttention` output shape is incorrect"); + assert(dropout_p_ >= 0.0 && dropout_p_ <= 1.0 && + "`ScaledDotProductAttention` requires `dropout_p` in [0, 1]"); + } + + ScaledDotProductAttention(const Tensor query, const Tensor key, + const Tensor value, Tensor out) + : ScaledDotProductAttention{query, key, value, std::nullopt, 0.0, + false, std::nullopt, false, out} {} + + virtual void operator()(const Tensor query, const Tensor key, + const Tensor value, + const std::optional attn_mask, + double dropout_p, bool is_causal, + const std::optional scale, bool enable_gqa, + Tensor out) const = 0; + + void operator()(const Tensor query, const Tensor key, const Tensor value, + Tensor out) const { + (*this)(query, key, value, std::nullopt, 0.0, false, std::nullopt, false, + out); + } + + template + static auto MakeReturnValue( + const TensorLike& query, const TensorLike& key, const TensorLike& value, + const std::optional attn_mask = std::nullopt, + double dropout_p = 0.0, bool is_causal = false, + const std::optional scale = std::nullopt, + bool enable_gqa = false) { + (void)key; + (void)attn_mask; + (void)dropout_p; + (void)is_causal; + (void)scale; + (void)enable_gqa; + + auto out_shape = query.shape(); + out_shape.back() = value.size(-1); + return TensorLike::Empty(out_shape, query.dtype(), query.device()); + } + + protected: + Tensor::Shape query_shape_; + + Tensor::Shape key_shape_; + + Tensor::Shape value_shape_; + + Tensor::Shape attn_mask_shape_; + + Tensor::Shape out_shape_; + + Tensor::Strides query_strides_; + + Tensor::Strides key_strides_; + + Tensor::Strides value_strides_; + + Tensor::Strides attn_mask_strides_; + + Tensor::Strides out_strides_; + + DataType query_type_; + + DataType attn_mask_type_; + + double dropout_p_{0.0}; + + bool is_causal_{false}; + + std::optional scale_; + + bool enable_gqa_{false}; + + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/scaled_softmax.h b/src/base/scaled_softmax.h index 707bf1ffc..671d052ee 100644 --- a/src/base/scaled_softmax.h +++ b/src/base/scaled_softmax.h @@ -11,6 +11,8 @@ namespace infini::ops { +/// \deprecated Use `ScaledSoftmaxInfinilm`. This interface will be removed in +/// a future release. class ScaledSoftmax : public Operator { public: ScaledSoftmax(const Tensor input, double scale, Tensor out) diff --git a/src/base/scaled_softmax_infinilm.h b/src/base/scaled_softmax_infinilm.h new file mode 100644 index 000000000..1ca081be1 --- /dev/null +++ b/src/base/scaled_softmax_infinilm.h @@ -0,0 +1,58 @@ +#ifndef INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ + +#include +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +class ScaledSoftmaxInfinilm : public Operator { + public: + ScaledSoftmaxInfinilm(const Tensor input, double scale, Tensor out) + : scale_{scale}, + batch_size_{input.size(0)}, + vocab_size_{input.size(1)}, + dtype_{input.dtype()}, + input_strides_{input.strides()}, + out_strides_{out.strides()} { + assert(input.ndim() == 2 && + "`ScaledSoftmaxInfinilm` currently supports 2D `[batch, vocab]` " + "input"); + assert(input.shape() == out.shape() && + "`ScaledSoftmaxInfinilm` requires `input` and `out` to have the " + "same shape"); + assert(input.dtype() == out.dtype() && + "`ScaledSoftmaxInfinilm` requires `input` and `out` to have the " + "same dtype"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && + "`ScaledSoftmaxInfinilm` requires a floating point dtype"); + assert(std::isfinite(scale_) && + "`ScaledSoftmaxInfinilm` requires a finite `scale`"); + } + + virtual void operator()(const Tensor input, double scale, + Tensor out) const = 0; + + protected: + double scale_{1.0}; + + Tensor::Size batch_size_{0}; + + Tensor::Size vocab_size_{0}; + + DataType dtype_; + + Tensor::Strides input_strides_; + + Tensor::Strides out_strides_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ diff --git a/src/base/silu_and_mul.h b/src/base/silu_and_mul.h new file mode 100644 index 000000000..603c132d2 --- /dev/null +++ b/src/base/silu_and_mul.h @@ -0,0 +1,74 @@ +#ifndef INFINI_OPS_BASE_SILU_AND_MUL_H_ +#define INFINI_OPS_BASE_SILU_AND_MUL_H_ + +#include "data_type.h" +#include "operator.h" + +namespace infini::ops { + +class SiluAndMul : public Operator { + public: + SiluAndMul(const Tensor input, Tensor out) + : ndim_{out.ndim()}, + output_size_{out.numel()}, + input_type_{input.dtype()}, + out_type_{out.dtype()}, + input_shape_{input.shape()}, + out_shape_{out.shape()}, + input_strides_{input.strides()}, + out_strides_{out.strides()}, + hidden_size_{out.size(-1)}, + is_input_contiguous_{input.IsContiguous()}, + is_out_contiguous_{out.IsContiguous()} { + assert(input.ndim() == out.ndim() && + "`SiluAndMul` requires input and output ranks to match"); + assert(input_type_ == out_type_ && + "`SiluAndMul` requires input and output dtypes to match"); + assert(input.size(-1) == 2 * hidden_size_ && + "`SiluAndMul` requires input last dimension to be twice output " + "last dimension"); + for (Tensor::Size i = 0; i + 1 < ndim_; ++i) { + assert(input.size(i) == out.size(i) && + "`SiluAndMul` requires matching leading dimensions"); + } + } + + virtual void operator()(const Tensor input, Tensor out) const = 0; + + template + static auto MakeReturnValue(const TensorLike& input) { + auto out_shape = input.shape(); + assert(!out_shape.empty() && out_shape.back() % 2 == 0 && + "`SiluAndMul` requires an even input last dimension"); + out_shape.back() /= 2; + + return TensorLike::Empty(out_shape, input.dtype(), input.device()); + } + + protected: + Tensor::Size ndim_{0}; + + Tensor::Size output_size_{0}; + + DataType input_type_; + + DataType out_type_; + + Tensor::Shape input_shape_; + + Tensor::Shape out_shape_; + + Tensor::Strides input_strides_; + + Tensor::Strides out_strides_; + + Tensor::Size hidden_size_{0}; + + bool is_input_contiguous_{false}; + + bool is_out_contiguous_{false}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/swiglu.h b/src/base/swiglu.h index 023b14a21..145748ced 100644 --- a/src/base/swiglu.h +++ b/src/base/swiglu.h @@ -7,6 +7,8 @@ namespace infini::ops { +/// \deprecated Use `SiluAndMul`. This interface will be removed in a future +/// release. class Swiglu : public Operator { public: Swiglu(const Tensor input, const Tensor gate, Tensor out) diff --git a/src/base/top_k_top_p_sample_infinilm.h b/src/base/top_k_top_p_sample_infinilm.h new file mode 100644 index 000000000..2d9b092d9 --- /dev/null +++ b/src/base/top_k_top_p_sample_infinilm.h @@ -0,0 +1,81 @@ +#ifndef INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ +#define INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ + +#include +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +class TopKTopPSampleInfinilm : public Operator { + public: + TopKTopPSampleInfinilm(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, + uint64_t offset, Tensor out) + : batch_size_{logits.size(0)}, + vocab_size_{logits.size(1)}, + dtype_{logits.dtype()} { + (void)seed; + (void)offset; + assert(logits.ndim() == 2 && + "`TopKTopPSampleInfinilm` requires 2D `[batch_size, vocab_size]` " + "logits"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && + "`TopKTopPSampleInfinilm` requires floating-point logits"); + assert(out.ndim() == 1 && + "`TopKTopPSampleInfinilm` requires 1D `[batch_size]` output"); + assert(out.size(0) == batch_size_ && + "`TopKTopPSampleInfinilm` requires output batch size to match " + "logits"); + assert(out.dtype() == DataType::kInt32 && + "`TopKTopPSampleInfinilm` requires int32 output"); + + ValidateK(k); + ValidateP(p); + } + + virtual void operator()(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, + uint64_t offset, Tensor out) const = 0; + + protected: + void ValidateK(std::optional k) const { + if (!k.has_value()) return; + + assert(k->ndim() == 1 && + "`TopKTopPSampleInfinilm` requires `k` to be 1D when provided"); + assert((k->size(0) == 1 || k->size(0) == batch_size_) && + "`TopKTopPSampleInfinilm` requires `k` shape [1] or [batch_size]"); + assert((k->dtype() == DataType::kInt32 || k->dtype() == DataType::kInt64) && + "`TopKTopPSampleInfinilm` requires int32 or int64 `k`"); + } + + void ValidateP(std::optional p) const { + if (!p.has_value()) return; + + assert(p->ndim() == 1 && + "`TopKTopPSampleInfinilm` requires `p` to be 1D when provided"); + assert((p->size(0) == 1 || p->size(0) == batch_size_) && + "`TopKTopPSampleInfinilm` requires `p` shape [1] or [batch_size]"); + assert((p->dtype() == DataType::kFloat16 || + p->dtype() == DataType::kBFloat16 || + p->dtype() == DataType::kFloat32 || + p->dtype() == DataType::kFloat64) && + "`TopKTopPSampleInfinilm` requires floating-point `p`"); + } + + Tensor::Size batch_size_{0}; + + Tensor::Size vocab_size_{0}; + + DataType dtype_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ diff --git a/src/base/top_k_top_p_sampler.h b/src/base/top_k_top_p_sampler.h index 085a2b937..55e4fa0bb 100644 --- a/src/base/top_k_top_p_sampler.h +++ b/src/base/top_k_top_p_sampler.h @@ -10,10 +10,11 @@ namespace infini::ops { -// `TopKTopPSampler` samples token ids from 2D `logits` after optional rank and -// nucleus filtering. The name and tensor boundary follow vLLM's -// `TopKTopPSampler`; temperature scaling is intentionally handled by callers. +// Legacy sampler for 2D `logits` after optional rank and nucleus filtering. +// Temperature scaling is intentionally handled by callers. // The optional `k` and `p` tensors may be shaped as `[1]` or `[batch_size]`. +/// \deprecated Use `TopKTopPSampleInfinilm`. This interface will be removed in +/// a future release. class TopKTopPSampler : public Operator { public: TopKTopPSampler(const Tensor logits, std::optional k, diff --git a/src/native/ascend/ops/add/kernel.h b/src/native/ascend/ops/add/kernel.h index 8bf4a2564..84aa2f4b0 100644 --- a/src/native/ascend/ops/add/kernel.h +++ b/src/native/ascend/ops/add/kernel.h @@ -15,8 +15,9 @@ namespace infini::ops { template <> class Operator : public Add { public: - Operator(const Tensor input, const Tensor other, Tensor out) - : Add(input, other, out), + Operator(const Tensor input, const Tensor other, const double alpha, + Tensor out) + : Add(input, other, alpha, out), in_cache_(input), oth_cache_(other), out_cache_(out) { @@ -25,12 +26,19 @@ class Operator : public Add { // The alpha scalar type must match the tensor dtype: use int64 for integer // dtypes and float for floating-point dtypes. if (ascend::IsIntegerDtype(input.dtype())) { + assert(alpha == static_cast(alpha) && + "operator `Add` requires integral `alpha` for integer tensors"); + alpha_int_storage_ = static_cast(alpha); alpha_ = aclCreateScalar(&alpha_int_storage_, ACL_INT64); } else { + alpha_float_storage_ = static_cast(alpha); alpha_ = aclCreateScalar(&alpha_float_storage_, ACL_FLOAT); } } + Operator(const Tensor input, const Tensor other, Tensor out) + : Operator(input, other, 1.0, out) {} + ~Operator() { if (!ascend::IsAclRuntimeAlive()) return; @@ -46,7 +54,7 @@ class Operator : public Add { } void operator()(const Tensor input, const Tensor other, - Tensor out) const override { + const double /*alpha*/, Tensor out) const override { auto stream = static_cast(stream_); auto t_in = in_cache_.get(const_cast(input.data())); auto t_oth = oth_cache_.get(const_cast(other.data())); diff --git a/src/native/ascend/ops/cat/kernel.h b/src/native/ascend/ops/cat/kernel.h index 32948181a..6dd18ab6c 100644 --- a/src/native/ascend/ops/cat/kernel.h +++ b/src/native/ascend/ops/cat/kernel.h @@ -17,14 +17,11 @@ namespace infini::ops { template <> class Operator : public Cat { public: - Operator(const Tensor first_input, std::vector rest_inputs, - int64_t dim, Tensor out) - : Cat(first_input, rest_inputs, dim, out), out_cache_(out) { - // Build `AclTensorCache` for each input tensor. + Operator(const std::vector tensors, const int64_t dim, Tensor out) + : Cat(tensors, dim, out), out_cache_(out) { in_caches_.reserve(input_count_); - in_caches_.emplace_back(first_input); - for (const auto& t : rest_inputs) { - in_caches_.emplace_back(t); + for (const auto& tensor : tensors) { + in_caches_.emplace_back(tensor); } } @@ -44,18 +41,9 @@ class Operator : public Cat { if (tensor_list_) aclDestroyTensorList(tensor_list_); } - void operator()(const Tensor first_input, std::vector rest_inputs, - int64_t /*dim*/, Tensor out) const override { + void operator()(const std::vector tensors, const int64_t /*dim*/, + Tensor out) const override { auto stream = static_cast(stream_); - - // Collect all input tensors in order. - std::vector inputs; - inputs.reserve(input_count_); - inputs.push_back(&first_input); - for (const auto& t : rest_inputs) { - inputs.push_back(&t); - } - auto t_out = out_cache_.get(out.data()); if (!executor_) { @@ -63,7 +51,7 @@ class Operator : public Cat { std::vector acl_tensors(input_count_); for (size_t i = 0; i < input_count_; ++i) { acl_tensors[i] = - in_caches_[i].get(const_cast(inputs[i]->data())); + in_caches_[i].get(const_cast(tensors[i].data())); } tensor_list_ = @@ -79,7 +67,7 @@ class Operator : public Cat { // `aclTensor*` objects inside `tensor_list_`, so updating their data // pointers is sufficient — no `aclSetInputTensorAddr` needed. for (size_t i = 0; i < input_count_; ++i) { - in_caches_[i].get(const_cast(inputs[i]->data())); + in_caches_[i].get(const_cast(tensors[i].data())); } aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); } diff --git a/src/native/ascend/ops/embedding/kernel.h b/src/native/ascend/ops/embedding/kernel.h index 4b7b0524a..84bc90582 100644 --- a/src/native/ascend/ops/embedding/kernel.h +++ b/src/native/ascend/ops/embedding/kernel.h @@ -16,9 +16,11 @@ namespace infini::ops { template <> class Operator : public Embedding { public: - Operator(const Tensor input_ids, const Tensor weight, Tensor out) - : Embedding(input_ids, weight, out), - input_ids_cache_(input_ids), + Operator(const Tensor weight, const Tensor indices, const int64_t padding_idx, + const bool scale_grad_by_freq, const bool sparse, Tensor out) + : Embedding(weight, indices, padding_idx, scale_grad_by_freq, sparse, + out), + indices_cache_(indices), weight_cache_(weight), out_cache_(out) { assert((weight_dtype_ == DataType::kFloat16 || @@ -28,33 +30,37 @@ class Operator : public Embedding { "`float32` weights"); } + Operator(const Tensor weight, const Tensor indices, Tensor out) + : Operator(weight, indices, -1, false, false, out) {} + ~Operator() { if (!ascend::IsAclRuntimeAlive()) return; - input_ids_cache_.release(); + indices_cache_.release(); weight_cache_.release(); out_cache_.release(); } - void operator()(const Tensor input_ids, const Tensor weight, + void operator()(const Tensor weight, const Tensor indices, + const int64_t /*padding_idx*/, + const bool /*scale_grad_by_freq*/, const bool /*sparse*/, Tensor out) const override { auto stream = static_cast(stream_); auto t_weight = weight_cache_.get(const_cast(weight.data())); - auto t_input_ids = - input_ids_cache_.get(const_cast(input_ids.data())); + auto t_indices = indices_cache_.get(const_cast(indices.data())); auto t_out = out_cache_.get(out.data()); if (!executor_) { - auto ret = aclnnEmbeddingGetWorkspaceSize(t_weight, t_input_ids, t_out, + auto ret = aclnnEmbeddingGetWorkspaceSize(t_weight, t_indices, t_out, &ws_size_, &executor_); assert(ret == ACL_SUCCESS && "`aclnnEmbeddingGetWorkspaceSize` failed"); aclSetAclOpExecutorRepeatable(executor_); } else { aclSetInputTensorAddr(executor_, 0, t_weight, const_cast(weight.data())); - aclSetInputTensorAddr(executor_, 1, t_input_ids, - const_cast(input_ids.data())); + aclSetInputTensorAddr(executor_, 1, t_indices, + const_cast(indices.data())); aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); } @@ -64,7 +70,7 @@ class Operator : public Embedding { } private: - mutable ascend::AclTensorCache input_ids_cache_; + mutable ascend::AclTensorCache indices_cache_; mutable ascend::AclTensorCache weight_cache_; diff --git a/src/native/ascend/ops/linear/kernel.h b/src/native/ascend/ops/linear/kernel.h index d6ce2a6da..faa995dd8 100644 --- a/src/native/ascend/ops/linear/kernel.h +++ b/src/native/ascend/ops/linear/kernel.h @@ -4,7 +4,6 @@ #include "acl/acl.h" #include "aclnn/aclnn_base.h" #include "aclnnop/aclnn_addmm.h" -#include "aclnnop/aclnn_baddbmm.h" #include "aclnnop/aclnn_matmul.h" #include "base/linear.h" #include "native/ascend/common.h" @@ -16,13 +15,16 @@ namespace infini::ops { template <> class Operator : public Linear { public: - Operator(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) - : Linear(a, b, bias, trans_a, trans_b, out), - batched_{out.ndim() > 2}, - a_cache_(a, trans_a), - b_cache_(b, trans_b), - out_cache_(out) { + Operator(const Tensor input, const Tensor weight, std::optional bias, + Tensor out) + : Linear(input, weight, bias, out), + input_cache_( + {static_cast(rows_), static_cast(input.size(-1))}, + ascend::ToAclDtype(input.dtype()), nullptr), + weight_cache_(weight, true), + out_cache_( + {static_cast(rows_), static_cast(weight.size(0))}, + ascend::ToAclDtype(out.dtype()), nullptr) { if (has_bias_) { bias_cache_ = ascend::AclTensorCache(*bias); alpha_scalar_ = aclCreateScalar(&alpha_storage_, ACL_FLOAT); @@ -35,59 +37,53 @@ class Operator : public Linear { // Null cached descriptors — see `AclTensorCache::release()`. bias_cache_.release(); - a_cache_.release(); - b_cache_.release(); + input_cache_.release(); + weight_cache_.release(); out_cache_.release(); if (alpha_scalar_) aclDestroyScalar(alpha_scalar_); if (beta_scalar_) aclDestroyScalar(beta_scalar_); } - void operator()(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) const override { + void operator()(const Tensor input, const Tensor weight, + std::optional bias, Tensor out) const override { auto stream = static_cast(stream_); - auto t_a = a_cache_.get(const_cast(a.data())); - auto t_b = b_cache_.get(const_cast(b.data())); + auto t_input = input_cache_.get(const_cast(input.data())); + auto t_weight = weight_cache_.get(const_cast(weight.data())); auto t_out = out_cache_.get(out.data()); if (has_bias_) { auto t_bias = bias_cache_.get(const_cast(bias->data())); if (!executor_) { - if (batched_) { - aclnnBaddbmmGetWorkspaceSize(t_bias, t_a, t_b, beta_scalar_, - alpha_scalar_, t_out, 0, &ws_size_, - &executor_); - } else { - aclnnAddmmGetWorkspaceSize(t_bias, t_a, t_b, beta_scalar_, - alpha_scalar_, t_out, 0, &ws_size_, - &executor_); - } + aclnnAddmmGetWorkspaceSize(t_bias, t_input, t_weight, beta_scalar_, + alpha_scalar_, t_out, 0, &ws_size_, + &executor_); aclSetAclOpExecutorRepeatable(executor_); } else { aclSetInputTensorAddr(executor_, 0, t_bias, const_cast(bias->data())); - aclSetInputTensorAddr(executor_, 1, t_a, const_cast(a.data())); - aclSetInputTensorAddr(executor_, 2, t_b, const_cast(b.data())); + aclSetInputTensorAddr(executor_, 1, t_input, + const_cast(input.data())); + aclSetInputTensorAddr(executor_, 2, t_weight, + const_cast(weight.data())); aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); } auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_); - if (batched_) { - aclnnBaddbmm(arena.buf, ws_size_, executor_, stream); - } else { - aclnnAddmm(arena.buf, ws_size_, executor_, stream); - } + aclnnAddmm(arena.buf, ws_size_, executor_, stream); } else { if (!executor_) { int8_t cube_math_type = 1; - aclnnMatmulGetWorkspaceSize(t_a, t_b, t_out, cube_math_type, &ws_size_, - &executor_); + aclnnMatmulGetWorkspaceSize(t_input, t_weight, t_out, cube_math_type, + &ws_size_, &executor_); aclSetAclOpExecutorRepeatable(executor_); } else { - aclSetInputTensorAddr(executor_, 0, t_a, const_cast(a.data())); - aclSetInputTensorAddr(executor_, 1, t_b, const_cast(b.data())); + aclSetInputTensorAddr(executor_, 0, t_input, + const_cast(input.data())); + aclSetInputTensorAddr(executor_, 1, t_weight, + const_cast(weight.data())); aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); } @@ -97,13 +93,11 @@ class Operator : public Linear { } private: - bool batched_; - mutable ascend::AclTensorCache bias_cache_; - mutable ascend::AclTensorCache a_cache_; + mutable ascend::AclTensorCache input_cache_; - mutable ascend::AclTensorCache b_cache_; + mutable ascend::AclTensorCache weight_cache_; mutable ascend::AclTensorCache out_cache_; diff --git a/src/native/ascend/ops/matmul/kernel.h b/src/native/ascend/ops/matmul/kernel.h index c574e4e55..bcf2d40aa 100644 --- a/src/native/ascend/ops/matmul/kernel.h +++ b/src/native/ascend/ops/matmul/kernel.h @@ -14,37 +14,39 @@ namespace infini::ops { template <> class Operator : public Matmul { public: - Operator(const Tensor a, const Tensor b, Tensor c, bool trans_a, bool trans_b) - : Matmul(a, b, c, trans_a, trans_b), - a_cache_(a, trans_a), - b_cache_(b, trans_b), - out_cache_(c) {} + Operator(const Tensor input, const Tensor other, Tensor out) + : Matmul(input, other, out), + input_cache_(input), + other_cache_(other), + out_cache_(out) {} ~Operator() { if (!ascend::IsAclRuntimeAlive()) return; // Null cached descriptors — see `AclTensorCache::release()`. - a_cache_.release(); - b_cache_.release(); + input_cache_.release(); + other_cache_.release(); out_cache_.release(); } - void operator()(const Tensor a, const Tensor b, Tensor c, bool trans_a, - bool trans_b) const override { + void operator()(const Tensor input, const Tensor other, + Tensor out) const override { auto stream = static_cast(stream_); - auto t_a = a_cache_.get(const_cast(a.data())); - auto t_b = b_cache_.get(const_cast(b.data())); - auto t_out = out_cache_.get(c.data()); + auto t_input = input_cache_.get(const_cast(input.data())); + auto t_other = other_cache_.get(const_cast(other.data())); + auto t_out = out_cache_.get(out.data()); if (!executor_) { int8_t cube_math_type = 1; - aclnnMatmulGetWorkspaceSize(t_a, t_b, t_out, cube_math_type, &ws_size_, - &executor_); + aclnnMatmulGetWorkspaceSize(t_input, t_other, t_out, cube_math_type, + &ws_size_, &executor_); aclSetAclOpExecutorRepeatable(executor_); } else { - aclSetInputTensorAddr(executor_, 0, t_a, const_cast(a.data())); - aclSetInputTensorAddr(executor_, 1, t_b, const_cast(b.data())); - aclSetOutputTensorAddr(executor_, 0, t_out, c.data()); + aclSetInputTensorAddr(executor_, 0, t_input, + const_cast(input.data())); + aclSetInputTensorAddr(executor_, 1, t_other, + const_cast(other.data())); + aclSetOutputTensorAddr(executor_, 0, t_out, out.data()); } auto& arena = ascend::GetWorkspacePool().Ensure(stream, ws_size_); @@ -52,9 +54,9 @@ class Operator : public Matmul { } private: - mutable ascend::AclTensorCache a_cache_; + mutable ascend::AclTensorCache input_cache_; - mutable ascend::AclTensorCache b_cache_; + mutable ascend::AclTensorCache other_cache_; mutable ascend::AclTensorCache out_cache_; diff --git a/src/native/ascend/ops/scaled_softmax_infinilm/kernel.h b/src/native/ascend/ops/scaled_softmax_infinilm/kernel.h new file mode 100644 index 000000000..c4d11f516 --- /dev/null +++ b/src/native/ascend/ops/scaled_softmax_infinilm/kernel.h @@ -0,0 +1,125 @@ +#ifndef INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ + +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_mul.h" +#include "aclnn_softmax.h" +#include "base/scaled_softmax_infinilm.h" +#include "data_type.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { + +template <> +class Operator + : public ScaledSoftmaxInfinilm { + public: + Operator(const Tensor input, double scale, Tensor out) + : ScaledSoftmaxInfinilm(input, scale, out), + in_cache_(input), + out_cache_(out), + temp_cache_(input), + scale_storage_(static_cast(scale)), + needs_scale_(std::fabs(scale - 1.0) > 1e-6) { + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32) && + "`ScaledSoftmaxInfinilm` Ascend path requires float16, bfloat16, or " + "float32 input"); + assert(input.IsContiguous() && + "`ScaledSoftmaxInfinilm` Ascend path requires contiguous input"); + assert(out.IsContiguous() && + "`ScaledSoftmaxInfinilm` Ascend path requires contiguous output"); + + temp_size_ = input.numel() * kDataTypeToSize.at(dtype_); + scale_scalar_ = aclCreateScalar(&scale_storage_, ACL_FLOAT); + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + in_cache_.release(); + out_cache_.release(); + temp_cache_.release(); + + if (scale_scalar_) aclDestroyScalar(scale_scalar_); + } + + void operator()(const Tensor input, double scale, Tensor out) const override { + assert(scale == scale_ && + "`ScaledSoftmaxInfinilm` scale changed after descriptor creation"); + + auto stream = static_cast(stream_); + auto t_in = in_cache_.get(const_cast(input.data())); + auto t_out = out_cache_.get(out.data()); + aclTensor* t_softmax_in = t_in; + void* softmax_in_data = const_cast(input.data()); + + if (needs_scale_) { + auto& temp = + ascend::GetWorkspacePool().Ensure(stream, temp_size_, "temp"); + auto t_temp = temp_cache_.get(temp.buf); + + if (!muls_exec_) { + aclnnMulsGetWorkspaceSize(t_in, scale_scalar_, t_temp, &muls_ws_, + &muls_exec_); + aclSetAclOpExecutorRepeatable(muls_exec_); + } else { + aclSetInputTensorAddr(muls_exec_, 0, t_in, + const_cast(input.data())); + aclSetOutputTensorAddr(muls_exec_, 0, t_temp, temp.buf); + } + + auto& muls_arena = ascend::GetWorkspacePool().Ensure(stream, muls_ws_); + aclnnMuls(muls_arena.buf, muls_ws_, muls_exec_, stream); + + t_softmax_in = t_temp; + softmax_in_data = temp.buf; + } + + if (!softmax_exec_) { + constexpr int64_t kLastDim = -1; + aclnnSoftmaxGetWorkspaceSize(t_softmax_in, kLastDim, t_out, &softmax_ws_, + &softmax_exec_); + aclSetAclOpExecutorRepeatable(softmax_exec_); + } else { + aclSetInputTensorAddr(softmax_exec_, 0, t_softmax_in, softmax_in_data); + aclSetOutputTensorAddr(softmax_exec_, 0, t_out, out.data()); + } + + auto& softmax_arena = + ascend::GetWorkspacePool().Ensure(stream, softmax_ws_); + aclnnSoftmax(softmax_arena.buf, softmax_ws_, softmax_exec_, stream); + } + + private: + mutable ascend::AclTensorCache in_cache_; + + mutable ascend::AclTensorCache out_cache_; + + mutable ascend::AclTensorCache temp_cache_; + + float scale_storage_{1.0f}; + + aclScalar* scale_scalar_ = nullptr; + + bool needs_scale_{false}; + + uint64_t temp_size_{0}; + + mutable aclOpExecutor* muls_exec_ = nullptr; + + mutable uint64_t muls_ws_ = 0; + + mutable aclOpExecutor* softmax_exec_ = nullptr; + + mutable uint64_t softmax_ws_ = 0; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ diff --git a/src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h b/src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h new file mode 100644 index 000000000..bf820a7ca --- /dev/null +++ b/src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h @@ -0,0 +1,291 @@ +#ifndef INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLE_INFINILM_KERNEL_H_ +#define INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLE_INFINILM_KERNEL_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnnop/aclnn_cast.h" +#include "aclnnop/aclnn_top_k_top_p_sample.h" +#include "base/top_k_top_p_sample_infinilm.h" +#include "data_type.h" +#include "native/ascend/common.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public TopKTopPSampleInfinilm { + public: + Operator(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, uint64_t offset, Tensor out) + : TopKTopPSampleInfinilm(logits, k, p, seed, offset, out) { + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16) && + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires float16 or " + "bfloat16 " + "logits"); + assert(logits.IsContiguous() && + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "logits"); + assert(out.IsContiguous() && + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "output"); + ValidateHostTensor(k); + ValidateHostTensor(p); + + logits_cache_ = ascend::AclTensorCache(logits); + top_k_cache_ = ascend::AclTensorCache({static_cast(batch_size_)}, + ACL_INT32, nullptr); + top_p_cache_ = ascend::AclTensorCache({static_cast(batch_size_)}, + ascend::ToAclDtype(dtype_), nullptr); + q_cache_ = ascend::AclTensorCache( + {static_cast(batch_size_), static_cast(vocab_size_)}, + ACL_FLOAT, nullptr); + selected_idx_cache_ = ascend::AclTensorCache( + {static_cast(batch_size_)}, ACL_INT64, nullptr); + selected_logits_cache_ = ascend::AclTensorCache( + {static_cast(batch_size_), static_cast(vocab_size_)}, + ACL_FLOAT, nullptr); + out_cache_ = ascend::AclTensorCache(out); + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + logits_cache_.release(); + top_k_cache_.release(); + top_p_cache_.release(); + q_cache_.release(); + selected_idx_cache_.release(); + selected_logits_cache_.release(); + out_cache_.release(); + } + + void operator()(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, uint64_t offset, + Tensor out) const override { + assert(logits.IsContiguous() && + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "logits"); + assert(out.IsContiguous() && + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "output"); + + auto stream = static_cast(stream_); + auto top_k_bytes = batch_size_ * kDataTypeToSize.at(DataType::kInt32); + auto top_p_bytes = batch_size_ * kDataTypeToSize.at(dtype_); + auto q_bytes = + batch_size_ * vocab_size_ * kDataTypeToSize.at(DataType::kFloat32); + auto selected_idx_bytes = + batch_size_ * kDataTypeToSize.at(DataType::kInt64); + auto selected_logits_bytes = + batch_size_ * vocab_size_ * kDataTypeToSize.at(DataType::kFloat32); + + FillParams(k, p, seed, offset); + + auto& top_k_arena = ascend::GetWorkspacePool().Ensure( + stream, top_k_bytes, "top_k_top_p_sample_top_k"); + auto& top_p_arena = ascend::GetWorkspacePool().Ensure( + stream, top_p_bytes, "top_k_top_p_sample_top_p"); + auto& q_arena = ascend::GetWorkspacePool().Ensure(stream, q_bytes, + "top_k_top_p_sample_q"); + auto ret = aclrtMemcpy(top_k_arena.buf, top_k_bytes, top_k_host_.data(), + top_k_bytes, ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "`TopKTopPSampleInfinilm`: copying `top_k` to Ascend failed"); + ret = aclrtMemcpy(top_p_arena.buf, top_p_bytes, top_p_host_.data(), + top_p_bytes, ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "`TopKTopPSampleInfinilm`: copying `top_p` to Ascend failed"); + ret = aclrtMemcpy(q_arena.buf, q_bytes, q_host_.data(), q_bytes, + ACL_MEMCPY_HOST_TO_DEVICE); + assert(ret == ACL_SUCCESS && + "`TopKTopPSampleInfinilm`: copying `q` to Ascend failed"); + + auto& selected_idx_arena = ascend::GetWorkspacePool().Ensure( + stream, selected_idx_bytes, "top_k_top_p_sample_idx"); + auto& selected_logits_arena = ascend::GetWorkspacePool().Ensure( + stream, selected_logits_bytes, "top_k_top_p_sample_logits"); + + auto t_logits = logits_cache_.get(const_cast(logits.data())); + auto t_top_k = top_k_cache_.get(top_k_arena.buf); + auto t_top_p = top_p_cache_.get(top_p_arena.buf); + auto t_q = q_cache_.get(q_arena.buf); + auto t_selected_idx = selected_idx_cache_.get(selected_idx_arena.buf); + auto t_selected_logits = + selected_logits_cache_.get(selected_logits_arena.buf); + + if (!sample_exec_) { + ret = aclnnTopKTopPSampleGetWorkspaceSize( + t_logits, t_top_k, t_top_p, t_q, /*eps=*/1e-8, /*isNeedLogits=*/false, + /*topKGuess=*/32, t_selected_idx, t_selected_logits, &sample_ws_size_, + &sample_exec_); + assert(ret == ACL_SUCCESS && + "`aclnnTopKTopPSampleGetWorkspaceSize` failed"); + aclSetAclOpExecutorRepeatable(sample_exec_); + } else { + aclSetInputTensorAddr(sample_exec_, 0, t_logits, + const_cast(logits.data())); + aclSetInputTensorAddr(sample_exec_, 1, t_top_k, top_k_arena.buf); + aclSetInputTensorAddr(sample_exec_, 2, t_top_p, top_p_arena.buf); + aclSetInputTensorAddr(sample_exec_, 3, t_q, q_arena.buf); + aclSetOutputTensorAddr(sample_exec_, 0, t_selected_idx, + selected_idx_arena.buf); + aclSetOutputTensorAddr(sample_exec_, 1, t_selected_logits, + selected_logits_arena.buf); + } + + auto& sample_ws_arena = ascend::GetWorkspacePool().Ensure( + stream, sample_ws_size_, "top_k_top_p_sample_workspace"); + ret = aclnnTopKTopPSample(sample_ws_arena.buf, sample_ws_size_, + sample_exec_, stream); + assert(ret == ACL_SUCCESS && "`aclnnTopKTopPSample` failed"); + + CastSelectedIdx(selected_idx_arena.buf, out); + } + + private: + void ValidateHostTensor(std::optional tensor) const { + if (!tensor.has_value()) return; + + assert(tensor->device().type() == Device::Type::kCpu && + "`TopKTopPSampleInfinilm` Ascend path currently requires host-side " + "`k`/`p` tensors"); + assert(tensor->IsContiguous() && + "`TopKTopPSampleInfinilm` Ascend path requires contiguous `k`/`p` " + "tensors"); + } + + void CastSelectedIdx(void* selected_idx, Tensor out) const { + auto stream = static_cast(stream_); + auto t_selected_idx = selected_idx_cache_.get(selected_idx); + auto t_out = out_cache_.get(out.data()); + + if (!cast_exec_) { + auto ret = aclnnCastGetWorkspaceSize(t_selected_idx, ACL_INT32, t_out, + &cast_ws_size_, &cast_exec_); + assert(ret == ACL_SUCCESS && "`aclnnCastGetWorkspaceSize` failed"); + aclSetAclOpExecutorRepeatable(cast_exec_); + } else { + aclSetInputTensorAddr(cast_exec_, 0, t_selected_idx, selected_idx); + aclSetOutputTensorAddr(cast_exec_, 0, t_out, out.data()); + } + + auto& cast_ws_arena = ascend::GetWorkspacePool().Ensure( + stream, cast_ws_size_, "top_k_top_p_sample_cast_workspace"); + auto ret = aclnnCast(cast_ws_arena.buf, cast_ws_size_, cast_exec_, stream); + assert(ret == ACL_SUCCESS && "`aclnnCast` failed"); + } + + void FillParams(std::optional k, std::optional p, + uint64_t seed, uint64_t offset) const { + top_k_host_.resize(batch_size_); + top_p_host_.resize(batch_size_ * kDataTypeToSize.at(dtype_)); + q_host_.resize(batch_size_ * vocab_size_); + std::exponential_distribution dist(1.0F); + std::mt19937_64 rng(seed); + rng.discard(offset); + + for (Tensor::Size row = 0; row < batch_size_; ++row) { + top_k_host_[row] = static_cast(GetK(k, row)); + auto value = static_cast(GetP(p, row)); + auto* dst = top_p_host_.data() + row * kDataTypeToSize.at(dtype_); + + if (dtype_ == DataType::kFloat16) { + auto converted = Float16::FromFloat(value); + std::memcpy(dst, &converted, sizeof(converted)); + } else { + auto converted = BFloat16::FromFloat(value); + std::memcpy(dst, &converted, sizeof(converted)); + } + } + + // CANN samples by taking argmax of softmax probabilities divided by + // exponential noise. + for (auto& value : q_host_) { + value = dist(rng); + } + } + + int64_t GetK(std::optional k, Tensor::Size row) const { + if (!k.has_value()) return static_cast(vocab_size_); + + const auto offset = k->size(0) == 1 ? 0 : row; + int64_t value = 0; + if (k->dtype() == DataType::kInt32) { + value = static_cast(k->data())[offset]; + } else { + value = static_cast(k->data())[offset]; + } + + if (value <= 0) return static_cast(vocab_size_); + return std::min(value, static_cast(vocab_size_)); + } + + double GetP(std::optional p, Tensor::Size row) const { + if (!p.has_value()) return 1.0; + + const auto offset = p->size(0) == 1 ? 0 : row; + double value = 1.0; + switch (p->dtype()) { + case DataType::kFloat16: + value = static_cast(p->data())[offset].ToFloat(); + break; + case DataType::kBFloat16: + value = static_cast(p->data())[offset].ToFloat(); + break; + case DataType::kFloat32: + value = static_cast(p->data())[offset]; + break; + case DataType::kFloat64: + value = static_cast(p->data())[offset]; + break; + default: + assert(false && "`TopKTopPSampleInfinilm` has unsupported `p` dtype"); + } + + if (value <= 0.0 || value > 1.0) return 1.0; + return value; + } + + mutable ascend::AclTensorCache logits_cache_; + + mutable ascend::AclTensorCache top_k_cache_; + + mutable ascend::AclTensorCache top_p_cache_; + + mutable ascend::AclTensorCache q_cache_; + + mutable ascend::AclTensorCache selected_idx_cache_; + + mutable ascend::AclTensorCache selected_logits_cache_; + + mutable ascend::AclTensorCache out_cache_; + + mutable std::vector top_k_host_; + + mutable std::vector top_p_host_; + + mutable std::vector q_host_; + + mutable aclOpExecutor* sample_exec_ = nullptr; + + mutable uint64_t sample_ws_size_ = 0; + + mutable aclOpExecutor* cast_exec_ = nullptr; + + mutable uint64_t cast_ws_size_ = 0; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLE_INFINILM_KERNEL_H_ diff --git a/src/native/cpu/ops/add/add.h b/src/native/cpu/ops/add/add.h index 8580c6df0..7528a9a5c 100644 --- a/src/native/cpu/ops/add/add.h +++ b/src/native/cpu/ops/add/add.h @@ -13,25 +13,30 @@ template <> class Operator : public Add, Caster { public: - Operator(const Tensor input, const Tensor other, Tensor out) - : Add{input, other, out} { + Operator(const Tensor input, const Tensor other, const double alpha, + Tensor out) + : Add{input, other, alpha, out} { // TODO: Check constraints. } - void operator()(const Tensor input, const Tensor other, + Operator(const Tensor input, const Tensor other, Tensor out) + : Operator{input, other, 1.0, out} {} + + void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { DispatchFunc( out_type_, [&](auto tag) { using T = typename decltype(tag)::type; - Compute(input, other, out); + Compute(input, other, alpha, out); }, "`Operator::operator()`"); } private: template - void Compute(const Tensor input, const Tensor other, Tensor out) const { + void Compute(const Tensor input, const Tensor other, const double alpha, + Tensor out) const { using ComputeType = std::conditional_t || IsFP16, float, T>; @@ -55,7 +60,8 @@ class Operator : public Add, out_strides_.data()); out_ptr[out_idx] = Cast(Cast(input_ptr[input_idx]) + - Cast(other_ptr[other_idx])); + Cast(alpha) * + Cast(other_ptr[other_idx])); } } }; diff --git a/src/native/cpu/ops/cat/cat.h b/src/native/cpu/ops/cat/cat.h index a2f600e5c..d7e99ad3c 100644 --- a/src/native/cpu/ops/cat/cat.h +++ b/src/native/cpu/ops/cat/cat.h @@ -5,63 +5,44 @@ #include #include "base/cat.h" -#include "native/cpu/caster_.h" namespace infini::ops { template <> class Operator : public Cat { public: - Operator(const Tensor first_input, std::vector rest_inputs, - int64_t dim, Tensor out) - : Cat{first_input, rest_inputs, dim, out} {} + Operator(const std::vector tensors, const int64_t dim, Tensor out) + : Cat{tensors, dim, out} {} - void operator()(const Tensor first_input, std::vector rest_inputs, - int64_t /*dim*/, Tensor out) const override { - // Collect all input tensors. - std::vector inputs; - inputs.reserve(input_count_); - inputs.push_back(&first_input); - for (const auto& t : rest_inputs) { - inputs.push_back(&t); - } - - // Use normalized `dim_` from base class (handles negative dim). - auto dim = dim_; + void operator()(const std::vector tensors, const int64_t /*dim*/, + Tensor out) const override { auto elem_size = kDataTypeToSize.at(out.dtype()); - auto ndim = out.ndim(); - auto out_shape = out.shape(); - - // Compute outer and inner sizes relative to the cat dimension. - Tensor::Size outer = 1; - for (int64_t i = 0; i < dim; ++i) { - outer *= out_shape[i]; - } - - Tensor::Size inner = 1; - for (size_t i = static_cast(dim) + 1; i < ndim; ++i) { - inner *= out_shape[i]; - } - auto* out_ptr = static_cast(out.data()); - Tensor::Size out_dim_size = out_shape[dim]; - - // For each outer index, copy slices from each input along the cat dim. - for (Tensor::Size o = 0; o < outer; ++o) { - Tensor::Size offset_in_dim = 0; - - for (size_t t = 0; t < input_count_; ++t) { - auto in_dim = inputs[t]->shape()[dim]; - auto in_ptr = static_cast(inputs[t]->data()); - - auto src_offset = (o * in_dim) * inner * elem_size; - auto dst_offset = - (o * out_dim_size + offset_in_dim) * inner * elem_size; - auto copy_size = in_dim * inner * elem_size; - - std::memcpy(out_ptr + dst_offset, in_ptr + src_offset, copy_size); - offset_in_dim += in_dim; + Tensor::Size cat_offset = 0; + + for (const auto& tensor : tensors) { + auto* input_ptr = static_cast(tensor.data()); + + for (Tensor::Size flat_index = 0; flat_index < tensor.numel(); + ++flat_index) { + auto remaining = flat_index; + Tensor::Stride input_offset = 0; + Tensor::Stride output_offset = 0; + + for (Tensor::Size axis = tensor.ndim(); axis-- > 0;) { + auto coordinate = remaining % tensor.size(axis); + remaining /= tensor.size(axis); + input_offset += coordinate * tensor.stride(axis); + output_offset += + (coordinate + + (axis == static_cast(dim_) ? cat_offset : 0)) * + out.stride(axis); + } + + std::memcpy(out_ptr + output_offset * elem_size, + input_ptr + input_offset * elem_size, elem_size); } + cat_offset += tensor.size(dim_); } } }; diff --git a/src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h b/src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h new file mode 100644 index 000000000..94d8709b5 --- /dev/null +++ b/src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h @@ -0,0 +1,84 @@ +#ifndef INFINI_OPS_CPU_CAUSAL_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_CPU_CAUSAL_SOFTMAX_INFINILM_H_ + +#include + +#include "base/causal_softmax_infinilm.h" +#include "common/generic_utils.h" +#include "data_type.h" +#include "native/cpu/caster_.h" +#include "tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public CausalSoftmaxInfinilm, Caster { + public: + Operator(const Tensor input, Tensor out) + : CausalSoftmaxInfinilm{input, out} {} + + void operator()(const Tensor input, Tensor out) const override { + DispatchFunc( + out.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(input, out); + }, + "`Operator::operator()`"); + } + + private: + template + void Compute(const Tensor input, Tensor out) const { + auto* out_ptr = static_cast(out.data()); + const auto* input_ptr = static_cast(input.data()); + + auto out_stride_b = ndim_ == 3 ? out_strides_[0] : 0; + auto out_stride_i = out_strides_[ndim_ - 2]; + auto out_stride_j = out_strides_[ndim_ - 1]; + auto input_stride_b = ndim_ == 3 ? input_strides_[0] : 0; + auto input_stride_i = input_strides_[ndim_ - 2]; + auto input_stride_j = input_strides_[ndim_ - 1]; + + for (Tensor::Size bi = 0; bi < batch_size_; ++bi) { + for (Tensor::Size i = 0; i < seq_len_; ++i) { + ptrdiff_t out_offset = bi * out_stride_b + i * out_stride_i; + ptrdiff_t input_offset = bi * input_stride_b + i * input_stride_i; + T* out_row = out_ptr + out_offset; + const T* input_row = input_ptr + input_offset; + + Tensor::Size valid_len = total_seq_len_ - seq_len_ + i + 1; + + for (Tensor::Size j = valid_len; j < total_seq_len_; ++j) { + out_row[j * out_stride_j] = Cast(0.0f); + } + + float max_val = Cast(input_row[0]); + for (Tensor::Size j = 1; j < valid_len; ++j) { + float v = Cast(input_row[j * input_stride_j]); + if (v > max_val) { + max_val = v; + } + } + + float sum = 0.0f; + for (Tensor::Size j = 0; j < valid_len; ++j) { + float v = + std::exp(Cast(input_row[j * input_stride_j]) - max_val); + out_row[j * out_stride_j] = Cast(v); + sum += v; + } + + for (Tensor::Size j = 0; j < valid_len; ++j) { + out_row[j * out_stride_j] = + Cast(Cast(out_row[j * out_stride_j]) / sum); + } + } + } + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cpu/ops/linear/linear.h b/src/native/cpu/ops/linear/linear.h index aa152658a..638fac1e7 100644 --- a/src/native/cpu/ops/linear/linear.h +++ b/src/native/cpu/ops/linear/linear.h @@ -13,90 +13,64 @@ template <> class Operator : public Linear, Caster { public: - Operator(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) - : Linear{a, b, bias, trans_a, trans_b, out} {} + Operator(const Tensor input, const Tensor weight, std::optional bias, + Tensor out) + : Linear{input, weight, bias, out} {} - void operator()(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) const override { + void operator()(const Tensor input, const Tensor weight, + std::optional bias, Tensor out) const override { DispatchFunc( out.dtype(), [&](auto tag) { using T = typename decltype(tag)::type; - Compute(a, b, bias, trans_a, trans_b, out); + Compute(input, weight, bias, out); }, "`Operator::operator()`"); } private: template - void Compute(const Tensor a, const Tensor b, std::optional bias, - bool trans_a, bool trans_b, Tensor out) const { - const auto* a_ptr = static_cast(a.data()); - const auto* b_ptr = static_cast(b.data()); + void Compute(const Tensor input, const Tensor weight, + std::optional bias, Tensor out) const { + const auto* input_ptr = static_cast(input.data()); + const auto* weight_ptr = static_cast(weight.data()); auto* out_ptr = static_cast(out.data()); const T* bias_ptr = bias ? static_cast(bias->data()) : nullptr; - auto ndim_a = a_shape_.size(); - auto ndim_b = b_shape_.size(); - auto ndim_out = out_shape_.size(); + auto in_features = input.size(-1); + auto out_features = weight.size(0); - Tensor::Size m = out_shape_[ndim_out - 2]; - Tensor::Size n = out_shape_[ndim_out - 1]; - Tensor::Size k = trans_a ? a_shape_[ndim_a - 2] : a_shape_[ndim_a - 1]; + for (Tensor::Size row = 0; row < rows_; ++row) { + auto remaining = row; + Tensor::Stride input_base = 0; + Tensor::Stride output_base = 0; - // Compute strides for the inner matrix dimensions after transpose. - Tensor::Stride stride_a_m = - trans_a ? a_strides_[ndim_a - 1] : a_strides_[ndim_a - 2]; - Tensor::Stride stride_a_k = - trans_a ? a_strides_[ndim_a - 2] : a_strides_[ndim_a - 1]; - Tensor::Stride stride_b_k = - trans_b ? b_strides_[ndim_b - 1] : b_strides_[ndim_b - 2]; - Tensor::Stride stride_b_n = - trans_b ? b_strides_[ndim_b - 2] : b_strides_[ndim_b - 1]; - Tensor::Stride stride_out_m = out_strides_[ndim_out - 2]; - Tensor::Stride stride_out_n = out_strides_[ndim_out - 1]; - - // Batch dimensions. - Tensor::Size batch_count = 1; - for (size_t i = 0; i + 2 < ndim_out; ++i) { - batch_count *= out_shape_[i]; - } - - Tensor::Stride batch_stride_a = ndim_a > 2 ? a_strides_[ndim_a - 3] : 0; - Tensor::Stride batch_stride_b = ndim_b > 2 ? b_strides_[ndim_b - 3] : 0; - Tensor::Stride batch_stride_out = - ndim_out > 2 ? out_strides_[ndim_out - 3] : 0; - - // Bias stride: for 1D bias `[n]`, stride is 1. For batched bias, use last - // stride. - Tensor::Stride bias_stride = 0; - if (bias_ptr) { - auto ndim_bias = bias->shape().size(); - bias_stride = bias->strides()[ndim_bias - 1]; - } - - for (Tensor::Size batch = 0; batch < batch_count; ++batch) { - const auto* a_batch = a_ptr + batch * batch_stride_a; - const auto* b_batch = b_ptr + batch * batch_stride_b; - auto* out_batch = out_ptr + batch * batch_stride_out; - - for (Tensor::Size i = 0; i < m; ++i) { - for (Tensor::Size j = 0; j < n; ++j) { - float sum = 0.0f; - - for (Tensor::Size l = 0; l < k; ++l) { - float a_val = Cast(a_batch[i * stride_a_m + l * stride_a_k]); - float b_val = Cast(b_batch[l * stride_b_k + j * stride_b_n]); - sum += a_val * b_val; - } + for (Tensor::Size axis = input.ndim() - 1; axis > 0; --axis) { + auto leading_axis = axis - 1; + auto coordinate = remaining % input.size(leading_axis); + remaining /= input.size(leading_axis); + input_base += coordinate * input.stride(leading_axis); + output_base += coordinate * out.stride(leading_axis); + } - if (bias_ptr) { - sum += Cast(bias_ptr[j * bias_stride]); - } + for (Tensor::Size output_feature = 0; output_feature < out_features; + ++output_feature) { + float sum = 0.0f; + + for (Tensor::Size input_feature = 0; input_feature < in_features; + ++input_feature) { + auto input_value = + input_ptr[input_base + input_feature * input.stride(-1)]; + auto weight_value = weight_ptr[output_feature * weight.stride(0) + + input_feature * weight.stride(1)]; + sum += Cast(input_value) * Cast(weight_value); + } - out_batch[i * stride_out_m + j * stride_out_n] = Cast(sum); + if (bias_ptr) { + sum += Cast(bias_ptr[output_feature * bias->stride(0)]); } + + out_ptr[output_base + output_feature * out.stride(-1)] = Cast(sum); } } } diff --git a/src/native/cpu/ops/silu_and_mul/silu_and_mul.h b/src/native/cpu/ops/silu_and_mul/silu_and_mul.h new file mode 100644 index 000000000..2da25e506 --- /dev/null +++ b/src/native/cpu/ops/silu_and_mul/silu_and_mul.h @@ -0,0 +1,67 @@ +#ifndef INFINI_OPS_CPU_SILU_AND_MUL_SILU_AND_MUL_H_ +#define INFINI_OPS_CPU_SILU_AND_MUL_SILU_AND_MUL_H_ + +#include + +#include "base/silu_and_mul.h" +#include "common/generic_utils.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +template <> +class Operator : public SiluAndMul, + Caster { + public: + using SiluAndMul::SiluAndMul; + + void operator()(const Tensor input, Tensor out) const override { + DispatchFunc( + out_type_, + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(input, out); + }, + "Operator::operator()"); + } + + private: + template + void Compute(const Tensor input, Tensor out) const { + using ComputeType = std::conditional_t || + IsFP16, + float, T>; + + const auto* input_ptr = static_cast(input.data()); + auto* out_ptr = static_cast(out.data()); + + auto get_idx = [&](Tensor::Size i, bool is_contig, const auto* shape, + const auto* strides) { + return is_contig ? i : utils::IndexToOffset(i, ndim_, shape, strides); + }; + +#pragma omp parallel for + for (Tensor::Size i = 0; i < output_size_; ++i) { + const auto row = i / hidden_size_; + const auto column = i % hidden_size_; + const auto gate_logical_idx = row * 2 * hidden_size_ + column; + const auto up_logical_idx = gate_logical_idx + hidden_size_; + auto gate_idx = get_idx(gate_logical_idx, is_input_contiguous_, + input_shape_.data(), input_strides_.data()); + auto up_idx = get_idx(up_logical_idx, is_input_contiguous_, + input_shape_.data(), input_strides_.data()); + auto out_idx = get_idx(i, is_out_contiguous_, out_shape_.data(), + out_strides_.data()); + const ComputeType gate_val = Cast(input_ptr[gate_idx]); + const ComputeType sigmoid_gate = static_cast( + 1.0 / (1.0 + std::exp(-static_cast(gate_val)))); + const ComputeType swish_gate = gate_val * sigmoid_gate; + out_ptr[out_idx] = + Cast(swish_gate * Cast(input_ptr[up_idx])); + } + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h b/src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h new file mode 100644 index 000000000..409fc960b --- /dev/null +++ b/src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h @@ -0,0 +1,163 @@ +#ifndef INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLE_INFINILM_H_ +#define INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLE_INFINILM_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "base/top_k_top_p_sample_infinilm.h" +#include "data_type.h" +#include "native/cpu/caster_.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public TopKTopPSampleInfinilm, Caster { + public: + Operator(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, uint64_t offset, Tensor out) + : TopKTopPSampleInfinilm(logits, k, p, seed, offset, out) {} + + void operator()(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, uint64_t offset, + Tensor out) const override { + DispatchFunc( + logits.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(logits, k, p, seed, offset, out); + }, + "`Operator::operator()`"); + } + + private: + template + void Compute(const Tensor logits, std::optional k, + std::optional p, uint64_t seed, uint64_t offset, + Tensor out) const { + const auto* logits_ptr = static_cast(logits.data()); + auto* out_ptr = static_cast(out.data()); + + for (Tensor::Size row = 0; row < batch_size_; ++row) { + const int64_t top_k = GetK(k, row); + const double top_p = GetP(p, row); + out_ptr[row * out.stride(0)] = + SampleRow(logits_ptr + row * logits.stride(0), logits.stride(1), + top_k, top_p, seed, offset + row); + } + } + + template + int32_t SampleRow(const T* row, Tensor::Stride stride, int64_t top_k, + double top_p, uint64_t seed, uint64_t offset) const { + std::vector indices(vocab_size_); + std::iota(indices.begin(), indices.end(), 0); + std::sort(indices.begin(), indices.end(), [&](int64_t a, int64_t b) { + return Cast(row[a * stride]) > Cast(row[b * stride]); + }); + + Tensor::Size keep_count = NormalizeTopK(top_k); + if (top_p > 0.0 && top_p < 1.0) { + keep_count = ApplyTopP(row, stride, top_p, indices, keep_count); + } + + if (keep_count == 1) { + return static_cast(indices[0]); + } + + std::vector weights(keep_count); + float max_val = -std::numeric_limits::infinity(); + for (Tensor::Size i = 0; i < keep_count; ++i) { + const auto v = Cast(row[indices[i] * stride]); + if (v > max_val) max_val = v; + } + + for (Tensor::Size i = 0; i < keep_count; ++i) { + weights[i] = std::exp(Cast(row[indices[i] * stride]) - max_val); + } + + std::discrete_distribution dist(weights.begin(), + weights.end()); + std::mt19937_64 rng(seed); + rng.discard(offset); + return static_cast(indices[dist(rng)]); + } + + template + Tensor::Size ApplyTopP(const T* row, Tensor::Stride stride, double top_p, + const std::vector& indices, + Tensor::Size keep_count) const { + float max_val = -std::numeric_limits::infinity(); + for (Tensor::Size i = 0; i < keep_count; ++i) { + const auto v = Cast(row[indices[i] * stride]); + if (v > max_val) max_val = v; + } + + double sum = 0.0; + std::vector probs(keep_count); + for (Tensor::Size i = 0; i < keep_count; ++i) { + probs[i] = std::exp(Cast(row[indices[i] * stride]) - max_val); + sum += probs[i]; + } + + double cumulative = 0.0; + for (Tensor::Size i = 0; i < keep_count; ++i) { + cumulative += probs[i] / sum; + if (cumulative >= top_p) { + return i + 1; + } + } + + return keep_count; + } + + Tensor::Size NormalizeTopK(int64_t top_k) const { + if (top_k <= 0 || static_cast(top_k) > + static_cast(vocab_size_)) { + return vocab_size_; + } + return static_cast(top_k); + } + + int64_t GetK(std::optional k, Tensor::Size row) const { + if (!k.has_value()) return static_cast(vocab_size_); + + const auto offset = (k->size(0) == 1 ? 0 : row) * k->stride(0); + if (k->dtype() == DataType::kInt32) { + return static_cast(k->data())[offset]; + } + return static_cast(k->data())[offset]; + } + + double GetP(std::optional p, Tensor::Size row) const { + if (!p.has_value()) return 1.0; + + const auto offset = (p->size(0) == 1 ? 0 : row) * p->stride(0); + switch (p->dtype()) { + case DataType::kFloat16: + return Cast(static_cast(p->data())[offset]); + case DataType::kBFloat16: + return Cast(static_cast(p->data())[offset]); + case DataType::kFloat32: + return static_cast(p->data())[offset]; + case DataType::kFloat64: + return static_cast(p->data())[offset]; + default: + assert(false && "`TopKTopPSampleInfinilm` has unsupported `p` dtype."); + return 1.0; + } + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLE_INFINILM_H_ diff --git a/src/native/cuda/iluvatar/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/iluvatar/ops/causal_softmax_infinilm/kernel.h new file mode 100644 index 000000000..dc3669efe --- /dev/null +++ b/src/native/cuda/iluvatar/ops/causal_softmax_infinilm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmaxInfinilm> { + public: + using CudaCausalSoftmaxInfinilm< + Runtime>::CudaCausalSoftmaxInfinilm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/iluvatar/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/iluvatar/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..b2ac7c91a --- /dev/null +++ b/src/native/cuda/iluvatar/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_ILUVATAR_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/fused_add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFusedAddRmsNorm> { + public: + using CudaFusedAddRmsNorm< + Runtime>::CudaFusedAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/iluvatar/ops/silu_and_mul/kernel.h b/src/native/cuda/iluvatar/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..b4a8fc18f --- /dev/null +++ b/src/native/cuda/iluvatar/ops/silu_and_mul/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_ILUVATAR_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_SILU_AND_MUL_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/silu_and_mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSiluAndMul> { + public: + using CudaSiluAndMul>::CudaSiluAndMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/metax/ops/causal_softmax_infinilm/kernel.h new file mode 100644 index 000000000..bc191034c --- /dev/null +++ b/src/native/cuda/metax/ops/causal_softmax_infinilm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_METAX_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_METAX_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmaxInfinilm> { + public: + using CudaCausalSoftmaxInfinilm< + Runtime>::CudaCausalSoftmaxInfinilm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/metax/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..c79d67df0 --- /dev/null +++ b/src/native/cuda/metax/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_METAX_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/fused_add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFusedAddRmsNorm> { + public: + using CudaFusedAddRmsNorm>::CudaFusedAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/silu_and_mul/kernel.h b/src/native/cuda/metax/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..6b9ba9f58 --- /dev/null +++ b/src/native/cuda/metax/ops/silu_and_mul/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_METAX_SILU_AND_MUL_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/silu_and_mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSiluAndMul> { + public: + using CudaSiluAndMul>::CudaSiluAndMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/moore/ops/causal_softmax_infinilm/kernel.h new file mode 100644 index 000000000..f8bff9aa5 --- /dev/null +++ b/src/native/cuda/moore/ops/causal_softmax_infinilm/kernel.h @@ -0,0 +1,33 @@ +#ifndef INFINI_OPS_MOORE_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_MOORE_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ + +// clang-format off +#include +// clang-format on + +// clang-format off +#include "native/cuda/moore/device_.h" +// clang-format on + +#include "native/cuda/moore/caster.cuh" +#include "native/cuda/moore/runtime_.h" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.h" + +namespace infini::ops { + +struct MooreCausalSoftmaxInfinilmBackend : Runtime { + // Moore's causal softmax kernel should not dispatch block sizes above 1024. + static constexpr int max_block_size = 1024; +}; + +template <> +class Operator + : public CudaCausalSoftmaxInfinilm { + public: + using CudaCausalSoftmaxInfinilm< + MooreCausalSoftmaxInfinilmBackend>::CudaCausalSoftmaxInfinilm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/moore/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..fccbee43b --- /dev/null +++ b/src/native/cuda/moore/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,25 @@ +#ifndef INFINI_OPS_MOORE_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_MOORE_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include + +// clang-format off +#include +// clang-format on + +#include "native/cuda/moore/caster.cuh" +#include "native/cuda/moore/runtime_.h" +#include "native/cuda/ops/fused_add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFusedAddRmsNorm> { + public: + using CudaFusedAddRmsNorm>::CudaFusedAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/silu_and_mul/kernel.h b/src/native/cuda/moore/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..0ac8815bc --- /dev/null +++ b/src/native/cuda/moore/ops/silu_and_mul/kernel.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_MOORE_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_MOORE_SILU_AND_MUL_KERNEL_H_ + +#include + +// clang-format off +#include "native/cuda/moore/polyfills.cuh" +// clang-format on + +#include "native/cuda/moore/caster.cuh" +#include "native/cuda/moore/polyfills.cuh" +#include "native/cuda/moore/runtime_.h" +#include "native/cuda/ops/silu_and_mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSiluAndMul> { + public: + using CudaSiluAndMul>::CudaSiluAndMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/nvidia/ops/causal_softmax_infinilm/kernel.h new file mode 100644 index 000000000..2d85028f7 --- /dev/null +++ b/src/native/cuda/nvidia/ops/causal_softmax_infinilm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmaxInfinilm> { + public: + using CudaCausalSoftmaxInfinilm< + Runtime>::CudaCausalSoftmaxInfinilm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..c4664ea8b --- /dev/null +++ b/src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_NVIDIA_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/fused_add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaFusedAddRmsNorm> { + public: + using CudaFusedAddRmsNorm< + Runtime>::CudaFusedAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.cuh b/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.cuh new file mode 100644 index 000000000..2ccf2e71d --- /dev/null +++ b/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.cuh @@ -0,0 +1,76 @@ +#ifndef INFINI_OPS_NVIDIA_RESHAPE_AND_CACHE_KERNEL_CUH_ +#define INFINI_OPS_NVIDIA_RESHAPE_AND_CACHE_KERNEL_CUH_ + +#include + +#include +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/nvidia/caster.cuh" + +namespace infini::ops { + +template +__device__ __forceinline__ uint8_t ToFp8Cache(T value, float scale, + bool use_e5m2) { + float converted = + Caster::template Cast(value) / scale; + auto format = use_e5m2 ? __NV_E5M2 : __NV_E4M3; + return static_cast( + __nv_cvt_float_to_fp8(converted, __NV_SATFINITE, format)); +} + +template +__global__ void ReshapeAndCacheKernel( + const T* __restrict__ key, const T* __restrict__ value, + TCache* __restrict__ key_cache, TCache* __restrict__ value_cache, + const int64_t* __restrict__ slot_mapping, int64_t key_token_stride, + int64_t value_token_stride, int64_t key_head_stride, + int64_t value_head_stride, size_t num_heads, size_t head_size, + size_t block_size, size_t x, const float* __restrict__ k_scale, + const float* __restrict__ v_scale) { + const size_t token_idx = blockIdx.x; + const int64_t slot_idx = slot_mapping[token_idx]; + if (slot_idx < 0) { + return; + } + + const size_t block_idx = static_cast(slot_idx) / block_size; + const size_t block_offset = static_cast(slot_idx) % block_size; + const size_t head_blocks = head_size / x; + + for (size_t i = threadIdx.x; i < num_heads * head_size; i += blockDim.x) { + const size_t head_idx = i / head_size; + const size_t head_offset = i % head_size; + const size_t head_block = head_offset / x; + const size_t x_offset = head_offset % x; + const size_t key_src_idx = + token_idx * key_token_stride + head_idx * key_head_stride + head_offset; + const size_t value_src_idx = token_idx * value_token_stride + + head_idx * value_head_stride + head_offset; + const size_t key_dst_idx = + (((block_idx * num_heads + head_idx) * head_blocks + head_block) * + block_size + + block_offset) * + x + + x_offset; + const size_t value_dst_idx = + ((block_idx * num_heads + head_idx) * head_size + head_offset) * + block_size + + block_offset; + + if constexpr (kQuantized) { + key_cache[key_dst_idx] = ToFp8Cache(key[key_src_idx], *k_scale, kUseE5M2); + value_cache[value_dst_idx] = + ToFp8Cache(value[value_src_idx], *v_scale, kUseE5M2); + } else { + key_cache[key_dst_idx] = key[key_src_idx]; + value_cache[value_dst_idx] = value[value_src_idx]; + } + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.h b/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.h new file mode 100644 index 000000000..096009059 --- /dev/null +++ b/src/native/cuda/nvidia/ops/reshape_and_cache/kernel.h @@ -0,0 +1,73 @@ +#ifndef INFINI_OPS_NVIDIA_RESHAPE_AND_CACHE_KERNEL_H_ +#define INFINI_OPS_NVIDIA_RESHAPE_AND_CACHE_KERNEL_H_ + +#include +#include +#include + +#include "base/reshape_and_cache.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/nvidia/ops/reshape_and_cache/kernel.cuh" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { + +template <> +class Operator + : public ReshapeAndCache { + public: + using ReshapeAndCache::ReshapeAndCache; + + void operator()(const Tensor key, const Tensor value, Tensor key_cache, + Tensor value_cache, const Tensor slot_mapping, + const std::string kv_cache_dtype, const Tensor k_scale, + const Tensor v_scale) const override { + if (num_tokens_ == 0) { + return; + } + + auto cuda_stream = static_cast::Stream>( + stream_ ? stream_ : 0); + const bool quantized = kv_cache_dtype != "auto"; + const bool use_e5m2 = kv_cache_dtype == "fp8_e5m2"; + + DispatchFunc, ReducedFloatTypes>>( + key.dtype(), + [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + auto launch = [&](auto cache_tag, auto quantized_tag, auto e5m2_tag) { + using TCache = typename decltype(cache_tag)::type; + constexpr bool kQuantized = decltype(quantized_tag)::value; + constexpr bool kUseE5M2 = decltype(e5m2_tag)::value; + ReshapeAndCacheKernel + <<(num_tokens_), 256, 0, cuda_stream>>>( + reinterpret_cast(key.data()), + reinterpret_cast(value.data()), + reinterpret_cast(key_cache.data()), + reinterpret_cast(value_cache.data()), + reinterpret_cast(slot_mapping.data()), + key_strides_[0], value_strides_[0], key_strides_[1], + value_strides_[1], num_heads_, head_size_, block_size_, x_, + reinterpret_cast(k_scale.data()), + reinterpret_cast(v_scale.data())); + }; + + if (quantized) { + if (use_e5m2) { + launch(TypeTag{}, std::true_type{}, std::true_type{}); + } else { + launch(TypeTag{}, std::true_type{}, std::false_type{}); + } + } else { + launch(type_tag, std::false_type{}, std::false_type{}); + } + }, + "ReshapeAndCache::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h b/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h new file mode 100644 index 000000000..2ee381c8a --- /dev/null +++ b/src/native/cuda/nvidia/ops/rotary_embedding/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_NVIDIA_ROTARY_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_NVIDIA_ROTARY_EMBEDDING_KERNEL_H_ + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/rotary_embedding/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaRotaryEmbedding> { + public: + using CudaRotaryEmbedding< + Runtime>::CudaRotaryEmbedding; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h b/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..85d311717 --- /dev/null +++ b/src/native/cuda/nvidia/ops/silu_and_mul/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_NVIDIA_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_NVIDIA_SILU_AND_MUL_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/silu_and_mul/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSiluAndMul> { + public: + using CudaSiluAndMul>::CudaSiluAndMul; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/add/kernel.cuh b/src/native/cuda/ops/add/kernel.cuh index ec11fa7ca..e858ac5cd 100644 --- a/src/native/cuda/ops/add/kernel.cuh +++ b/src/native/cuda/ops/add/kernel.cuh @@ -9,15 +9,20 @@ template struct AddOp { static constexpr std::size_t num_inputs = 2; + double alpha; + template __device__ __forceinline__ T operator()(const T& input, const T& other) const { if constexpr (IsFP16 || IsBFloat16) { - return __hadd(input, other); + auto input_float = Caster::template Cast(input); + auto other_float = Caster::template Cast(other); + return Caster::template Cast( + input_float + static_cast(alpha) * other_float); } else if constexpr (std::is_same_v) { - return __fadd_rn(input, other); + return __fadd_rn(input, __fmul_rn(static_cast(alpha), other)); } else { - return input + other; + return input + static_cast(alpha) * other; } } }; @@ -32,7 +37,8 @@ __global__ void AddKernel(T* __restrict__ out, const T* __restrict__ input, const ptrdiff_t* __restrict__ input_strides, const ptrdiff_t* __restrict__ other_strides, size_t output_size, size_t ndim, bool out_contiguous, - bool input_contiguous, bool other_contiguous) { + bool input_contiguous, bool other_contiguous, + double alpha) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < output_size) { @@ -45,7 +51,7 @@ __global__ void AddKernel(T* __restrict__ out, const T* __restrict__ input, other_contiguous ? idx : IndexToOffset(idx, ndim, other_shape, other_strides); - out[out_idx] = AddOp{}(input[input_idx], other[other_idx]); + out[out_idx] = AddOp{alpha}(input[input_idx], other[other_idx]); } } diff --git a/src/native/cuda/ops/add/kernel.h b/src/native/cuda/ops/add/kernel.h index 07a111261..65d0ea7fb 100644 --- a/src/native/cuda/ops/add/kernel.h +++ b/src/native/cuda/ops/add/kernel.h @@ -17,8 +17,9 @@ namespace infini::ops { template class CudaAdd : public Add { public: - CudaAdd(const Tensor input, const Tensor other, Tensor out) - : Add{input, other, out} { + CudaAdd(const Tensor input, const Tensor other, const double alpha, + Tensor out) + : Add{input, other, alpha, out} { size_t shape_size = ndim_ * sizeof(*d_input_shape_); size_t strides_size = ndim_ * sizeof(*d_input_strides_); const size_t metadata_size = 3 * (shape_size + strides_size); @@ -54,9 +55,12 @@ class CudaAdd : public Add { Backend::kMemcpyHostToDevice); } + CudaAdd(const Tensor input, const Tensor other, Tensor out) + : CudaAdd{input, other, 1.0, out} {} + ~CudaAdd() { Backend::Free(d_metadata_); } - void operator()(const Tensor input, const Tensor other, + void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override { int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( @@ -80,7 +84,7 @@ class CudaAdd : public Add { d_out, d_input, d_other, d_out_shape_, d_input_shape_, d_other_shape_, d_out_strides_, d_input_strides_, d_other_strides_, output_size_, ndim_, is_out_contiguous_, - is_input_contiguous_, is_other_contiguous_); + is_input_contiguous_, is_other_contiguous_, alpha); }, "CudaAdd::operator()"); } diff --git a/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh new file mode 100644 index 000000000..8ffb079ca --- /dev/null +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh @@ -0,0 +1,107 @@ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_CUH_ + +#include +#include +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { + +namespace { + +template +__device__ __forceinline__ Data ExpAndCastInfinilm(Compute x) { + return Caster::template Cast( + expf(Caster::template Cast(x))); +} + +struct BlockMaxInfinilmOp { + template + __device__ __forceinline__ T operator()(const T& a, const T& b) const { + return (a > b) ? a : b; + } +}; + +template +__device__ __forceinline__ Data BlockMaxInfinilm(const Data* data_ptr, + size_t count) { + Data thread_max = count > 0 ? data_ptr[0] : Data{}; + for (size_t i = threadIdx.x; i < count; i += block_size) { + Data v = data_ptr[i]; + thread_max = (v > thread_max) ? v : thread_max; + } + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + return BlockReduce(temp_storage).Reduce(thread_max, BlockMaxInfinilmOp()); +} + +template +__device__ __forceinline__ Compute BlockSumInfinilm(const Data* data_ptr, + size_t count) { + Compute thread_sum = 0; + for (size_t i = threadIdx.x; i < count; i += block_size) { + thread_sum += Caster::template Cast(data_ptr[i]); + } + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + return BlockReduce(temp_storage).Sum(thread_sum); +} + +} // namespace + +template +__global__ void CausalSoftmaxInfinilmKernel( + Data* __restrict__ out_ptr, const Data* __restrict__ input_ptr, + size_t batch_size, size_t seq_len, size_t total_seq_len, + int64_t stride_out_batch, int64_t stride_out_row, + int64_t stride_input_batch, int64_t stride_input_row) { + size_t row_idx = blockIdx.x; + size_t batch_idx = blockIdx.y; + + Data* out_row = + out_ptr + batch_idx * stride_out_batch + row_idx * stride_out_row; + const Data* input_row = + input_ptr + batch_idx * stride_input_batch + row_idx * stride_input_row; + + size_t valid_len = total_seq_len - seq_len + row_idx + 1; + + __shared__ Data max_val; + Data block_max = BlockMaxInfinilm(input_row, valid_len); + if (threadIdx.x == 0) { + max_val = block_max; + } + __syncthreads(); + + for (size_t col = threadIdx.x; col < total_seq_len; col += block_size) { + if (col < valid_len) { + Compute diff = Caster::template Cast(input_row[col]) - + Caster::template Cast(max_val); + out_row[col] = ExpAndCastInfinilm(diff); + } else { + out_row[col] = Caster::template Cast(0.0f); + } + } + __syncthreads(); + + __shared__ Compute sum_val; + Compute block_sum = + BlockSumInfinilm(out_row, total_seq_len); + if (threadIdx.x == 0) { + sum_val = block_sum; + } + __syncthreads(); + + for (size_t col = threadIdx.x; col < total_seq_len; col += block_size) { + Compute quot = Caster::template Cast(out_row[col]) / sum_val; + out_row[col] = Caster::template Cast(quot); + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h new file mode 100644 index 000000000..4e1a97924 --- /dev/null +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h @@ -0,0 +1,64 @@ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ + +#include +#include +#include + +#include "base/causal_softmax_infinilm.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaCausalSoftmaxInfinilm : public CausalSoftmaxInfinilm { + public: + using CausalSoftmaxInfinilm::CausalSoftmaxInfinilm; + + void operator()(const Tensor input, Tensor out) const override { + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + + auto stride_input_batch = ndim_ == 3 ? input_strides_[0] : 0; + auto stride_input_row = input_strides_[ndim_ - 2]; + auto stride_out_batch = ndim_ == 3 ? out_strides_[0] : 0; + auto stride_out_row = out_strides_[ndim_ - 2]; + + dim3 grid(static_cast(seq_len_), + static_cast(batch_size_)); + + assert(out.dtype() == input.dtype()); + + constexpr int kMaxBlockSize = BackendMaxBlockSize::value; + int block_size = + std::min(RuntimeUtils::GetOptimalBlockSize(), + kMaxBlockSize); + + DispatchFunc< + ConcatType, ReducedFloatTypes>, + SupportedCudaBlockSizesType::value>>( + // TODO: Output dtype should use the one passed in during construction. + {static_cast(out.dtype()), block_size}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + constexpr int kBlockSize = ListGet<1>(list_tag); + + CausalSoftmaxInfinilmKernel + <<>>( + reinterpret_cast(out.data()), + reinterpret_cast(input.data()), batch_size_, + seq_len_, total_seq_len_, stride_out_batch, stride_out_row, + stride_input_batch, stride_input_row); + }, + "CudaCausalSoftmaxInfinilm::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/embedding/kernel.h b/src/native/cuda/ops/embedding/kernel.h index 7e8186383..2fab2df2f 100644 --- a/src/native/cuda/ops/embedding/kernel.h +++ b/src/native/cuda/ops/embedding/kernel.h @@ -18,19 +18,22 @@ namespace infini::ops { template class CudaEmbedding : public Embedding { public: - CudaEmbedding(const Tensor input, const Tensor weight, Tensor out) - : Embedding{input, weight, out}, - input_ndim_{input.ndim()}, + CudaEmbedding(const Tensor weight, const Tensor indices, + const int64_t padding_idx, const bool scale_grad_by_freq, + const bool sparse, Tensor out) + : Embedding{weight, indices, padding_idx, scale_grad_by_freq, + sparse, out}, + indices_ndim_{indices.ndim()}, out_ndim_{out.ndim()}, - is_input_contiguous_{input.IsContiguous()}, + are_indices_contiguous_{indices.IsContiguous()}, is_out_contiguous_{out.IsContiguous()}, weight_row_stride_{weight.stride(0)}, weight_col_stride_{weight.stride(1)} { - size_t input_shape_size = input_ndim_ * sizeof(*d_input_shape_); - size_t input_strides_size = input_ndim_ * sizeof(*d_input_strides_); + size_t indices_shape_size = indices_ndim_ * sizeof(*d_indices_shape_); + size_t indices_strides_size = indices_ndim_ * sizeof(*d_indices_strides_); size_t out_shape_size = out_ndim_ * sizeof(*d_out_shape_); size_t out_strides_size = out_ndim_ * sizeof(*d_out_strides_); - const size_t metadata_size = input_shape_size + input_strides_size + + const size_t metadata_size = indices_shape_size + indices_strides_size + out_shape_size + out_strides_size; std::vector metadata(metadata_size); @@ -38,15 +41,16 @@ class CudaEmbedding : public Embedding { Backend::Malloc((void**)&d_metadata_, metadata_size); size_t offset = 0; - d_input_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_shape_.data(), - input_shape_size); - offset += input_shape_size; + d_indices_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, indices_shape_.data(), + indices_shape_size); + offset += indices_shape_size; - d_input_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, input_strides_.data(), - input_strides_size); - offset += input_strides_size; + d_indices_strides_ = + reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, indices_strides_.data(), + indices_strides_size); + offset += indices_strides_size; d_out_shape_ = reinterpret_cast(d_metadata_ + offset); std::memcpy(metadata.data() + offset, out_shape_.data(), out_shape_size); @@ -60,9 +64,14 @@ class CudaEmbedding : public Embedding { Backend::kMemcpyHostToDevice); } + CudaEmbedding(const Tensor weight, const Tensor indices, Tensor out) + : CudaEmbedding(weight, indices, -1, false, false, out) {} + ~CudaEmbedding() { Backend::Free(d_metadata_); } - void operator()(const Tensor input, const Tensor weight, + void operator()(const Tensor weight, const Tensor indices, + const int64_t /*padding_idx*/, + const bool /*scale_grad_by_freq*/, const bool /*sparse*/, Tensor out) const override { if (num_indices_ == 0) { return; @@ -82,7 +91,7 @@ class CudaEmbedding : public Embedding { DispatchFunc, ConcatType, ReducedFloatTypes>>( - {static_cast(input_dtype_), + {static_cast(indices_dtype_), static_cast(weight_dtype_)}, [&](auto list_tag) { using IndexT = @@ -92,22 +101,22 @@ class CudaEmbedding : public Embedding { EmbeddingKernel <<>>( reinterpret_cast(out.data()), - reinterpret_cast(input.data()), + reinterpret_cast(indices.data()), reinterpret_cast(weight.data()), num_indices_, - input_ndim_, d_input_shape_, d_input_strides_, out_ndim_, - d_out_shape_, d_out_strides_, weight_row_stride_, + indices_ndim_, d_indices_shape_, d_indices_strides_, + out_ndim_, d_out_shape_, d_out_strides_, weight_row_stride_, weight_col_stride_, embedding_dim_, vocab_size_, - is_input_contiguous_, is_out_contiguous_); + are_indices_contiguous_, is_out_contiguous_); }, "CudaEmbedding::operator()"); } private: - Tensor::Size input_ndim_{0}; + Tensor::Size indices_ndim_{0}; Tensor::Size out_ndim_{0}; - bool is_input_contiguous_{false}; + bool are_indices_contiguous_{false}; bool is_out_contiguous_{false}; @@ -117,9 +126,9 @@ class CudaEmbedding : public Embedding { std::byte* d_metadata_{nullptr}; - Tensor::Size* d_input_shape_{nullptr}; + Tensor::Size* d_indices_shape_{nullptr}; - Tensor::Stride* d_input_strides_{nullptr}; + Tensor::Stride* d_indices_strides_{nullptr}; Tensor::Size* d_out_shape_{nullptr}; diff --git a/src/native/cuda/ops/fused_add_rms_norm/kernel.cuh b/src/native/cuda/ops/fused_add_rms_norm/kernel.cuh new file mode 100644 index 000000000..5b21050f4 --- /dev/null +++ b/src/native/cuda/ops/fused_add_rms_norm/kernel.cuh @@ -0,0 +1,69 @@ +#ifndef INFINI_OPS_CUDA_FUSED_ADD_RMS_NORM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_FUSED_ADD_RMS_NORM_KERNEL_CUH_ + +#include +#include +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { +namespace fused_add_rms_norm_detail { + +// Same as `native/cuda/ops/rms_norm/kernel.cuh`. +template +__device__ __forceinline__ TCompute SumSquared(const TData* data_ptr, + size_t count) { + TCompute ss = 0; + for (size_t i = threadIdx.x; i < count; i += block_size) { + TCompute value = Caster::template Cast(data_ptr[i]); + ss += value * value; + } + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage temp_storage; + return BlockReduce(temp_storage).Sum(ss); +} + +} // namespace fused_add_rms_norm_detail + +template +__global__ void FusedAddRmsNormKernel(TData* input, int64_t stride_input, + TData* residual, int64_t stride_residual, + const TData* __restrict__ weight, + size_t dim, float epsilon) { + auto input_ptr = input + blockIdx.x * stride_input; + auto residual_ptr = residual + blockIdx.x * stride_residual; + + for (size_t i = threadIdx.x; i < dim; i += block_size) { + TCompute sum_val = Caster::template Cast(input_ptr[i]) + + Caster::template Cast(residual_ptr[i]); + residual_ptr[i] = Caster::template Cast(sum_val); + } + + TCompute sum_squared = + fused_add_rms_norm_detail::SumSquared( + residual_ptr, dim); + + __shared__ TCompute rms; + if (threadIdx.x == 0) { + rms = Caster::template Cast(rsqrtf( + sum_squared / Caster::template Cast(dim) + epsilon)); + } + __syncthreads(); + + for (size_t i = threadIdx.x; i < dim; i += block_size) { + TCompute value = Caster::template Cast(residual_ptr[i]); + value *= rms; + if (weight != nullptr) { + value *= Caster::template Cast(weight[i]); + } + input_ptr[i] = Caster::template Cast(value); + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/fused_add_rms_norm/kernel.h b/src/native/cuda/ops/fused_add_rms_norm/kernel.h new file mode 100644 index 000000000..fd22c970e --- /dev/null +++ b/src/native/cuda/ops/fused_add_rms_norm/kernel.h @@ -0,0 +1,53 @@ +#ifndef INFINI_OPS_CUDA_FUSED_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_CUDA_FUSED_ADD_RMS_NORM_KERNEL_H_ + +#include +#include +#include + +#include "base/fused_add_rms_norm.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/fused_add_rms_norm/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaFusedAddRmsNorm : public FusedAddRmsNorm { + public: + using FusedAddRmsNorm::FusedAddRmsNorm; + + void operator()(Tensor input, Tensor residual, + const std::optional weight, float) const override { + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + + int block_size = RuntimeUtils::GetOptimalBlockSize(); + + DispatchFunc, ReducedFloatTypes>, + AllCudaBlockSizes>( + {static_cast(input.dtype()), block_size}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + constexpr int kBlockSize = ListGet<1>(list_tag); + + auto weight_data = weight.has_value() + ? reinterpret_cast(weight->data()) + : nullptr; + FusedAddRmsNormKernel + <<(num_tokens_), kBlockSize, 0, + cuda_stream>>>(reinterpret_cast(input.data()), + input_strides_[input_strides_.size() - 2], + reinterpret_cast(residual.data()), + residual_strides_[residual_strides_.size() - 2], + weight_data, dim_, epsilon_); + }, + "CudaFusedAddRmsNorm::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/rotary_embedding/kernel.cuh b/src/native/cuda/ops/rotary_embedding/kernel.cuh new file mode 100644 index 000000000..f94804195 --- /dev/null +++ b/src/native/cuda/ops/rotary_embedding/kernel.cuh @@ -0,0 +1,66 @@ +#ifndef INFINI_OPS_CUDA_ROTARY_EMBEDDING_KERNEL_CUH_ +#define INFINI_OPS_CUDA_ROTARY_EMBEDDING_KERNEL_CUH_ + +#include +#include + +#include "native/cuda/caster.cuh" + +namespace infini::ops { + +template +__device__ __forceinline__ void ApplyRotaryPair(T* data, const TCache* cache, + size_t rot_offset, + size_t embed_dim, + bool inverse) { + const size_t x_index = kIsNeox ? rot_offset : 2 * rot_offset; + const size_t y_index = kIsNeox ? embed_dim + rot_offset : 2 * rot_offset + 1; + const size_t cache_index = kIsNeox ? x_index : x_index / 2; + float cos_value = Caster::template Cast(cache[cache_index]); + float sin_value = + Caster::template Cast(cache[embed_dim + cache_index]); + if (inverse) { + sin_value = -sin_value; + } + + const float x = Caster::template Cast(data[x_index]); + const float y = Caster::template Cast(data[y_index]); + data[x_index] = Caster::template Cast(x * cos_value - y * sin_value); + data[y_index] = Caster::template Cast(y * cos_value + x * sin_value); +} + +template +__global__ void RotaryEmbeddingKernel( + const int64_t* __restrict__ positions, T* query, T* key, + const TCache* __restrict__ cos_sin_cache, int64_t cache_token_stride, + int64_t query_token_stride, int64_t key_token_stride, + int64_t query_head_stride, int64_t key_head_stride, size_t num_heads, + size_t num_kv_heads, size_t rot_dim, size_t rope_dim_offset, bool inverse) { + const size_t token_idx = blockIdx.x; + const int64_t position = positions[token_idx]; + const TCache* cache = cos_sin_cache + position * cache_token_stride; + const size_t embed_dim = rot_dim / 2; + + for (size_t i = threadIdx.x; i < num_heads * embed_dim; i += blockDim.x) { + const size_t head_idx = i / embed_dim; + T* data = query + token_idx * query_token_stride + + head_idx * query_head_stride + rope_dim_offset; + ApplyRotaryPair(data, cache, i % embed_dim, + embed_dim, inverse); + } + + if (key != nullptr) { + for (size_t i = threadIdx.x; i < num_kv_heads * embed_dim; + i += blockDim.x) { + const size_t head_idx = i / embed_dim; + T* data = key + token_idx * key_token_stride + + head_idx * key_head_stride + rope_dim_offset; + ApplyRotaryPair(data, cache, i % embed_dim, + embed_dim, inverse); + } + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/rotary_embedding/kernel.h b/src/native/cuda/ops/rotary_embedding/kernel.h new file mode 100644 index 000000000..f10936909 --- /dev/null +++ b/src/native/cuda/ops/rotary_embedding/kernel.h @@ -0,0 +1,65 @@ +#ifndef INFINI_OPS_CUDA_ROTARY_EMBEDDING_KERNEL_H_ +#define INFINI_OPS_CUDA_ROTARY_EMBEDDING_KERNEL_H_ + +#include +#include +#include + +#include "base/rotary_embedding.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/ops/rotary_embedding/kernel.cuh" + +namespace infini::ops { + +template +class CudaRotaryEmbedding : public RotaryEmbedding { + public: + using RotaryEmbedding::RotaryEmbedding; + + void operator()(const Tensor positions, Tensor query, + std::optional key, int64_t, + const Tensor cos_sin_cache, bool, int64_t, + bool) const override { + if (num_tokens_ == 0) { + return; + } + + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + DispatchFunc, ReducedFloatTypes>, + ConcatType, ReducedFloatTypes>>( + {static_cast(query_type_), + static_cast(cos_sin_cache_type_)}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + using TCache = + TypeMapType(list_tag)>; + auto key_data = + key.has_value() ? reinterpret_cast(key->data()) : nullptr; + auto launch = [&](auto neox_tag) { + constexpr bool kIsNeox = decltype(neox_tag)::value; + RotaryEmbeddingKernel + <<(num_tokens_), 256, 0, cuda_stream>>>( + reinterpret_cast(positions.data()), + reinterpret_cast(query.data()), key_data, + reinterpret_cast(cos_sin_cache.data()), + cos_sin_cache_strides_[0], query_token_stride_, + key_token_stride_, query_head_stride_, key_head_stride_, + num_heads_, num_kv_heads_, rot_dim_, rope_dim_offset_, + inverse_); + }; + + if (is_neox_) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } + }, + "CudaRotaryEmbedding::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/silu_and_mul/kernel.cuh b/src/native/cuda/ops/silu_and_mul/kernel.cuh new file mode 100644 index 000000000..c5ad1ce1f --- /dev/null +++ b/src/native/cuda/ops/silu_and_mul/kernel.cuh @@ -0,0 +1,85 @@ +#ifndef INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_CUH_ +#define INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_CUH_ + +#include + +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { + +namespace detail { + +// Optimized sigmoid function with support for FP16 and BF16 types. +// TODO: The unified FP16/BF16 branch uses `Caster` and scalar float +// arithmetic instead of native vectorized intrinsics (e.g. `h2rcp`, +// `__hmul2`). Profile and restore specialized paths if needed. +template +__device__ __forceinline__ T Sigmoid(const T& x) { + if constexpr (IsFP16 || IsBFloat16) { + float xf = Caster::template Cast(x); + return Caster::template Cast( + __frcp_rn(__fadd_rn(1.0f, __expf(-xf)))); + } else if constexpr (std::is_same_v) { + return __frcp_rn(__fadd_rn(1.0f, __expf(-x))); + } else { + return 1.0f / (1.0f + expf(-x)); + } +} + +} // namespace detail + +template +__global__ void SiluAndMulKernel(T* __restrict__ out, + const T* __restrict__ input, + const size_t* __restrict__ out_shape, + const size_t* __restrict__ input_shape, + const ptrdiff_t* __restrict__ out_strides, + const ptrdiff_t* __restrict__ input_strides, + size_t output_size, size_t ndim, + size_t hidden_size, bool out_contiguous, + bool input_contiguous) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + + if (idx < output_size) { + size_t out_idx; + + if (out_contiguous) { + out_idx = idx; + } else { + out_idx = IndexToOffset(idx, ndim, out_shape, out_strides); + } + + const size_t row = idx / hidden_size; + const size_t column = idx % hidden_size; + const size_t gate_logical_idx = row * 2 * hidden_size + column; + const size_t up_logical_idx = gate_logical_idx + hidden_size; + const size_t gate_idx = + input_contiguous + ? gate_logical_idx + : IndexToOffset(gate_logical_idx, ndim, input_shape, input_strides); + const size_t up_idx = + input_contiguous + ? up_logical_idx + : IndexToOffset(up_logical_idx, ndim, input_shape, input_strides); + + T gate = input[gate_idx]; + T up = input[up_idx]; + + if constexpr (IsFP16 || IsBFloat16) { + float gatef = Caster::template Cast(gate); + float upf = Caster::template Cast(up); + float sigf = __frcp_rn(__fadd_rn(1.0f, __expf(-gatef))); + out[out_idx] = Caster::template Cast( + __fmul_rn(__fmul_rn(gatef, sigf), upf)); + } else if constexpr (std::is_same_v) { + out[out_idx] = + __fmul_rn(__fmul_rn(gate, detail::Sigmoid(gate)), up); + } else { + out[out_idx] = gate * detail::Sigmoid(gate) * up; + } + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/silu_and_mul/kernel.h b/src/native/cuda/ops/silu_and_mul/kernel.h new file mode 100644 index 000000000..0e9b5e620 --- /dev/null +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -0,0 +1,89 @@ +#ifndef INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/silu_and_mul.h" +#include "common/generic_utils.h" +#include "native/cuda/ops/silu_and_mul/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaSiluAndMul : public SiluAndMul { + public: + CudaSiluAndMul(const Tensor input, Tensor out) : SiluAndMul{input, out} { + size_t shape_size = ndim_ * sizeof(*d_input_shape_); + size_t strides_size = ndim_ * sizeof(*d_input_strides_); + + const size_t metadata_size = 2 * (shape_size + strides_size); + std::vector metadata(metadata_size); + + Backend::Malloc((void**)&d_metadata_, metadata_size); + + size_t offset = 0; + d_input_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); + offset += shape_size; + + d_out_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size); + offset += shape_size; + + d_input_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); + offset += strides_size; + + d_out_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size); + + Backend::Memcpy(d_metadata_, metadata.data(), metadata_size, + Backend::kMemcpyHostToDevice); + } + + ~CudaSiluAndMul() { Backend::Free(d_metadata_); } + + void operator()(const Tensor input, Tensor out) const override { + int block_size = RuntimeUtils::GetOptimalBlockSize(); + DispatchFunc( + {static_cast(out_type_), block_size}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + constexpr int kBlockSize = ListGet<1>(list_tag); + + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + dim3 blockDims( + std::min(static_cast(block_size), output_size_)); + dim3 gridDims(utils::CeilDiv(output_size_, blockDims.x)); + + T* d_out = reinterpret_cast(out.data()); + const T* d_input = reinterpret_cast(input.data()); + SiluAndMulKernel + <<>>( + d_out, d_input, d_out_shape_, d_input_shape_, d_out_strides_, + d_input_strides_, output_size_, ndim_, hidden_size_, + is_out_contiguous_, is_input_contiguous_); + }, + "CudaSiluAndMul::operator()"); + } + + private: + std::byte* d_metadata_{nullptr}; + + Tensor::Size* d_input_shape_{nullptr}; + + Tensor::Size* d_out_shape_{nullptr}; + + Tensor::Stride* d_input_strides_{nullptr}; + + Tensor::Stride* d_out_strides_{nullptr}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 1571e2981..679e3b12c 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -26,6 +26,10 @@ std::unordered_map BuildTorchNameMap( } // namespace detail inline DataType DataTypeFromString(const std::string& name) { + // InfiniRT has no bool dtype; carry bool tensor storage as byte data and + // restore bool semantics in operators that accept bool tensors. + if (name == "bool") return DataType::kUInt8; + return kStringToDataType.at(name); } diff --git a/src/torch/ops/add/add.cc b/src/torch/ops/add/add.cc index 8e6530d78..cc36f5bc0 100644 --- a/src/torch/ops/add/add.cc +++ b/src/torch/ops/add/add.cc @@ -4,14 +4,19 @@ namespace infini::ops { +template +Operator::Operator(const Tensor input, const Tensor other, + const double alpha, Tensor out) + : Add{input, other, alpha, out}, device_index_{out.device().index()} {} + template Operator::Operator(const Tensor input, const Tensor other, Tensor out) - : Add{input, other, out}, device_index_{out.device().index()} {} + : Operator{input, other, 1.0, out} {} template void Operator::operator()(const Tensor input, const Tensor other, - Tensor out) const { + const double alpha, Tensor out) const { auto at_input = ToAtenTensor(const_cast(input.data()), input_shape_, input_strides_, input_type_, device_index_); @@ -21,7 +26,10 @@ void Operator::operator()(const Tensor input, const Tensor other, auto at_out = ToAtenTensor(out.data(), out_shape_, out_strides_, out_type_, device_index_); - at::add_out(at_out, at_input, at_other); + auto at_alpha = at_input.is_floating_point() + ? at::Scalar(alpha) + : at::Scalar(static_cast(alpha)); + at::add_out(at_out, at_input, at_other, at_alpha); } template class Operator; diff --git a/src/torch/ops/add/add.h b/src/torch/ops/add/add.h index 5a7ea230c..f5624f521 100644 --- a/src/torch/ops/add/add.h +++ b/src/torch/ops/add/add.h @@ -8,9 +8,12 @@ namespace infini::ops { template class Operator : public Add { public: + Operator(const Tensor input, const Tensor other, const double alpha, + Tensor out); + Operator(const Tensor input, const Tensor other, Tensor out); - void operator()(const Tensor input, const Tensor other, + void operator()(const Tensor input, const Tensor other, const double alpha, Tensor out) const override; private: diff --git a/src/torch/ops/reshape_and_cache/reshape_and_cache.cc b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc new file mode 100644 index 000000000..c1d5f45f2 --- /dev/null +++ b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc @@ -0,0 +1,89 @@ +#include "torch/ops/reshape_and_cache/reshape_and_cache.h" + +#include "torch/tensor_.h" + +namespace infini::ops { +namespace { + +at::Tensor PrepareCacheValues(const at::Tensor& values, const at::Tensor& scale, + const std::string& kv_cache_dtype) { + if (kv_cache_dtype == "auto") { + return values; + } + + const auto fp8_type = + kv_cache_dtype == "fp8_e5m2" ? at::kFloat8_e5m2 : at::kFloat8_e4m3fn; + return (values / scale).to(fp8_type).view(at::kByte); +} + +} // namespace + +template +Operator::Operator( + const Tensor key, const Tensor value, Tensor key_cache, Tensor value_cache, + const Tensor slot_mapping, const std::string kv_cache_dtype, + const Tensor k_scale, const Tensor v_scale) + : ReshapeAndCache{key, value, key_cache, value_cache, + slot_mapping, kv_cache_dtype, k_scale, v_scale} {} + +template +void Operator::operator()( + const Tensor key, const Tensor value, Tensor key_cache, Tensor value_cache, + const Tensor slot_mapping, const std::string kv_cache_dtype, + const Tensor k_scale, const Tensor v_scale) const { + auto at_key = ToAtenTensor(const_cast(key.data()), key_shape_, + key_strides_, key_type_, device_index_); + auto at_value = + ToAtenTensor(const_cast(value.data()), value_shape_, + value_strides_, key_type_, device_index_); + auto at_key_cache = + ToAtenTensor(key_cache.data(), key_cache_shape_, key_cache_strides_, + key_cache_type_, device_index_); + auto at_value_cache = ToAtenTensor( + value_cache.data(), value_cache_shape_, value_cache_strides_, + value_cache_type_, device_index_); + auto at_slot_mapping = ToAtenTensor( + const_cast(slot_mapping.data()), slot_mapping_shape_, + slot_mapping_strides_, DataType::kInt64, device_index_); + auto at_k_scale = + ToAtenTensor(const_cast(k_scale.data()), k_scale_shape_, + k_scale_strides_, k_scale_type_, device_index_); + auto at_v_scale = + ToAtenTensor(const_cast(v_scale.data()), v_scale_shape_, + v_scale_strides_, v_scale_type_, device_index_); + + auto token_indices = at::nonzero(at_slot_mapping >= 0).reshape({-1}); + if (token_indices.numel() == 0) { + return; + } + + auto slots = at_slot_mapping.index_select(0, token_indices); + auto block_indices = + at::floor_divide(slots, static_cast(block_size_)); + auto block_offsets = at::remainder(slots, static_cast(block_size_)); + auto key_values = + at_key.index_select(0, token_indices) + .view({token_indices.numel(), static_cast(num_heads_), + static_cast(head_size_ / x_), + static_cast(x_)}); + auto value_values = at_value.index_select(0, token_indices); + key_values = PrepareCacheValues(key_values, at_k_scale, kv_cache_dtype); + value_values = PrepareCacheValues(value_values, at_v_scale, kv_cache_dtype); + + using torch::indexing::Slice; + at_key_cache.index_put_( + {block_indices, Slice(), Slice(), block_offsets, Slice()}, key_values); + at_value_cache.index_put_({block_indices, Slice(), Slice(), block_offsets}, + value_values); +} + +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; + +} // namespace infini::ops diff --git a/src/torch/ops/reshape_and_cache/reshape_and_cache.h b/src/torch/ops/reshape_and_cache/reshape_and_cache.h new file mode 100644 index 000000000..d9f3d97b9 --- /dev/null +++ b/src/torch/ops/reshape_and_cache/reshape_and_cache.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_TORCH_RESHAPE_AND_CACHE_H_ +#define INFINI_OPS_TORCH_RESHAPE_AND_CACHE_H_ + +#include + +#include "base/reshape_and_cache.h" + +namespace infini::ops { + +template +class Operator : public ReshapeAndCache { + public: + Operator(const Tensor key, const Tensor value, Tensor key_cache, + Tensor value_cache, const Tensor slot_mapping, + const std::string kv_cache_dtype, const Tensor k_scale, + const Tensor v_scale); + + void operator()(const Tensor key, const Tensor value, Tensor key_cache, + Tensor value_cache, const Tensor slot_mapping, + const std::string kv_cache_dtype, const Tensor k_scale, + const Tensor v_scale) const override; +}; + +} // namespace infini::ops + +#endif diff --git a/src/torch/ops/rotary_embedding/rotary_embedding.cc b/src/torch/ops/rotary_embedding/rotary_embedding.cc new file mode 100644 index 000000000..96a9c53ef --- /dev/null +++ b/src/torch/ops/rotary_embedding/rotary_embedding.cc @@ -0,0 +1,87 @@ +#include "torch/ops/rotary_embedding/rotary_embedding.h" + +#include "torch/tensor_.h" + +namespace infini::ops { +namespace { + +void ApplyRotaryEmbedding(at::Tensor data, const at::Tensor& positions, + const at::Tensor& cos_sin_cache, int64_t token_stride, + int64_t head_stride, int64_t num_heads, + int64_t head_size, int64_t rot_dim, + int64_t rope_dim_offset, bool is_neox, bool inverse) { + const int64_t num_tokens = positions.numel(); + const int64_t embed_dim = rot_dim / 2; + auto data_view = data.as_strided({num_tokens, num_heads, head_size}, + {token_stride, head_stride, 1}); + auto cache = cos_sin_cache.index_select(0, positions.reshape({-1})); + auto cos = cache.slice(1, 0, embed_dim).unsqueeze(1).to(at::kFloat); + auto sin = cache.slice(1, embed_dim, rot_dim).unsqueeze(1).to(at::kFloat); + if (inverse) { + sin = -sin; + } + + auto rotary = data_view.slice(2, rope_dim_offset, rope_dim_offset + rot_dim); + if (is_neox) { + auto x = rotary.slice(2, 0, embed_dim).to(at::kFloat).clone(); + auto y = rotary.slice(2, embed_dim, rot_dim).to(at::kFloat).clone(); + rotary.slice(2, 0, embed_dim).copy_(x * cos - y * sin); + rotary.slice(2, embed_dim, rot_dim).copy_(y * cos + x * sin); + } else { + auto x = rotary.slice(2, 0, rot_dim, 2).to(at::kFloat).clone(); + auto y = rotary.slice(2, 1, rot_dim, 2).to(at::kFloat).clone(); + rotary.slice(2, 0, rot_dim, 2).copy_(x * cos - y * sin); + rotary.slice(2, 1, rot_dim, 2).copy_(y * cos + x * sin); + } +} + +} // namespace + +template +Operator::Operator( + const Tensor positions, Tensor query, std::optional key, + int64_t head_size, const Tensor cos_sin_cache, bool is_neox, + int64_t rope_dim_offset, bool inverse) + : RotaryEmbedding{positions, query, key, + head_size, cos_sin_cache, is_neox, + rope_dim_offset, inverse} {} + +template +void Operator::operator()( + const Tensor positions, Tensor query, std::optional key, int64_t, + const Tensor cos_sin_cache, bool, int64_t, bool) const { + if (num_tokens_ == 0) { + return; + } + + auto at_positions = + ToAtenTensor(const_cast(positions.data()), positions_shape_, + positions_strides_, positions_type_, device_index_); + auto at_query = ToAtenTensor(query.data(), query_shape_, query_strides_, + query_type_, device_index_); + auto at_cache = ToAtenTensor( + const_cast(cos_sin_cache.data()), cos_sin_cache_shape_, + cos_sin_cache_strides_, cos_sin_cache_type_, device_index_); + + ApplyRotaryEmbedding(at_query, at_positions, at_cache, query_token_stride_, + query_head_stride_, num_heads_, head_size_, rot_dim_, + rope_dim_offset_, is_neox_, inverse_); + if (key.has_value()) { + auto at_key = ToAtenTensor(key->data(), key_shape_, key_strides_, + query_type_, device_index_); + ApplyRotaryEmbedding(at_key, at_positions, at_cache, key_token_stride_, + key_head_stride_, num_kv_heads_, head_size_, rot_dim_, + rope_dim_offset_, is_neox_, inverse_); + } +} + +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; + +} // namespace infini::ops diff --git a/src/torch/ops/rotary_embedding/rotary_embedding.h b/src/torch/ops/rotary_embedding/rotary_embedding.h new file mode 100644 index 000000000..916db723c --- /dev/null +++ b/src/torch/ops/rotary_embedding/rotary_embedding.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_TORCH_ROTARY_EMBEDDING_H_ +#define INFINI_OPS_TORCH_ROTARY_EMBEDDING_H_ + +#include + +#include "base/rotary_embedding.h" + +namespace infini::ops { + +template +class Operator : public RotaryEmbedding { + public: + Operator(const Tensor positions, Tensor query, std::optional key, + int64_t head_size, const Tensor cos_sin_cache, bool is_neox, + int64_t rope_dim_offset = 0, bool inverse = false); + + void operator()(const Tensor positions, Tensor query, + std::optional key, int64_t head_size, + const Tensor cos_sin_cache, bool is_neox, + int64_t rope_dim_offset = 0, + bool inverse = false) const override; +}; + +} // namespace infini::ops + +#endif diff --git a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc new file mode 100644 index 000000000..81489cf8c --- /dev/null +++ b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc @@ -0,0 +1,82 @@ +#include "torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.h" + +#include "torch/tensor_.h" + +namespace infini::ops { + +template +Operator::Operator( + const Tensor query, const Tensor key, const Tensor value, + const std::optional attn_mask, double dropout_p, bool is_causal, + const std::optional scale, bool enable_gqa, Tensor out) + : ScaledDotProductAttention{query, key, value, + attn_mask, dropout_p, is_causal, + scale, enable_gqa, out} {} + +template +Operator::Operator(const Tensor query, + const Tensor key, + const Tensor value, + Tensor out) + : Operator{query, key, value, std::nullopt, 0.0, + false, std::nullopt, false, out} {} + +template +void Operator::operator()( + const Tensor query, const Tensor key, const Tensor value, + const std::optional attn_mask, double dropout_p, bool is_causal, + const std::optional scale, bool enable_gqa, Tensor out) const { + auto at_query = + ToAtenTensor(const_cast(query.data()), query_shape_, + query_strides_, query_type_, device_index_); + auto at_key = ToAtenTensor(const_cast(key.data()), key_shape_, + key_strides_, query_type_, device_index_); + auto at_value = + ToAtenTensor(const_cast(value.data()), value_shape_, + value_strides_, query_type_, device_index_); + auto at_out = ToAtenTensor(out.data(), out_shape_, out_strides_, + query_type_, device_index_); + + c10::optional at_attn_mask; + if (attn_mask.has_value()) { + const auto dtype_override = attn_mask_type_ == DataType::kUInt8 + ? std::optional{at::kBool} + : std::nullopt; + at_attn_mask.emplace(ToAtenTensor( + const_cast(attn_mask->data()), attn_mask_shape_, + attn_mask_strides_, attn_mask_type_, device_index_, dtype_override)); + } + + c10::optional at_scale; + if (scale.has_value()) { + at_scale = *scale; + } + +#if TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 5) + auto result = at::scaled_dot_product_attention( + at_query, at_key, at_value, at_attn_mask, dropout_p, is_causal, at_scale, + enable_gqa); +#else + if (enable_gqa) { + at_key = at_key.repeat_interleave(at_query.size(-3) / at_key.size(-3), -3); + at_value = + at_value.repeat_interleave(at_query.size(-3) / at_value.size(-3), -3); + } + + auto result = at::scaled_dot_product_attention( + at_query, at_key, at_value, at_attn_mask, dropout_p, is_causal, at_scale); +#endif + at_out.copy_(result); +} + +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; +template class Operator; + +} // namespace infini::ops diff --git a/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.h b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.h new file mode 100644 index 000000000..0848ca62e --- /dev/null +++ b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.h @@ -0,0 +1,30 @@ +#ifndef INFINI_OPS_TORCH_SCALED_DOT_PRODUCT_ATTENTION_H_ +#define INFINI_OPS_TORCH_SCALED_DOT_PRODUCT_ATTENTION_H_ + +#include + +#include "base/scaled_dot_product_attention.h" + +namespace infini::ops { + +template +class Operator + : public ScaledDotProductAttention { + public: + Operator(const Tensor query, const Tensor key, const Tensor value, + const std::optional attn_mask, double dropout_p, + bool is_causal, const std::optional scale, bool enable_gqa, + Tensor out); + + Operator(const Tensor query, const Tensor key, const Tensor value, + Tensor out); + + void operator()(const Tensor query, const Tensor key, const Tensor value, + const std::optional attn_mask, double dropout_p, + bool is_causal, const std::optional scale, + bool enable_gqa, Tensor out) const override; +}; + +} // namespace infini::ops + +#endif diff --git a/src/torch/tensor_.h b/src/torch/tensor_.h index 5780babd9..e199bf704 100644 --- a/src/torch/tensor_.h +++ b/src/torch/tensor_.h @@ -4,6 +4,8 @@ #include #include +#include + #include "tensor.h" #include "torch/device_.h" @@ -82,13 +84,15 @@ inline at::ScalarType ToAtenDataType(DataType dtype) { // shape/strides from the `Tensor` parameter, which may have been moved-from // by the `Call()` dispatch path (see `operator.h`). template -inline at::Tensor ToAtenTensor(void* data, const Tensor::Shape& shape, - const Tensor::Strides& strides, DataType dtype, - int device_index = 0) { +inline at::Tensor ToAtenTensor( + void* data, const Tensor::Shape& shape, const Tensor::Strides& strides, + DataType dtype, int device_index = 0, + std::optional dtype_override = std::nullopt) { std::vector at_shape(shape.begin(), shape.end()); std::vector at_strides(strides.begin(), strides.end()); - auto options = at::TensorOptions().dtype(ToAtenDataType(dtype)); + auto options = + at::TensorOptions().dtype(dtype_override.value_or(ToAtenDataType(dtype))); if constexpr (kDev != Device::Type::kCpu) { std::string device_str = @@ -97,6 +101,10 @@ inline at::Tensor ToAtenTensor(void* data, const Tensor::Shape& shape, options = options.device(device_str); } + for (auto dim : shape) { + if (dim == 0) return at::empty_strided(at_shape, at_strides, options); + } + return at::from_blob(data, at_shape, at_strides, options); } diff --git a/tests/conftest.py b/tests/conftest.py index f0be7431f..c1552d93a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -193,11 +193,13 @@ def _is_smoke_item(item): "tests/test_cast.py": _is_smoke_cast_case, "tests/test_cat.py": _is_smoke_cat_case, "tests/test_causal_softmax.py": _is_smoke_causal_softmax_case, + "tests/test_causal_softmax_infinilm.py": _is_smoke_causal_softmax_case, "tests/test_gemm.py": _is_smoke_gemm_case, "tests/test_linear.py": _is_smoke_linear_case, "tests/test_matmul.py": _is_smoke_matmul_case, "tests/test_mul.py": _is_smoke_mul_case, "tests/test_rms_norm.py": _is_smoke_rms_norm_case, + "tests/test_silu_and_mul.py": _is_smoke_silu_and_mul_case, "tests/test_swiglu.py": _is_smoke_swiglu_case, "tests/test_torch_ops.py": _is_smoke_torch_op_case, } @@ -336,26 +338,21 @@ def _is_smoke_matmul_case(params): ((2, 4, 64), (2, 64, 32), (2, 4, 32)), } - return ( - _is_float32(params) - and params.get("trans_a") is False - and params.get("trans_b") is False - and _shape_case(params, "a_shape", "b_shape", "c_shape") in cases + return _is_float32(params) and ( + _shape_case(params, "a_shape", "b_shape", "c_shape") in cases ) def _is_smoke_linear_case(params): cases = { - (((4, 64), (64, 32), (4, 32)), False), - (((2, 4, 64), (2, 64, 32), (2, 4, 32)), True), + (((4, 64), (32, 64), (4, 32)), False), + (((2, 4, 64), (32, 64), (2, 4, 32)), True), } return ( _is_float32(params) - and params.get("trans_a") is False - and params.get("trans_b") is False and ( - _shape_case(params, "a_shape", "b_shape", "out_shape"), + _shape_case(params, "input_shape", "weight_shape", "out_shape"), params.get("has_bias"), ) in cases @@ -383,6 +380,24 @@ def _is_smoke_rms_norm_case(params): ) +def _is_smoke_silu_and_mul_case(params): + cases = { + ((13, 4), None, None), + ((16, 5632), None, None), + } + + return ( + _is_float32(params) + and _shape_case( + params, + "out_shape", + "input_strides", + "out_strides", + ) + in cases + ) + + def _is_smoke_swiglu_case(params): cases = { ((13, 4), None, None, None), diff --git a/tests/test_add.py b/tests/test_add.py index e2266c30d..ae826cc30 100644 --- a/tests/test_add.py +++ b/tests/test_add.py @@ -109,3 +109,57 @@ def _torch_add(input, other, out): out.copy_(res.to(out.dtype)) return out + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + "input_shape, other_shape, out_shape", + ( + ((3, 1), (1, 5), (3, 5)), + ((5,), (2, 3, 5), (2, 3, 5)), + ), +) +@pytest.mark.parametrize("alpha", (0.0, 0.5, 1.0, 2.0)) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_add_alpha_and_broadcast( + input_shape, + other_shape, + out_shape, + alpha, + dtype, + device, + implementation_index, +): + input = randn_strided(input_shape, None, dtype=dtype, device=device) + other = randn_strided(other_shape, None, dtype=dtype, device=device) + out = empty_strided(out_shape, None, dtype=dtype, device=device) + + return Payload( + lambda *args: _add_alpha( + *args, alpha=alpha, implementation_index=implementation_index + ), + lambda *args: _torch_add_alpha(*args, alpha=alpha), + (input, other, out), + {}, + rtol=1e-2 if dtype != torch.float32 else 1e-6, + atol=1e-2 if dtype != torch.float32 else 1e-6, + ) + + +def _add_alpha(input, other, out, *, alpha, implementation_index): + infini.ops.add( + input, + other, + alpha, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_add_alpha(input, other, out, *, alpha): + torch.add(input, other, alpha=alpha, out=out) + + return out diff --git a/tests/test_cat.py b/tests/test_cat.py index 85428b53f..0c0f4ac10 100644 --- a/tests/test_cat.py +++ b/tests/test_cat.py @@ -23,6 +23,10 @@ (((2, 4, 32), (2, 4, 64)), 2, (2, 4, 96)), # 4 inputs, dim=1 (((1, 1024), (1, 1024), (1, 1024), (1, 1024)), 1, (1, 4096)), + # 1 input + (((4, 64),), 0, (4, 64)), + # Empty input dimension + (((0, 4), (3, 4)), 0, (3, 4)), ), ) @pytest.mark.parametrize( @@ -51,10 +55,7 @@ def _cat(*args, dim): inputs = list(args[:-1]) out = args[-1] - first = inputs[0] - rest = inputs[1:] - - infini.ops.cat(first, rest, dim, out, stream=get_stream(first.device)) + infini.ops.cat(inputs, dim, out, stream=get_stream(inputs[0].device)) return out @@ -67,3 +68,29 @@ def _torch_cat(*args, dim): out.copy_(result) return out + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-7, 1e-7), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 1e-2, 5e-3), + ), +) +def test_cat_noncontiguous(dtype, device, rtol, atol): + inputs = [ + randn_strided((2, 3, 4), (4, 8, 1), dtype=dtype, device=device) + for _ in range(2) + ] + out = empty_strided((2, 6, 4), None, dtype=dtype, device=device) + + return Payload( + lambda *args: _cat(*args, dim=1), + lambda *args: _torch_cat(*args, dim=1), + (*inputs, out), + {}, + rtol=rtol, + atol=atol, + ) diff --git a/tests/test_causal_softmax_infinilm.py b/tests/test_causal_softmax_infinilm.py new file mode 100644 index 000000000..15cfc0d39 --- /dev/null +++ b/tests/test_causal_softmax_infinilm.py @@ -0,0 +1,57 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randn_strided + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + "shape, input_strides, out_strides", + ( + ((3, 3), None, None), + ((3, 5), None, None), + ((32, 512), None, None), + ((32, 512), (1024, 1), (1024, 1)), + ((4, 20, 512), None, None), + ((4, 20, 512), (20480, 512, 1), None), + ), +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 1e-2, 1e-2), + ), +) +def test_causal_softmax_infinilm( + shape, input_strides, out_strides, dtype, device, rtol, atol +): + input_tensor = randn_strided(shape, input_strides, dtype=dtype, device=device) + out = empty_strided(shape, out_strides, dtype=dtype, device=device) + + return Payload( + _causal_softmax_infinilm, + _torch_causal_softmax, + (input_tensor, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _causal_softmax_infinilm(input, out): + infini.ops.causal_softmax_infinilm(input, out, stream=get_stream(input.device)) + + return out + + +def _torch_causal_softmax(input, out): + input_cpu = input.detach().cpu().to(torch.float32) + mask = torch.tril(torch.ones_like(input_cpu), diagonal=-1).flip(dims=[-2, -1]) + masked = torch.where(mask == 1, -torch.inf, input_cpu) + result = torch.nn.functional.softmax(masked, dim=-1) + out.copy_(result.to(device=out.device, dtype=out.dtype)) + + return out diff --git a/tests/test_embedding.py b/tests/test_embedding.py index 98fe6bad3..7bd22a3bd 100644 --- a/tests/test_embedding.py +++ b/tests/test_embedding.py @@ -69,27 +69,27 @@ def test_embedding( *args, **kwargs, implementation_index=implementation_index ), _torch_embedding, - (input, weight), + (weight, input), {"out": out}, rtol=rtol, atol=atol, ) -def _embedding(input, weight, *, out=None, implementation_index=0): +def _embedding(weight, indices, *, out=None, implementation_index=0): infini.ops.embedding( - input, weight, + indices, out, implementation_index=implementation_index, - stream=get_stream(input.device), + stream=get_stream(indices.device), ) return out -def _torch_embedding(input, weight, *, out=None): - result = torch.nn.functional.embedding(input, weight) +def _torch_embedding(weight, indices, *, out=None): + result = torch.nn.functional.embedding(indices, weight) if out is not None: out.copy_(result) @@ -97,3 +97,82 @@ def _torch_embedding(input, weight, *, out=None): out = result return out + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("padding_idx", (-1, 0)) +@pytest.mark.parametrize("scale_grad_by_freq", (False, True)) +@pytest.mark.parametrize("sparse", (False, True)) +def test_embedding_parameters( + padding_idx, + scale_grad_by_freq, + sparse, + device, + implementation_index, +): + weight = randn_strided((8, 4), None, dtype=torch.float32, device=device) + indices = randint_strided(0, 8, (2, 3), None, dtype=torch.int64, device=device) + out = empty_strided((2, 3, 4), None, dtype=torch.float32, device=device) + + return Payload( + lambda *args: _embedding_with_parameters( + *args, + padding_idx=padding_idx, + scale_grad_by_freq=scale_grad_by_freq, + sparse=sparse, + implementation_index=implementation_index, + ), + lambda *args: _torch_embedding_with_parameters( + *args, + padding_idx=padding_idx, + scale_grad_by_freq=scale_grad_by_freq, + sparse=sparse, + ), + (weight, indices, out), + {}, + ) + + +def _embedding_with_parameters( + weight, + indices, + out, + *, + padding_idx, + scale_grad_by_freq, + sparse, + implementation_index, +): + infini.ops.embedding( + weight, + indices, + padding_idx, + scale_grad_by_freq, + sparse, + out, + implementation_index=implementation_index, + stream=get_stream(indices.device), + ) + + return out + + +def _torch_embedding_with_parameters( + weight, + indices, + out, + *, + padding_idx, + scale_grad_by_freq, + sparse, +): + result = torch.nn.functional.embedding( + indices, + weight, + padding_idx=padding_idx, + scale_grad_by_freq=scale_grad_by_freq, + sparse=sparse, + ) + out.copy_(result) + + return out diff --git a/tests/test_fused_add_rms_norm.py b/tests/test_fused_add_rms_norm.py new file mode 100644 index 000000000..77684bf5a --- /dev/null +++ b/tests/test_fused_add_rms_norm.py @@ -0,0 +1,80 @@ +import infini.ops +import pytest +import torch + +from tests.utils import clone_strided, get_stream, randn_strided + + +_TEST_CASES = ( + ((1, 4), None, None), + ((2, 4), None, None), + ((2, 2, 4), None, None), + ((2, 2, 4), (16, 8, 1), None), + ((2, 2, 2, 8), None, None), + ((16, 2048), None, None), + ((15, 3584), None, None), +) + + +@pytest.mark.parametrize( + "shape, input_strides, residual_strides", + _TEST_CASES, +) +@pytest.mark.parametrize("has_weight", (False, True)) +@pytest.mark.parametrize("epsilon", (1e-6, 1e-5)) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-4, 1e-4), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 2e-2, 1e-2), + ), +) +def test_fused_add_rms_norm( + shape, + input_strides, + residual_strides, + has_weight, + epsilon, + implementation_index, + dtype, + device, + rtol, + atol, +): + input = randn_strided(shape, input_strides, dtype=dtype, device=device) + residual = randn_strided(shape, residual_strides, dtype=dtype, device=device) + weight = torch.randn(shape[-1], dtype=dtype, device=device) if has_weight else None + expected_input = clone_strided(input) + expected_residual = clone_strided(residual) + + _torch_fused_add_rms_norm( + expected_input, + expected_residual, + weight, + epsilon, + ) + result = infini.ops.fused_add_rms_norm( + input, + residual, + weight, + epsilon, + implementation_index=implementation_index, + stream=get_stream(input.device), + ) + + assert result is None + torch.testing.assert_close(input, expected_input, rtol=rtol, atol=atol) + torch.testing.assert_close(residual, expected_residual, rtol=rtol, atol=atol) + + +def _torch_fused_add_rms_norm(input, residual, weight, epsilon): + summed = (input.to(torch.float32) + residual.to(torch.float32)).to(input.dtype) + variance = summed.to(torch.float32).pow(2).mean(dim=-1, keepdim=True) + normalized = summed.to(torch.float32) * torch.rsqrt(variance + epsilon) + + if weight is not None: + normalized *= weight.to(torch.float32) + + residual.copy_(summed) + input.copy_(normalized.to(input.dtype)) diff --git a/tests/test_linear.py b/tests/test_linear.py index 364ba5fcf..ba692d323 100644 --- a/tests/test_linear.py +++ b/tests/test_linear.py @@ -7,17 +7,16 @@ @pytest.mark.auto_act_and_assert @pytest.mark.parametrize( - "a_shape, b_shape, out_shape", + "input_shape, weight_shape, out_shape", ( - ((4, 64), (64, 32), (4, 32)), - ((2, 128), (128, 256), (2, 256)), + ((64,), (32, 64), (32,)), + ((4, 64), (32, 64), (4, 32)), + ((2, 128), (256, 128), (2, 256)), ((1, 4096), (4096, 4096), (1, 4096)), - ((2, 4, 64), (2, 64, 32), (2, 4, 32)), - ((4, 8, 128), (4, 128, 64), (4, 8, 64)), + ((2, 4, 64), (32, 64), (2, 4, 32)), + ((2, 3, 4, 64), (32, 64), (2, 3, 4, 32)), ), ) -@pytest.mark.parametrize("trans_a", (False, True)) -@pytest.mark.parametrize("trans_b", (False, True)) @pytest.mark.parametrize("has_bias", (False, True)) @pytest.mark.parametrize( ("dtype", "rtol", "atol"), @@ -28,63 +27,64 @@ ), ) def test_linear( - a_shape, - b_shape, + input_shape, + weight_shape, out_shape, - trans_a, - trans_b, has_bias, dtype, device, rtol, atol, ): - a = randn_strided(a_shape, None, dtype=dtype, device=device) - b = randn_strided(b_shape, None, dtype=dtype, device=device) - - if trans_a: - a = a.transpose(-2, -1) - - if trans_b: - b = b.transpose(-2, -1) + input = randn_strided(input_shape, None, dtype=dtype, device=device) + weight = randn_strided(weight_shape, None, dtype=dtype, device=device) # Bias shape is [N], the last dim of the output. bias = None if has_bias: - N = out_shape[-1] - bias = randn_strided((N,), None, dtype=dtype, device=device) + bias = randn_strided((out_shape[-1],), None, dtype=dtype, device=device) out = empty_strided(out_shape, None, dtype=dtype, device=device) return Payload( - lambda *args: _linear(*args, trans_a=trans_a, trans_b=trans_b), - lambda *args: _torch_linear(*args, trans_a=trans_a, trans_b=trans_b), - (a, b, bias, out), + _linear, + _torch_linear, + (input, weight, bias, out), {}, rtol=rtol, atol=atol, ) -def _linear(a, b, bias, out, trans_a=False, trans_b=False): - infini.ops.linear(a, b, bias, trans_a, trans_b, out, stream=get_stream(a.device)) +def _linear(input, weight, bias, out): + infini.ops.linear(input, weight, bias, out, stream=get_stream(input.device)) return out -def _torch_linear(a, b, bias, out, trans_a=False, trans_b=False): - if trans_a: - a = a.transpose(-2, -1) +def _torch_linear(input, weight, bias, out): + result = torch.nn.functional.linear( + input.float(), weight.float(), None if bias is None else bias.float() + ) - if trans_b: - b = b.transpose(-2, -1) + out.copy_(result.to(out.dtype)) - result = torch.matmul(a.float(), b.float()) + return out - if bias is not None: - result = result + bias.float() - out.copy_(result.to(out.dtype)) +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) +def test_linear_noncontiguous_weight(dtype, device): + input = randn_strided((3, 8), None, dtype=dtype, device=device) + weight = randn_strided((5, 8), (1, 5), dtype=dtype, device=device) + out = empty_strided((3, 5), None, dtype=dtype, device=device) - return out + return Payload( + _linear, + _torch_linear, + (input, weight, None, out), + {}, + rtol=1e-2, + atol=1e-2, + ) diff --git a/tests/test_matmul.py b/tests/test_matmul.py index fea3822a8..1845b077e 100644 --- a/tests/test_matmul.py +++ b/tests/test_matmul.py @@ -11,12 +11,14 @@ ( ((4, 64), (64, 32), (4, 32)), ((2, 128), (128, 256), (2, 256)), + ((64,), (64,), ()), + ((64,), (64, 32), (32,)), + ((4, 64), (64,), (4,)), ((2, 4, 64), (2, 64, 32), (2, 4, 32)), ((4, 8, 128), (4, 128, 64), (4, 8, 64)), + ((2, 1, 4, 64), (3, 64, 32), (2, 3, 4, 32)), ), ) -@pytest.mark.parametrize("trans_a", (False, True)) -@pytest.mark.parametrize("trans_b", (False, True)) @pytest.mark.parametrize( ("dtype", "rtol", "atol"), ( @@ -29,8 +31,6 @@ def test_matmul( a_shape, b_shape, c_shape, - trans_a, - trans_b, dtype, device, rtol, @@ -39,17 +39,11 @@ def test_matmul( a = randn_strided(a_shape, None, dtype=dtype, device=device) b = randn_strided(b_shape, None, dtype=dtype, device=device) - if trans_a: - a = a.transpose(-2, -1) - - if trans_b: - b = b.transpose(-2, -1) - c = empty_strided(c_shape, None, dtype=dtype, device=device) return Payload( - lambda *args: _matmul(*args, trans_a=trans_a, trans_b=trans_b), - lambda *args: _torch_matmul(*args, trans_a=trans_a, trans_b=trans_b), + _matmul, + _torch_matmul, (a, b, c), {}, rtol=rtol, @@ -57,19 +51,13 @@ def test_matmul( ) -def _matmul(a, b, c, trans_a=False, trans_b=False): - infini.ops.matmul(a, b, c, trans_a, trans_b, stream=get_stream(a.device)) +def _matmul(a, b, c): + infini.ops.matmul(a, b, c, stream=get_stream(a.device)) return c -def _torch_matmul(a, b, c, trans_a=False, trans_b=False): - if trans_a: - a = a.transpose(-2, -1) - - if trans_b: - b = b.transpose(-2, -1) - +def _torch_matmul(a, b, c): result = torch.matmul(a.float(), b.float()).to(c.dtype) c.copy_(result) diff --git a/tests/test_reshape_and_cache.py b/tests/test_reshape_and_cache.py new file mode 100644 index 000000000..773b20ac7 --- /dev/null +++ b/tests/test_reshape_and_cache.py @@ -0,0 +1,123 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream, randn_strided + + +@pytest.mark.parametrize("padded_source", (False, True)) +@pytest.mark.parametrize( + "kv_cache_dtype", + ("auto", "fp8", "fp8_e4m3", "fp8_e5m2"), +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-6, 1e-6), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 1e-3, 1e-3), + ), +) +def test_reshape_and_cache( + padded_source, + kv_cache_dtype, + implementation_index, + dtype, + device, + rtol, + atol, +): + num_tokens, num_heads, head_size = 4, 2, 32 + num_blocks, block_size = 3, 4 + source_strides = ( + (num_heads * head_size + 16, head_size, 1) if padded_source else None + ) + key = randn_strided( + (num_tokens, num_heads, head_size), + source_strides, + dtype=dtype, + device=device, + ) + value = randn_strided( + (num_tokens, num_heads, head_size), + source_strides, + dtype=dtype, + device=device, + ) + slot_mapping = torch.tensor((0, -1, 5, 10), dtype=torch.int64, device=device) + quantized = kv_cache_dtype != "auto" + cache_dtype = torch.uint8 if quantized else dtype + x = 16 if quantized else 16 // key.element_size() + key_cache = torch.zeros( + (num_blocks, num_heads, head_size // x, block_size, x), + dtype=cache_dtype, + device=device, + ) + value_cache = torch.zeros( + (num_blocks, num_heads, head_size, block_size), + dtype=cache_dtype, + device=device, + ) + k_scale = torch.tensor((0.5,), dtype=torch.float32, device=device) + v_scale = torch.tensor((0.25,), dtype=torch.float32, device=device) + expected_key_cache = key_cache.clone() + expected_value_cache = value_cache.clone() + + _torch_reshape_and_cache( + key, + value, + expected_key_cache, + expected_value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + ) + result = infini.ops.reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, + implementation_index=implementation_index, + stream=get_stream(key.device), + ) + + assert result is None + torch.testing.assert_close(key_cache, expected_key_cache, rtol=rtol, atol=atol) + torch.testing.assert_close(value_cache, expected_value_cache, rtol=rtol, atol=atol) + + +def _torch_reshape_and_cache( + key, + value, + key_cache, + value_cache, + slot_mapping, + kv_cache_dtype, + k_scale, + v_scale, +): + token_indices = torch.nonzero(slot_mapping >= 0).flatten() + slots = slot_mapping[token_indices] + block_size = key_cache.shape[3] + block_indices = torch.div(slots, block_size, rounding_mode="floor") + block_offsets = torch.remainder(slots, block_size) + x = key_cache.shape[4] + key_values = key[token_indices].view( + token_indices.numel(), key.shape[1], key.shape[2] // x, x + ) + value_values = value[token_indices] + + if kv_cache_dtype != "auto": + fp8_dtype = ( + torch.float8_e5m2 if kv_cache_dtype == "fp8_e5m2" else torch.float8_e4m3fn + ) + key_values = (key_values / k_scale).to(fp8_dtype).view(torch.uint8) + value_values = (value_values / v_scale).to(fp8_dtype).view(torch.uint8) + + key_cache[block_indices, :, :, block_offsets, :] = key_values + value_cache[block_indices, :, :, block_offsets] = value_values diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py new file mode 100644 index 000000000..682ac3cd4 --- /dev/null +++ b/tests/test_rotary_embedding.py @@ -0,0 +1,166 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +_TEST_CASES = ( + ((4,), False, True, 0, False), + ((4,), True, False, 2, False), + ((2, 2), False, False, 0, True), + ((2, 2), True, True, 2, True), +) + + +@pytest.mark.parametrize( + "positions_shape, structured, is_neox, rope_dim_offset, inverse", + _TEST_CASES, +) +@pytest.mark.parametrize("has_key", (False, True)) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 2e-2, 1e-2), + ), +) +def test_rotary_embedding( + positions_shape, + structured, + is_neox, + rope_dim_offset, + inverse, + has_key, + implementation_index, + dtype, + device, + rtol, + atol, +): + num_heads, num_kv_heads = 4, 2 + head_size, rot_dim = 12, 8 + positions = torch.tensor((0, 3, 5, 7), dtype=torch.int64, device=device).view( + positions_shape + ) + token_shape = positions_shape + query_shape = ( + (*token_shape, num_heads, head_size) + if structured + else (*token_shape, num_heads * head_size) + ) + key_shape = ( + (*token_shape, num_kv_heads, head_size) + if structured + else (*token_shape, num_kv_heads * head_size) + ) + query = torch.randn(query_shape, dtype=dtype, device=device) + key = torch.randn(key_shape, dtype=dtype, device=device) if has_key else None + cos_sin_cache = torch.randn(16, rot_dim, dtype=dtype, device=device) + expected_query = query.clone() + expected_key = key.clone() if key is not None else None + + _torch_rotary_embedding( + positions, + expected_query, + expected_key, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, + ) + args = ( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + ) + if rope_dim_offset == 0 and not inverse: + result = infini.ops.rotary_embedding( + *args, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + else: + result = infini.ops.rotary_embedding( + *args, + rope_dim_offset, + inverse, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + + assert result is None + torch.testing.assert_close(query, expected_query, rtol=rtol, atol=atol) + if key is not None: + torch.testing.assert_close(key, expected_key, rtol=rtol, atol=atol) + + +def _torch_rotary_embedding( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, +): + _apply_rotary( + positions, + query, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, + ) + + if key is not None: + _apply_rotary( + positions, + key, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, + ) + + +def _apply_rotary( + positions, + data, + head_size, + cos_sin_cache, + is_neox, + rope_dim_offset, + inverse, +): + num_tokens = positions.numel() + rot_dim = cos_sin_cache.shape[1] + embed_dim = rot_dim // 2 + num_heads = data.numel() // num_tokens // head_size + data_view = data.view(num_tokens, num_heads, head_size) + cache = cos_sin_cache[positions.flatten()] + cos = cache[:, :embed_dim].unsqueeze(1).float() + sin = cache[:, embed_dim:].unsqueeze(1).float() + + if inverse: + sin = -sin + + rotary = data_view[..., rope_dim_offset : rope_dim_offset + rot_dim] + if is_neox: + x = rotary[..., :embed_dim].float().clone() + y = rotary[..., embed_dim:].float().clone() + rotary[..., :embed_dim].copy_(x * cos - y * sin) + rotary[..., embed_dim:].copy_(y * cos + x * sin) + else: + x = rotary[..., 0::2].float().clone() + y = rotary[..., 1::2].float().clone() + rotary[..., 0::2].copy_(x * cos - y * sin) + rotary[..., 1::2].copy_(y * cos + x * sin) diff --git a/tests/test_scaled_dot_product_attention.py b/tests/test_scaled_dot_product_attention.py new file mode 100644 index 000000000..f92dfa9f1 --- /dev/null +++ b/tests/test_scaled_dot_product_attention.py @@ -0,0 +1,139 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +_TEST_CASES = ( + (4, 4, "none", 0.0, False, None, False), + (3, 5, "float", 0.0, False, 0.5, False), + (4, 4, "bool", 0.0, False, None, False), + (4, 4, "none", 0.0, True, None, False), + (3, 5, "none", 0.0, False, None, True), + (4, 4, "none", 0.2, False, None, False), +) + + +@pytest.mark.parametrize( + "query_length, key_length, mask_kind, dropout_p, is_causal, scale, enable_gqa", + _TEST_CASES, +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 2e-2, 1e-2), + ), +) +def test_scaled_dot_product_attention( + query_length, + key_length, + mask_kind, + dropout_p, + is_causal, + scale, + enable_gqa, + implementation_index, + dtype, + device, + rtol, + atol, +): + batch_size, query_heads, head_size, value_size = 2, 4, 8, 6 + kv_heads = 2 if enable_gqa else query_heads + query = torch.randn( + batch_size, + query_heads, + query_length, + head_size, + dtype=dtype, + device=device, + ) + key = torch.randn( + batch_size, + kv_heads, + key_length, + head_size, + dtype=dtype, + device=device, + ) + value = torch.randn( + batch_size, + kv_heads, + key_length, + value_size, + dtype=dtype, + device=device, + ) + attn_mask = _make_mask(mask_kind, query_length, key_length, dtype, device) + out = torch.empty( + batch_size, + query_heads, + query_length, + value_size, + dtype=dtype, + device=device, + ) + + torch.manual_seed(1234) + expected = torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + ) + torch.manual_seed(1234) + if ( + attn_mask is None + and dropout_p == 0.0 + and not is_causal + and scale is None + and not enable_gqa + ): + result = infini.ops.scaled_dot_product_attention( + query, + key, + value, + out, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + else: + result = infini.ops.scaled_dot_product_attention( + query, + key, + value, + attn_mask, + dropout_p, + is_causal, + scale, + enable_gqa, + out, + implementation_index=implementation_index, + stream=get_stream(query.device), + ) + + assert result is None + torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) + + +def _make_mask(mask_kind, query_length, key_length, dtype, device): + if mask_kind == "none": + return None + + if mask_kind == "bool": + mask = torch.ones(query_length, key_length, dtype=torch.bool, device=device) + mask[0, -1] = False + + return mask + + mask = torch.zeros(query_length, key_length, dtype=dtype, device=device) + mask[0, -1] = -3.0 + + return mask diff --git a/tests/test_scaled_softmax_infinilm.py b/tests/test_scaled_softmax_infinilm.py new file mode 100644 index 000000000..cfb19700b --- /dev/null +++ b/tests/test_scaled_softmax_infinilm.py @@ -0,0 +1,66 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randn_strided + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + "shape", + ( + (1, 7), + (3, 11), + (16, 512), + ), +) +@pytest.mark.parametrize("scale", (1.0, 0.5, 1.7)) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-5, 1e-5), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 1e-2, 1e-2), + ), +) +def test_scaled_softmax_infinilm( + shape, + scale, + dtype, + device, + implementation_index, + rtol, + atol, +): + input_tensor = randn_strided(shape, None, dtype=dtype, device=device) + out = empty_strided(shape, None, dtype=dtype, device=device) + + return Payload( + _scaled_softmax_infinilm, + _torch_scaled_softmax, + (input_tensor, out), + {"scale": scale, "implementation_index": implementation_index}, + rtol=rtol, + atol=atol, + ) + + +def _scaled_softmax_infinilm(input_tensor, out, *, scale, implementation_index): + infini.ops.scaled_softmax_infinilm( + input_tensor, + scale, + out, + stream=get_stream(input_tensor.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_scaled_softmax(input_tensor, out, *, scale, implementation_index): + del implementation_index + + result = torch.nn.functional.softmax(input_tensor.to(torch.float32) * scale, dim=-1) + out.copy_(result.to(input_tensor.dtype)) + + return out diff --git a/tests/test_silu_and_mul.py b/tests/test_silu_and_mul.py new file mode 100644 index 000000000..ea6f0a798 --- /dev/null +++ b/tests/test_silu_and_mul.py @@ -0,0 +1,72 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, rand_strided + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + "out_shape, input_strides, out_strides", + ( + ((13, 4), None, None), + ((13, 4), (12, 1), (6, 1)), + ((13, 4, 4), None, None), + ((13, 4, 4), (48, 10, 1), (24, 6, 1)), + ((16, 5632), None, None), + ((16, 5632), (13312, 1), (6656, 1)), + ((4, 4, 5632), None, None), + ((4, 4, 5632), (53248, 13312, 1), (26624, 6656, 1)), + ), +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-7, 1e-7), + (torch.float16, 1e-3, 1e-3), + (torch.bfloat16, 1e-2, 5e-3), + ), +) +def test_silu_and_mul( + out_shape, + input_strides, + out_strides, + implementation_index, + dtype, + device, + rtol, + atol, +): + input_shape = (*out_shape[:-1], out_shape[-1] * 2) + input = rand_strided(input_shape, input_strides, dtype=dtype, device=device) + out = empty_strided(out_shape, out_strides, dtype=dtype, device=device) + + return Payload( + lambda *args, **kwargs: _silu_and_mul( + *args, **kwargs, implementation_index=implementation_index + ), + _torch_silu_and_mul, + (input, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _silu_and_mul(input, out, implementation_index=0): + infini.ops.silu_and_mul( + input, + out, + implementation_index=implementation_index, + stream=get_stream(input.device), + ) + + return out + + +def _torch_silu_and_mul(input, out): + hidden_size = input.shape[-1] // 2 + gate = input[..., :hidden_size] + up = input[..., hidden_size:] + + return torch.mul(torch.nn.functional.silu(gate), up, out=out) diff --git a/tests/test_top_k_top_p_sample_infinilm.py b/tests/test_top_k_top_p_sample_infinilm.py new file mode 100644 index 000000000..f9a4e2705 --- /dev/null +++ b/tests/test_top_k_top_p_sample_infinilm.py @@ -0,0 +1,92 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_top_k_top_p_sample_infinilm_reproducible( + dtype, + device, + implementation_index, +): + logits = torch.zeros((64, 16), dtype=dtype, device=device) + first = torch.empty((64,), dtype=torch.int32, device=device) + second = torch.empty_like(first) + different_seed = torch.empty_like(first) + + _top_k_top_p_sample_infinilm( + logits, None, None, 1234, 9, first, implementation_index + ) + _top_k_top_p_sample_infinilm( + logits, None, None, 1234, 9, second, implementation_index + ) + _top_k_top_p_sample_infinilm( + logits, None, None, 5678, 9, different_seed, implementation_index + ) + + assert torch.equal(first, second) + assert not torch.equal(first, different_seed) + assert first.dtype == torch.int32 + assert torch.all((first >= 0) & (first < logits.shape[1])) + + +@pytest.mark.parametrize( + "k_value, p_value, allowed", + ( + (3, None, (0, 1, 2)), + (None, 0.6, (0,)), + (3, 0.8, (0, 1)), + ), +) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_top_k_top_p_sample_infinilm_filters( + k_value, + p_value, + allowed, + dtype, + device, + implementation_index, +): + logits = torch.full((32, 16), -10.0, dtype=dtype, device=device) + logits[:, 0] = 5.0 + logits[:, 1] = 4.0 + logits[:, 2] = 3.0 + k = ( + torch.tensor((k_value,), dtype=torch.int64, device="cpu") + if k_value is not None + else None + ) + p = ( + torch.tensor((p_value,), dtype=torch.float32, device="cpu") + if p_value is not None + else None + ) + out = torch.empty((32,), dtype=torch.int32, device=device) + + _top_k_top_p_sample_infinilm(logits, k, p, 1234, 0, out, implementation_index) + + allowed_tensor = torch.tensor(allowed, dtype=torch.int32, device=device) + assert torch.all(torch.isin(out, allowed_tensor)) + + +def _top_k_top_p_sample_infinilm( + logits, + k, + p, + seed, + offset, + out, + implementation_index, +): + infini.ops.top_k_top_p_sample_infinilm( + logits, + k, + p, + seed, + offset, + out, + stream=get_stream(logits.device), + implementation_index=implementation_index, + )