From 98d179221a9c907fa607bb2c071c19d0fc9cdd6f Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 10:15:44 +0800 Subject: [PATCH 01/22] feat(add): align tensor overload with PyTorch --- src/base/add.h | 96 ++++++++++++++++++++++++++---- src/native/ascend/ops/add/kernel.h | 14 ++++- src/native/cpu/ops/add/add.h | 18 ++++-- src/native/cuda/ops/add/kernel.cuh | 16 +++-- src/native/cuda/ops/add/kernel.h | 12 ++-- src/torch/ops/add/add.cc | 14 ++++- src/torch/ops/add/add.h | 5 +- tests/test_add.py | 56 +++++++++++++++++ 8 files changed, 199 insertions(+), 32 deletions(-) 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/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/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/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/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/tests/test_add.py b/tests/test_add.py index e2266c30d..f54073e28 100644 --- a/tests/test_add.py +++ b/tests/test_add.py @@ -109,3 +109,59 @@ 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 From d0ba912405b1800749f78e10909b4d845b09e49b Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 10:43:45 +0800 Subject: [PATCH 02/22] feat(cat): align tensor-list interface with PyTorch --- scripts/torch_ops.yaml | 1 + src/base/cat.h | 38 ++++++++++++--- src/native/ascend/ops/cat/kernel.h | 28 ++++------- src/native/cpu/ops/cat/cat.h | 75 +++++++++++------------------- src/torch/tensor_.h | 4 ++ tests/test_cat.py | 35 ++++++++++++-- 6 files changed, 103 insertions(+), 78 deletions(-) diff --git a/scripts/torch_ops.yaml b/scripts/torch_ops.yaml index e0b7ece5f..8313da583 100644 --- a/scripts/torch_ops.yaml +++ b/scripts/torch_ops.yaml @@ -68,6 +68,7 @@ - bitwise_xor - bmm - bucketize +- cat - ceil - chain_matmul - cholesky 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/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/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/torch/tensor_.h b/src/torch/tensor_.h index 5780babd9..3a7396b0b 100644 --- a/src/torch/tensor_.h +++ b/src/torch/tensor_.h @@ -97,6 +97,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/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, + ) From 8a074530f3f50191a85cb9b795d4bde2d9df4905 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 10:59:05 +0800 Subject: [PATCH 03/22] feat(matmul): align interface with PyTorch --- scripts/torch_ops.yaml | 1 + src/base/matmul.h | 47 +++++++++++++++++---------- src/native/ascend/ops/matmul/kernel.h | 40 ++++++++++++----------- tests/conftest.py | 7 ++-- tests/test_matmul.py | 30 +++++------------ 5 files changed, 62 insertions(+), 63 deletions(-) diff --git a/scripts/torch_ops.yaml b/scripts/torch_ops.yaml index 8313da583..ef5eefb0c 100644 --- a/scripts/torch_ops.yaml +++ b/scripts/torch_ops.yaml @@ -260,6 +260,7 @@ - lu_solve - lu_unpack - masked_select +- matmul - matrix_power - max - max_pool2d_with_indices 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/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/tests/conftest.py b/tests/conftest.py index f0be7431f..83ea78833 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -336,11 +336,8 @@ 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 ) 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) From 22f6135c1c6cd36c49e89b3620a7c04a05d1d976 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 11:15:00 +0800 Subject: [PATCH 04/22] feat(linear): align weight semantics with PyTorch --- scripts/torch_ops.yaml | 1 + src/base/linear.h | 84 ++++++++++++++------- src/native/ascend/ops/linear/kernel.h | 70 ++++++++--------- src/native/cpu/ops/linear/linear.h | 104 ++++++++++---------------- tests/conftest.py | 8 +- tests/test_linear.py | 76 ++++++++++--------- 6 files changed, 170 insertions(+), 173 deletions(-) diff --git a/scripts/torch_ops.yaml b/scripts/torch_ops.yaml index ef5eefb0c..3fc72844c 100644 --- a/scripts/torch_ops.yaml +++ b/scripts/torch_ops.yaml @@ -236,6 +236,7 @@ - linalg_tensorsolve - linalg_vecdot - linalg_vector_norm +- linear - linspace - log - log10 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/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/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/tests/conftest.py b/tests/conftest.py index 83ea78833..230801748 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -343,16 +343,14 @@ def _is_smoke_matmul_case(params): 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 diff --git a/tests/test_linear.py b/tests/test_linear.py index 364ba5fcf..b51b6414d 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,68 @@ ), ) 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, + ) From 82ffc5c20863a0b655fe026fc532941fa2bc10ad Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 11:31:40 +0800 Subject: [PATCH 05/22] feat(embedding): align interface with PyTorch --- src/base/embedding.h | 70 +++++++++++++----- src/native/ascend/ops/embedding/kernel.h | 28 ++++--- src/native/cuda/ops/embedding/kernel.h | 59 ++++++++------- tests/test_embedding.py | 93 ++++++++++++++++++++++-- 4 files changed, 190 insertions(+), 60 deletions(-) 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/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/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/tests/test_embedding.py b/tests/test_embedding.py index 98fe6bad3..499a5776b 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,84 @@ 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 From fab45ff69e43ef8de5267f390f266c6201c6092d Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 13:47:08 +0800 Subject: [PATCH 06/22] feat(rms-norm): align fused add interface with vLLM --- src/base/add_rms_norm.h | 103 --------- src/base/fused_add_rms_norm.h | 78 +++++++ .../cuda/iluvatar/ops/add_rms_norm/kernel.h | 21 -- .../iluvatar/ops/fused_add_rms_norm/kernel.h | 22 ++ .../cuda/metax/ops/add_rms_norm/kernel.h | 21 -- .../metax/ops/fused_add_rms_norm/kernel.h | 21 ++ .../cuda/moore/ops/add_rms_norm/kernel.h | 25 --- .../moore/ops/fused_add_rms_norm/kernel.h | 25 +++ .../cuda/nvidia/ops/add_rms_norm/kernel.h | 21 -- .../nvidia/ops/fused_add_rms_norm/kernel.h | 22 ++ src/native/cuda/ops/add_rms_norm/kernel.cuh | 80 ------- src/native/cuda/ops/add_rms_norm/kernel.h | 84 -------- .../cuda/ops/fused_add_rms_norm/kernel.cuh | 69 ++++++ .../cuda/ops/fused_add_rms_norm/kernel.h | 53 +++++ tests/test_add_rms_norm.py | 199 ------------------ tests/test_fused_add_rms_norm.py | 82 ++++++++ 16 files changed, 372 insertions(+), 554 deletions(-) delete mode 100644 src/base/add_rms_norm.h create mode 100644 src/base/fused_add_rms_norm.h delete mode 100644 src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/iluvatar/ops/fused_add_rms_norm/kernel.h delete mode 100644 src/native/cuda/metax/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/metax/ops/fused_add_rms_norm/kernel.h delete mode 100644 src/native/cuda/moore/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/moore/ops/fused_add_rms_norm/kernel.h delete mode 100644 src/native/cuda/nvidia/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/nvidia/ops/fused_add_rms_norm/kernel.h delete mode 100644 src/native/cuda/ops/add_rms_norm/kernel.cuh delete mode 100644 src/native/cuda/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/ops/fused_add_rms_norm/kernel.cuh create mode 100644 src/native/cuda/ops/fused_add_rms_norm/kernel.h delete mode 100644 tests/test_add_rms_norm.py create mode 100644 tests/test_fused_add_rms_norm.py diff --git a/src/base/add_rms_norm.h b/src/base/add_rms_norm.h deleted file mode 100644 index 0241f7fce..000000000 --- a/src/base/add_rms_norm.h +++ /dev/null @@ -1,103 +0,0 @@ -#ifndef INFINI_OPS_BASE_ADD_RMS_NORM_H_ -#define INFINI_OPS_BASE_ADD_RMS_NORM_H_ - -#include -#include - -#include "operator.h" -#include "tensor.h" - -namespace infini::ops { - -// Fused add + RMSNorm aligned with vLLM `fused_add_rms_norm`. -class AddRmsNorm : public Operator { - public: - AddRmsNorm(const Tensor input, const Tensor residual, const Tensor weight, - std::optional eps, Tensor out, Tensor residual_out) - : input_shape_{input.shape()}, - out_shape_{out.shape()}, - input_strides_{input.strides()}, - residual_strides_{residual.strides()}, - out_strides_{out.strides()}, - residual_out_strides_{residual_out.strides()}, - eps_{eps.value_or(1e-6f)}, - dim_{out.size(-1)}, - ndim_{out.ndim()}, - batch_size_{ndim_ == 2 ? out.size(-2) : out.size(-3)}, - nhead_{ndim_ == 2 ? 1 : out.size(-2)} { - assert((ndim_ == 2 || ndim_ == 3) && - "`AddRmsNorm` supports 2D or 3D tensors only"); - assert(input.shape() == out.shape() && - "`AddRmsNorm` requires `input` and `out` to have the same shape"); - assert(input.shape() == residual.shape() && - "`AddRmsNorm` requires `input` and `residual` to have the same " - "shape"); - assert(input.shape() == residual_out.shape() && - "`AddRmsNorm` requires `input` and `residual_out` to have the " - "same shape"); - assert(weight.ndim() == 1 && weight.size(-1) == dim_ && - "`AddRmsNorm` requires 1D `weight` with size equal to the " - "normalized dimension"); - assert(input.dtype() == out.dtype() && - "`AddRmsNorm` requires `input` and `out` to have the same dtype"); - assert(input.dtype() == residual.dtype() && - "`AddRmsNorm` requires `input` and `residual` to have the same " - "dtype"); - assert(input.dtype() == residual_out.dtype() && - "`AddRmsNorm` requires `input` and `residual_out` to have the same " - "dtype"); - // The CUDA kernel indexes the normalized dimension with stride 1. - assert(input.stride(-1) == 1 && - "`AddRmsNorm` requires the last dimension of `input` to be " - "contiguous"); - assert(residual.stride(-1) == 1 && - "`AddRmsNorm` requires the last dimension of `residual` to be " - "contiguous"); - assert(out.stride(-1) == 1 && - "`AddRmsNorm` requires the last dimension of `out` to be " - "contiguous"); - assert(residual_out.stride(-1) == 1 && - "`AddRmsNorm` requires the last dimension of `residual_out` to be " - "contiguous"); - assert(weight.stride(-1) == 1 && - "`AddRmsNorm` requires the last dimension of `weight` to be " - "contiguous"); - } - - virtual void operator()(const Tensor input, const Tensor residual, - const Tensor weight, std::optional eps, - Tensor out, Tensor residual_out) const = 0; - - virtual void operator()(const Tensor input, const Tensor residual, - const Tensor weight, Tensor out, - Tensor residual_out) const { - return operator()(input, residual, weight, std::nullopt, out, residual_out); - } - - protected: - Tensor::Shape input_shape_; - - Tensor::Shape out_shape_; - - Tensor::Strides input_strides_; - - Tensor::Strides residual_strides_; - - Tensor::Strides out_strides_; - - Tensor::Strides residual_out_strides_; - - float eps_{1e-6f}; - - Tensor::Size dim_{0}; - - Tensor::Size ndim_{0}; - - Tensor::Size batch_size_{0}; - - Tensor::Size nhead_{1}; -}; - -} // namespace infini::ops - -#endif 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/native/cuda/iluvatar/ops/add_rms_norm/kernel.h b/src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h deleted file mode 100644 index 828b93bbc..000000000 --- a/src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_ILUVATAR_ADD_RMS_NORM_KERNEL_H_ -#define INFINI_OPS_ILUVATAR_ADD_RMS_NORM_KERNEL_H_ - -#include - -#include "native/cuda/iluvatar/caster.cuh" -#include "native/cuda/iluvatar/runtime_.h" -#include "native/cuda/ops/add_rms_norm/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaAddRmsNorm> { - public: - using CudaAddRmsNorm>::CudaAddRmsNorm; -}; - -} // 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/metax/ops/add_rms_norm/kernel.h b/src/native/cuda/metax/ops/add_rms_norm/kernel.h deleted file mode 100644 index 564ceba61..000000000 --- a/src/native/cuda/metax/ops/add_rms_norm/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_METAX_ADD_RMS_NORM_KERNEL_H_ -#define INFINI_OPS_METAX_ADD_RMS_NORM_KERNEL_H_ - -#include - -#include "native/cuda/metax/caster.cuh" -#include "native/cuda/metax/runtime_.h" -#include "native/cuda/ops/add_rms_norm/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaAddRmsNorm> { - public: - using CudaAddRmsNorm>::CudaAddRmsNorm; -}; - -} // 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/moore/ops/add_rms_norm/kernel.h b/src/native/cuda/moore/ops/add_rms_norm/kernel.h deleted file mode 100644 index 8d3f5cc79..000000000 --- a/src/native/cuda/moore/ops/add_rms_norm/kernel.h +++ /dev/null @@ -1,25 +0,0 @@ -#ifndef INFINI_OPS_MOORE_ADD_RMS_NORM_KERNEL_H_ -#define INFINI_OPS_MOORE_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/add_rms_norm/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaAddRmsNorm> { - public: - using CudaAddRmsNorm>::CudaAddRmsNorm; -}; - -} // 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/nvidia/ops/add_rms_norm/kernel.h b/src/native/cuda/nvidia/ops/add_rms_norm/kernel.h deleted file mode 100644 index 2bb6f6051..000000000 --- a/src/native/cuda/nvidia/ops/add_rms_norm/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_ADD_RMS_NORM_KERNEL_H_ -#define INFINI_OPS_NVIDIA_ADD_RMS_NORM_KERNEL_H_ - -#include - -#include "native/cuda/nvidia/caster.cuh" -#include "native/cuda/nvidia/runtime_.h" -#include "native/cuda/ops/add_rms_norm/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaAddRmsNorm> { - public: - using CudaAddRmsNorm>::CudaAddRmsNorm; -}; - -} // 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/ops/add_rms_norm/kernel.cuh b/src/native/cuda/ops/add_rms_norm/kernel.cuh deleted file mode 100644 index 7fd8215e1..000000000 --- a/src/native/cuda/ops/add_rms_norm/kernel.cuh +++ /dev/null @@ -1,80 +0,0 @@ -#ifndef INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_CUH_ -#define INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_CUH_ - -#include -#include -#include - -#include "native/cuda/caster.cuh" -#include "native/cuda/kernel_commons.cuh" - -namespace infini::ops { -namespace 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 add_rms_norm_detail - -template -__global__ void AddRmsNormKernel( - TData* __restrict__ y, int64_t stride_y_batch, int64_t stride_y_nhead, - TData* __restrict__ residual_out, int64_t stride_residual_out_batch, - int64_t stride_residual_out_nhead, const TData* __restrict__ input, - int64_t stride_input_batch, int64_t stride_input_nhead, - const TData* __restrict__ residual, int64_t stride_residual_batch, - int64_t stride_residual_nhead, const TWeight* __restrict__ w, size_t nhead, - size_t dim, float epsilon) { - size_t batch_idx = blockIdx.x / nhead; - size_t head_idx = blockIdx.x % nhead; - - auto y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; - auto input_ptr = - input + batch_idx * stride_input_batch + head_idx * stride_input_nhead; - auto residual_ptr = residual + batch_idx * stride_residual_batch + - head_idx * stride_residual_nhead; - auto w_ptr = w; - auto residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + - head_idx * stride_residual_out_nhead; - - 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_out_ptr[i] = Caster::template Cast(sum_val); - } - - TCompute sum_squared = - add_rms_norm_detail::SumSquared( - residual_out_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 sum_val = - Caster::template Cast(residual_out_ptr[i]); - y_ptr[i] = Caster::template Cast( - sum_val * Caster::template Cast(w_ptr[i]) * rms); - } -} - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/ops/add_rms_norm/kernel.h b/src/native/cuda/ops/add_rms_norm/kernel.h deleted file mode 100644 index a2c0c7cbf..000000000 --- a/src/native/cuda/ops/add_rms_norm/kernel.h +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_H_ -#define INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_H_ - -#include -#include -#include - -#include "base/add_rms_norm.h" -#include "data_type.h" -#include "dispatcher.h" -#include "native/cuda/kernel_commons.cuh" -#include "native/cuda/ops/add_rms_norm/kernel.cuh" -#include "native/cuda/runtime_utils.h" - -namespace infini::ops { - -template -class CudaAddRmsNorm : public AddRmsNorm { - public: - using AddRmsNorm::AddRmsNorm; - - void operator()(const Tensor input, const Tensor residual, - const Tensor weight, std::optional eps, Tensor out, - Tensor residual_out) const override { - const float kernel_eps = eps.value_or(eps_); - auto cuda_stream = - static_cast(stream_ ? stream_ : 0); - - auto stride_input_batch = input_strides_.size() > 1 ? input_strides_[0] : 0; - auto stride_input_nhead = - input_strides_.size() > 1 ? input_strides_[1] : input_strides_[0]; - auto stride_residual_batch = - residual_strides_.size() > 1 ? residual_strides_[0] : 0; - auto stride_residual_nhead = residual_strides_.size() > 1 - ? residual_strides_[1] - : residual_strides_[0]; - auto stride_out_batch = out_strides_.size() > 1 ? out_strides_[0] : 0; - auto stride_out_nhead = - out_strides_.size() > 1 ? out_strides_[1] : out_strides_[0]; - auto stride_residual_out_batch = - residual_out_strides_.size() > 1 ? residual_out_strides_[0] : 0; - auto stride_residual_out_nhead = residual_out_strides_.size() > 1 - ? residual_out_strides_[1] - : residual_out_strides_[0]; - - uint32_t num_blocks = static_cast(batch_size_ * nhead_); - - assert(out.dtype() == input.dtype() && out.dtype() == residual.dtype() && - out.dtype() == residual_out.dtype() && - "`CudaAddRmsNorm` requires input, residual, out and residual_out to " - "have the same dtype"); - - int block_size = RuntimeUtils::GetOptimalBlockSize(); - - DispatchFunc, ReducedFloatTypes>, - ConcatType, ReducedFloatTypes>, - AllCudaBlockSizes>( - {static_cast(out.dtype()), - static_cast(weight.dtype()), block_size}, - [&](auto list_tag) { - using T = TypeMapType(list_tag)>; - using TWeight = - TypeMapType(list_tag)>; - constexpr int kBlockSize = ListGet<2>(list_tag); - - AddRmsNormKernel - <<>>( - reinterpret_cast(out.data()), stride_out_batch, - stride_out_nhead, reinterpret_cast(residual_out.data()), - stride_residual_out_batch, stride_residual_out_nhead, - reinterpret_cast(input.data()), stride_input_batch, - stride_input_nhead, - reinterpret_cast(residual.data()), - stride_residual_batch, stride_residual_nhead, - reinterpret_cast(weight.data()), nhead_, dim_, - kernel_eps); - }, - "CudaAddRmsNorm::operator()"); - } -}; - -} // namespace infini::ops - -#endif 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/tests/test_add_rms_norm.py b/tests/test_add_rms_norm.py deleted file mode 100644 index 1d94f7455..000000000 --- a/tests/test_add_rms_norm.py +++ /dev/null @@ -1,199 +0,0 @@ -import infini.ops -import pytest -import torch - -from tests.utils import Payload, empty_strided, get_stream, randn_strided - - -# Format: (input_shape, weight_shape, input_strides, residual_strides, weight_strides, out_strides); -# input/residual/residual_out share input_shape. -_TEST_CASES = ( - ((1, 4), (4,), None, None, None, None), - ((2, 4), (4,), None, None, None, None), - ((2, 2, 4), (4,), None, None, None, None), - ((2, 2, 4), (4,), (12, 8, 1), (12, 8, 1), None, (12, 8, 1)), - ((16, 2048), (2048,), None, None, None, None), - ((16, 2048), (2048,), (4096, 1), (4096, 1), None, (4096, 1)), - ((15, 3584), (3584,), None, None, None, None), - ((4, 4, 2048), (2048,), None, None, None, None), - ((4, 4, 2048), (2048,), (2048, 8192, 1), (2048, 8192, 1), None, (2048, 8192, 1)), - ( - (4, 4, 2048), - (2048,), - (16384, 4096, 1), - (16384, 4096, 1), - None, - (16384, 4096, 1), - ), -) - - -@pytest.mark.auto_act_and_assert -@pytest.mark.parametrize("check_output", ("out", "residual_out")) -@pytest.mark.parametrize( - "input_shape, weight_shape, input_strides, residual_strides, weight_strides, out_strides", - _TEST_CASES, -) -@pytest.mark.parametrize("eps", (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_add_rms_norm( - check_output, - input_shape, - weight_shape, - input_strides, - residual_strides, - weight_strides, - out_strides, - eps, - implementation_index, - dtype, - device, - rtol, - atol, -): - input = randn_strided(input_shape, input_strides, dtype=dtype, device=device) - residual = randn_strided(input_shape, residual_strides, dtype=dtype, device=device) - weight = randn_strided(weight_shape, weight_strides, dtype=dtype, device=device) - residual_out = empty_strided(input_shape, input_strides, dtype=dtype, device=device) - out = empty_strided(input_shape, out_strides, dtype=dtype, device=device) - - if check_output == "out": - func = _add_rms_norm_out - ref = _torch_add_rms_norm_out - else: - func = _add_rms_norm_residual_out - ref = _torch_add_rms_norm_residual_out - - return Payload( - lambda *args, **kwargs: func( - *args, **kwargs, implementation_index=implementation_index - ), - ref, - (input, residual, weight), - {"eps": eps, "out": out, "residual_out": residual_out}, - rtol=rtol, - atol=atol, - ) - - -def _add_rms_norm( - input, - residual, - weight, - *, - eps=1e-6, - out=None, - residual_out=None, - implementation_index=0, -): - infini.ops.add_rms_norm( - input, - residual, - weight, - eps, - out, - residual_out, - implementation_index=implementation_index, - stream=get_stream(input.device), - ) - - -def _add_rms_norm_out( - input, - residual, - weight, - *, - eps=1e-6, - out=None, - residual_out=None, - implementation_index=0, -): - _add_rms_norm( - input, - residual, - weight, - eps=eps, - out=out, - residual_out=residual_out, - implementation_index=implementation_index, - ) - - return out - - -def _add_rms_norm_residual_out( - input, - residual, - weight, - *, - eps=1e-6, - out=None, - residual_out=None, - implementation_index=0, -): - _add_rms_norm( - input, - residual, - weight, - eps=eps, - out=out, - residual_out=residual_out, - implementation_index=implementation_index, - ) - - return residual_out - - -def _torch_add_rms_norm( - input, residual, weight, *, eps=1e-6, out=None, residual_out=None -): - """Reference aligned with vLLM `fused_add_rms_norm` (ignoring `variance_size`).""" - orig_dtype = input.dtype - x = input.to(torch.float32) - x = x + residual.to(torch.float32) - add_result = x.to(orig_dtype) - - variance = x.pow(2).mean(dim=-1, keepdim=True) - x = x * torch.rsqrt(variance + eps) - if weight is not None: - x = x.to(weight.dtype) * weight - normalized_result = x.to(orig_dtype) - - if out is not None: - out.copy_(normalized_result) - else: - out = normalized_result - - if residual_out is not None: - residual_out.copy_(add_result) - else: - residual_out = add_result - - return out, residual_out - - -def _torch_add_rms_norm_out( - input, residual, weight, *, eps=1e-6, out=None, residual_out=None -): - out, _ = _torch_add_rms_norm( - input, residual, weight, eps=eps, out=out, residual_out=residual_out - ) - - return out - - -def _torch_add_rms_norm_residual_out( - input, residual, weight, *, eps=1e-6, out=None, residual_out=None -): - _, residual_out = _torch_add_rms_norm( - input, residual, weight, eps=eps, out=out, residual_out=residual_out - ) - - return residual_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..377a7bd85 --- /dev/null +++ b/tests/test_fused_add_rms_norm.py @@ -0,0 +1,82 @@ +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)) From 7d862f72bc965da7715a9213f0b233d7d4debec3 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 17:05:12 +0800 Subject: [PATCH 07/22] feat(cache): align reshape and rotary interfaces with vLLM --- src/base/reshape_and_cache.h | 161 ++++++++++---- src/base/rotary_embedding.h | 198 +++++++++++++----- .../nvidia/ops/reshape_and_cache/kernel.cuh | 76 +++++++ .../nvidia/ops/reshape_and_cache/kernel.h | 73 +++++++ .../cuda/nvidia/ops/rotary_embedding/kernel.h | 20 ++ .../cuda/ops/rotary_embedding/kernel.cuh | 66 ++++++ src/native/cuda/ops/rotary_embedding/kernel.h | 65 ++++++ .../reshape_and_cache/reshape_and_cache.cc | 91 ++++++++ .../ops/reshape_and_cache/reshape_and_cache.h | 26 +++ .../ops/rotary_embedding/rotary_embedding.cc | 89 ++++++++ .../ops/rotary_embedding/rotary_embedding.h | 27 +++ tests/test_reshape_and_cache.py | 129 ++++++++++++ tests/test_rotary_embedding.py | 167 +++++++++++++++ 13 files changed, 1097 insertions(+), 91 deletions(-) create mode 100644 src/native/cuda/nvidia/ops/reshape_and_cache/kernel.cuh create mode 100644 src/native/cuda/nvidia/ops/reshape_and_cache/kernel.h create mode 100644 src/native/cuda/nvidia/ops/rotary_embedding/kernel.h create mode 100644 src/native/cuda/ops/rotary_embedding/kernel.cuh create mode 100644 src/native/cuda/ops/rotary_embedding/kernel.h create mode 100644 src/torch/ops/reshape_and_cache/reshape_and_cache.cc create mode 100644 src/torch/ops/reshape_and_cache/reshape_and_cache.h create mode 100644 src/torch/ops/rotary_embedding/rotary_embedding.cc create mode 100644 src/torch/ops/rotary_embedding/rotary_embedding.h create mode 100644 tests/test_reshape_and_cache.py create mode 100644 tests/test_rotary_embedding.py 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..0f11d2f45 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, + const 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_{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, + const 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/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/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..8566303a5 --- /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, + const 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/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..c91c05363 --- /dev/null +++ b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc @@ -0,0 +1,91 @@ +#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; +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..b30f8c33b --- /dev/null +++ b/src/torch/ops/rotary_embedding/rotary_embedding.cc @@ -0,0 +1,89 @@ +#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, const 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, const 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; +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..ec1d6811a --- /dev/null +++ b/src/torch/ops/rotary_embedding/rotary_embedding.h @@ -0,0 +1,27 @@ +#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, + const 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, + const 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/tests/test_reshape_and_cache.py b/tests/test_reshape_and_cache.py new file mode 100644 index 000000000..b24d32813 --- /dev/null +++ b/tests/test_reshape_and_cache.py @@ -0,0 +1,129 @@ +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..a714769a1 --- /dev/null +++ b/tests/test_rotary_embedding.py @@ -0,0 +1,167 @@ +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_tokens = 4 + 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) From 4f7756ffb9db507c0bbc6156ee483c5eba885fbd Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 17:06:00 +0800 Subject: [PATCH 08/22] feat(activation): replace swiglu with silu_and_mul --- src/base/silu_and_mul.h | 74 +++++++++++++++++++ src/base/swiglu.h | 68 ----------------- .../swiglu.h => silu_and_mul/silu_and_mul.h} | 38 +++++----- .../cuda/iluvatar/ops/silu_and_mul/kernel.h | 21 ++++++ src/native/cuda/iluvatar/ops/swiglu/kernel.h | 21 ------ .../cuda/metax/ops/silu_and_mul/kernel.h | 21 ++++++ src/native/cuda/metax/ops/swiglu/kernel.h | 21 ------ .../cuda/moore/ops/silu_and_mul/kernel.h | 26 +++++++ src/native/cuda/moore/ops/swiglu/kernel.h | 26 ------- .../cuda/nvidia/ops/silu_and_mul/kernel.h | 21 ++++++ src/native/cuda/nvidia/ops/swiglu/kernel.h | 21 ------ .../ops/{swiglu => silu_and_mul}/kernel.cuh | 54 +++++++------- .../ops/{swiglu => silu_and_mul}/kernel.h | 45 ++++------- tests/test_silu_and_mul.py | 72 ++++++++++++++++++ tests/test_swiglu.py | 72 ------------------ 15 files changed, 295 insertions(+), 306 deletions(-) create mode 100644 src/base/silu_and_mul.h delete mode 100644 src/base/swiglu.h rename src/native/cpu/ops/{swiglu/swiglu.h => silu_and_mul/silu_and_mul.h} (52%) create mode 100644 src/native/cuda/iluvatar/ops/silu_and_mul/kernel.h delete mode 100644 src/native/cuda/iluvatar/ops/swiglu/kernel.h create mode 100644 src/native/cuda/metax/ops/silu_and_mul/kernel.h delete mode 100644 src/native/cuda/metax/ops/swiglu/kernel.h create mode 100644 src/native/cuda/moore/ops/silu_and_mul/kernel.h delete mode 100644 src/native/cuda/moore/ops/swiglu/kernel.h create mode 100644 src/native/cuda/nvidia/ops/silu_and_mul/kernel.h delete mode 100644 src/native/cuda/nvidia/ops/swiglu/kernel.h rename src/native/cuda/ops/{swiglu => silu_and_mul}/kernel.cuh (56%) rename src/native/cuda/ops/{swiglu => silu_and_mul}/kernel.h (63%) create mode 100644 tests/test_silu_and_mul.py delete mode 100644 tests/test_swiglu.py 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 deleted file mode 100644 index 023b14a21..000000000 --- a/src/base/swiglu.h +++ /dev/null @@ -1,68 +0,0 @@ -#ifndef INFINI_OPS_BASE_SWIGLU_H_ -#define INFINI_OPS_BASE_SWIGLU_H_ - -#include - -#include "operator.h" - -namespace infini::ops { - -class Swiglu : public Operator { - public: - Swiglu(const Tensor input, const Tensor gate, Tensor out) - : ndim_{out.ndim()}, - output_size_{out.numel()}, - input_type_{input.dtype()}, - gate_type_{gate.dtype()}, - out_type_{out.dtype()}, - input_shape_{input.shape()}, - gate_shape_{gate.shape()}, - out_shape_{out.shape()}, - input_strides_{input.strides()}, - gate_strides_{gate.strides()}, - out_strides_{out.strides()}, - is_input_contiguous_{input.IsContiguous()}, - is_gate_contiguous_{gate.IsContiguous()}, - is_out_contiguous_{out.IsContiguous()} { - assert( - input_type_ == gate_type_ && gate_type_ == out_type_ && - "operator `Swiglu` requires all input and output tensors to have the " - "same dtype"); - } - - virtual void operator()(const Tensor input, const Tensor gate, - Tensor out) const = 0; - - protected: - Tensor::Size ndim_{0}; - - Tensor::Size output_size_{0}; - - const DataType input_type_; - - const DataType gate_type_; - - const DataType out_type_; - - Tensor::Shape input_shape_; - - Tensor::Shape gate_shape_; - - Tensor::Shape out_shape_; - - Tensor::Strides input_strides_; - - Tensor::Strides gate_strides_; - - Tensor::Strides out_strides_; - - bool is_input_contiguous_{false}; - - bool is_gate_contiguous_{false}; - - bool is_out_contiguous_{false}; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cpu/ops/swiglu/swiglu.h b/src/native/cpu/ops/silu_and_mul/silu_and_mul.h similarity index 52% rename from src/native/cpu/ops/swiglu/swiglu.h rename to src/native/cpu/ops/silu_and_mul/silu_and_mul.h index 37ae46767..2da25e506 100644 --- a/src/native/cpu/ops/swiglu/swiglu.h +++ b/src/native/cpu/ops/silu_and_mul/silu_and_mul.h @@ -1,40 +1,38 @@ -#ifndef INFINI_OPS_CPU_SWIGLU_SWIGLU_H_ -#define INFINI_OPS_CPU_SWIGLU_SWIGLU_H_ +#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/swiglu.h" +#include "base/silu_and_mul.h" #include "common/generic_utils.h" #include "native/cpu/caster_.h" namespace infini::ops { template <> -class Operator : public Swiglu, - Caster { +class Operator : public SiluAndMul, + Caster { public: - using Swiglu::Swiglu; + using SiluAndMul::SiluAndMul; - void operator()(const Tensor input, const Tensor gate, - Tensor out) const override { + void operator()(const Tensor input, Tensor out) const override { DispatchFunc( out_type_, [&](auto tag) { using T = typename decltype(tag)::type; - Compute(input, gate, out); + Compute(input, out); }, - "Operator::operator()"); + "Operator::operator()"); } private: template - void Compute(const Tensor input, const Tensor gate, Tensor out) const { + void Compute(const Tensor input, Tensor out) const { using ComputeType = std::conditional_t || IsFP16, float, T>; const auto* input_ptr = static_cast(input.data()); - const auto* gate_ptr = static_cast(gate.data()); auto* out_ptr = static_cast(out.data()); auto get_idx = [&](Tensor::Size i, bool is_contig, const auto* shape, @@ -44,18 +42,22 @@ class Operator : public Swiglu, #pragma omp parallel for for (Tensor::Size i = 0; i < output_size_; ++i) { - auto input_idx = get_idx(i, is_input_contiguous_, input_shape_.data(), - input_strides_.data()); - auto gate_idx = get_idx(i, is_gate_contiguous_, gate_shape_.data(), - gate_strides_.data()); + 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(gate_ptr[gate_idx]); + 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(Cast(input_ptr[input_idx]) * swish_gate); + Cast(swish_gate * Cast(input_ptr[up_idx])); } } }; 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/iluvatar/ops/swiglu/kernel.h b/src/native/cuda/iluvatar/ops/swiglu/kernel.h deleted file mode 100644 index b03ef19c9..000000000 --- a/src/native/cuda/iluvatar/ops/swiglu/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_ILUVATAR_SWIGLU_KERNEL_H_ -#define INFINI_OPS_ILUVATAR_SWIGLU_KERNEL_H_ - -#include - -#include "native/cuda/iluvatar/caster.cuh" -#include "native/cuda/iluvatar/runtime_.h" -#include "native/cuda/ops/swiglu/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaSwiglu> { - public: - using CudaSwiglu>::CudaSwiglu; -}; - -} // 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/metax/ops/swiglu/kernel.h b/src/native/cuda/metax/ops/swiglu/kernel.h deleted file mode 100644 index 0bcfd3198..000000000 --- a/src/native/cuda/metax/ops/swiglu/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_METAX_SWIGLU_KERNEL_H_ -#define INFINI_OPS_METAX_SWIGLU_KERNEL_H_ - -#include - -#include "native/cuda/metax/caster.cuh" -#include "native/cuda/metax/runtime_.h" -#include "native/cuda/ops/swiglu/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaSwiglu> { - public: - using CudaSwiglu>::CudaSwiglu; -}; - -} // 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/moore/ops/swiglu/kernel.h b/src/native/cuda/moore/ops/swiglu/kernel.h deleted file mode 100644 index ac54bff50..000000000 --- a/src/native/cuda/moore/ops/swiglu/kernel.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef INFINI_OPS_MOORE_SWIGLU_KERNEL_H_ -#define INFINI_OPS_MOORE_SWIGLU_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/swiglu/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaSwiglu> { - public: - using CudaSwiglu>::CudaSwiglu; -}; - -} // 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/nvidia/ops/swiglu/kernel.h b/src/native/cuda/nvidia/ops/swiglu/kernel.h deleted file mode 100644 index fd58f29ef..000000000 --- a/src/native/cuda/nvidia/ops/swiglu/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_SWIGLU_KERNEL_H_ -#define INFINI_OPS_NVIDIA_SWIGLU_KERNEL_H_ - -#include - -#include "native/cuda/nvidia/caster.cuh" -#include "native/cuda/nvidia/runtime_.h" -#include "native/cuda/ops/swiglu/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaSwiglu> { - public: - using CudaSwiglu>::CudaSwiglu; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/ops/swiglu/kernel.cuh b/src/native/cuda/ops/silu_and_mul/kernel.cuh similarity index 56% rename from src/native/cuda/ops/swiglu/kernel.cuh rename to src/native/cuda/ops/silu_and_mul/kernel.cuh index da4bfb053..c5ad1ce1f 100644 --- a/src/native/cuda/ops/swiglu/kernel.cuh +++ b/src/native/cuda/ops/silu_and_mul/kernel.cuh @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_CUDA_SWIGLU_KERNEL_CUH_ -#define INFINI_OPS_CUDA_SWIGLU_KERNEL_CUH_ +#ifndef INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_CUH_ +#define INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_CUH_ #include @@ -28,23 +28,20 @@ __device__ __forceinline__ T Sigmoid(const T& x) { } // namespace detail -// SwiGLU(x, gate) = Swish(x) * gate = (x * sigmoid(x)) * gate. template -__global__ void SwigluKernel(T* __restrict__ out, const T* __restrict__ a, - const T* __restrict__ b, - const size_t* __restrict__ out_shape, - const size_t* __restrict__ input_shape, - const size_t* __restrict__ gate_shape, - const ptrdiff_t* __restrict__ out_strides, - const ptrdiff_t* __restrict__ input_strides, - const ptrdiff_t* __restrict__ gate_strides, - size_t output_size, size_t ndim, - bool out_contiguous, bool input_contiguous, - bool gate_contiguous) { +__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, input_idx, gate_idx; + size_t out_idx; if (out_contiguous) { out_idx = idx; @@ -52,20 +49,21 @@ __global__ void SwigluKernel(T* __restrict__ out, const T* __restrict__ a, out_idx = IndexToOffset(idx, ndim, out_shape, out_strides); } - if (input_contiguous) { - input_idx = idx; - } else { - input_idx = IndexToOffset(idx, ndim, input_shape, input_strides); - } - - if (gate_contiguous) { - gate_idx = idx; - } else { - gate_idx = IndexToOffset(idx, ndim, gate_shape, gate_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 up = a[input_idx]; - T gate = b[gate_idx]; + T gate = input[gate_idx]; + T up = input[up_idx]; if constexpr (IsFP16 || IsBFloat16) { float gatef = Caster::template Cast(gate); diff --git a/src/native/cuda/ops/swiglu/kernel.h b/src/native/cuda/ops/silu_and_mul/kernel.h similarity index 63% rename from src/native/cuda/ops/swiglu/kernel.h rename to src/native/cuda/ops/silu_and_mul/kernel.h index 486617f43..0e9b5e620 100644 --- a/src/native/cuda/ops/swiglu/kernel.h +++ b/src/native/cuda/ops/silu_and_mul/kernel.h @@ -1,27 +1,26 @@ -#ifndef INFINI_OPS_CUDA_SWIGLU_KERNEL_H_ -#define INFINI_OPS_CUDA_SWIGLU_KERNEL_H_ +#ifndef INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_H_ +#define INFINI_OPS_CUDA_SILU_AND_MUL_KERNEL_H_ #include #include #include #include -#include "base/swiglu.h" +#include "base/silu_and_mul.h" #include "common/generic_utils.h" -#include "native/cuda/ops/swiglu/kernel.cuh" +#include "native/cuda/ops/silu_and_mul/kernel.cuh" #include "native/cuda/runtime_utils.h" namespace infini::ops { template -class CudaSwiglu : public Swiglu { +class CudaSiluAndMul : public SiluAndMul { public: - CudaSwiglu(const Tensor input, const Tensor gate, Tensor out) - : Swiglu{input, gate, out} { + 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 = 3 * (shape_size + strides_size); + const size_t metadata_size = 2 * (shape_size + strides_size); std::vector metadata(metadata_size); Backend::Malloc((void**)&d_metadata_, metadata_size); @@ -31,10 +30,6 @@ class CudaSwiglu : public Swiglu { std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size); offset += shape_size; - d_gate_shape_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, gate_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; @@ -43,10 +38,6 @@ class CudaSwiglu : public Swiglu { std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size); offset += strides_size; - d_gate_strides_ = reinterpret_cast(d_metadata_ + offset); - std::memcpy(metadata.data() + offset, gate_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); @@ -54,10 +45,9 @@ class CudaSwiglu : public Swiglu { Backend::kMemcpyHostToDevice); } - ~CudaSwiglu() { Backend::Free(d_metadata_); } + ~CudaSiluAndMul() { Backend::Free(d_metadata_); } - void operator()(const Tensor input, const Tensor gate, - Tensor out) const override { + void operator()(const Tensor input, Tensor out) const override { int block_size = RuntimeUtils::GetOptimalBlockSize(); DispatchFunc( {static_cast(out_type_), block_size}, @@ -73,16 +63,13 @@ class CudaSwiglu : public Swiglu { T* d_out = reinterpret_cast(out.data()); const T* d_input = reinterpret_cast(input.data()); - const T* d_gate = reinterpret_cast(gate.data()); - - SwigluKernel + SiluAndMulKernel <<>>( - d_out, d_input, d_gate, d_out_shape_, d_input_shape_, - d_gate_shape_, d_out_strides_, d_input_strides_, - d_gate_strides_, output_size_, ndim_, is_out_contiguous_, - is_input_contiguous_, is_gate_contiguous_); + 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_); }, - "CudaSwiglu::operator()"); + "CudaSiluAndMul::operator()"); } private: @@ -90,14 +77,10 @@ class CudaSwiglu : public Swiglu { Tensor::Size* d_input_shape_{nullptr}; - Tensor::Size* d_gate_shape_{nullptr}; - Tensor::Size* d_out_shape_{nullptr}; Tensor::Stride* d_input_strides_{nullptr}; - Tensor::Stride* d_gate_strides_{nullptr}; - Tensor::Stride* d_out_strides_{nullptr}; }; 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_swiglu.py b/tests/test_swiglu.py deleted file mode 100644 index f159742ca..000000000 --- a/tests/test_swiglu.py +++ /dev/null @@ -1,72 +0,0 @@ -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( - "shape, input_strides, gate_strides, out_strides", - ( - ((13, 4), None, None, None), - ((13, 4), (10, 1), (10, 1), (10, 1)), - ((13, 4, 4), None, None, None), - ((13, 4, 4), (20, 4, 1), (20, 4, 1), (20, 4, 1)), - ((16, 5632), None, None, None), - ((16, 5632), (13312, 1), (13312, 1), (13312, 1)), - ((4, 4, 5632), None, None, None), - ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), (45056, 5632, 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_swiglu( - shape, - input_strides, - gate_strides, - out_strides, - implementation_index, - dtype, - device, - rtol, - atol, -): - input = rand_strided(shape, input_strides, dtype=dtype, device=device) - gate = rand_strided(shape, gate_strides, dtype=dtype, device=device) - out = empty_strided(shape, out_strides, dtype=dtype, device=device) - - return Payload( - lambda *args, **kwargs: _swiglu( - *args, **kwargs, implementation_index=implementation_index - ), - _torch_swiglu, - (input, gate, out), - {}, - rtol=rtol, - atol=atol, - ) - - -def _swiglu(input, gate, out, implementation_index=0): - infini.ops.swiglu( - input, - gate, - out, - implementation_index=implementation_index, - stream=get_stream(input.device), - ) - - return out - - -def _torch_swiglu(input, gate, out): - swish_x = gate * torch.sigmoid(gate) - - return torch.mul(input, swish_x, out=out) From 52fca201ebf173ab357559d54c77e228ec6c646f Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 17:07:01 +0800 Subject: [PATCH 09/22] feat(attention): add PyTorch scaled dot product attention --- src/base/flash_attention.h | 112 -------------- src/base/scaled_dot_product_attention.h | 134 +++++++++++++++++ src/pybind11_utils.h | 4 + .../scaled_dot_product_attention.cc | 67 +++++++++ .../scaled_dot_product_attention.h | 30 ++++ src/torch/tensor_.h | 12 +- tests/test_scaled_dot_product_attention.py | 139 ++++++++++++++++++ 7 files changed, 382 insertions(+), 116 deletions(-) delete mode 100644 src/base/flash_attention.h create mode 100644 src/base/scaled_dot_product_attention.h create mode 100644 src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc create mode 100644 src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.h create mode 100644 tests/test_scaled_dot_product_attention.py diff --git a/src/base/flash_attention.h b/src/base/flash_attention.h deleted file mode 100644 index e5952b514..000000000 --- a/src/base/flash_attention.h +++ /dev/null @@ -1,112 +0,0 @@ -#ifndef INFINI_OPS_BASE_FLASH_ATTENTION_H_ -#define INFINI_OPS_BASE_FLASH_ATTENTION_H_ - -#include -#include -#include - -#include "operator.h" - -namespace infini::ops { - -// Fused multi-head / grouped-query attention. -// -// Interface follows vLLM v1 `AttentionImpl.forward()`: -// `vllm.v1.attention.backends.abstract.AttentionImpl` -// -// 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. -class FlashAttention : public Operator { - public: - FlashAttention(const Tensor query, const Tensor key, const Tensor value, - std::optional cu_seqlens_q, - std::optional cu_seqlens_kv, - std::optional block_table, int64_t num_heads, - int64_t num_kv_heads, int64_t head_size, double scale, - bool causal, int64_t window_left, int64_t window_right, - int64_t block_size, Tensor output) - : num_tokens_{query.size(0)}, - num_heads_{num_heads}, - num_kv_heads_{num_kv_heads}, - head_size_{head_size}, - scale_{scale}, - causal_{causal}, - window_left_{window_left}, - window_right_{window_right}, - block_size_{block_size}, - dtype_{query.dtype()}, - query_shape_{query.shape()}, - key_shape_{key.shape()}, - value_shape_{value.shape()}, - output_shape_{output.shape()}, - query_strides_{query.strides()}, - key_strides_{key.strides()}, - value_strides_{value.strides()}, - output_strides_{output.strides()}, - has_cu_seqlens_q_{cu_seqlens_q.has_value()}, - has_cu_seqlens_kv_{cu_seqlens_kv.has_value()}, - has_block_table_{block_table.has_value()} { - assert(num_heads % num_kv_heads == 0 && - "`FlashAttention` requires num_heads divisible by num_kv_heads"); - assert(query.ndim() == 3 && - "`FlashAttention` requires query to be 3D [T, N, D]"); - } - - virtual void operator()(const Tensor query, const Tensor key, - const Tensor value, - std::optional cu_seqlens_q, - std::optional cu_seqlens_kv, - std::optional block_table, int64_t num_heads, - int64_t num_kv_heads, int64_t head_size, double scale, - bool causal, int64_t window_left, - int64_t window_right, int64_t block_size, - Tensor output) const = 0; - - protected: - Tensor::Size num_tokens_{0}; - - int64_t num_heads_{0}; - - int64_t num_kv_heads_{0}; - - int64_t head_size_{0}; - - double scale_{0.0}; - - bool causal_{false}; - - int64_t window_left_{-1}; - - int64_t window_right_{-1}; - - int64_t block_size_{0}; - - const DataType dtype_; - - Tensor::Shape query_shape_; - - Tensor::Shape key_shape_; - - Tensor::Shape value_shape_; - - Tensor::Shape output_shape_; - - Tensor::Strides query_strides_; - - Tensor::Strides key_strides_; - - Tensor::Strides value_strides_; - - Tensor::Strides output_strides_; - - bool has_cu_seqlens_q_{false}; - - bool has_cu_seqlens_kv_{false}; - - bool has_block_table_{false}; -}; - -} // namespace infini::ops - -#endif diff --git a/src/base/scaled_dot_product_attention.h b/src/base/scaled_dot_product_attention.h new file mode 100644 index 000000000..c6d24d4c7 --- /dev/null +++ b/src/base/scaled_dot_product_attention.h @@ -0,0 +1,134 @@ +#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"); + assert(out.shape() == ReturnShape(query, value) && + "`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; + + return TensorLike::Empty(ReturnShape(query, value), query.dtype(), + query.device()); + } + + protected: + template + static auto ReturnShape(const TensorLike& query, const TensorLike& value) { + auto shape = query.shape(); + shape.back() = value.size(-1); + + return shape; + } + + 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/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/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..169c1eacd --- /dev/null +++ b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc @@ -0,0 +1,67 @@ +#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_); + + std::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)); + } + + auto result = + at::scaled_dot_product_attention(at_query, at_key, at_value, at_attn_mask, + dropout_p, is_causal, scale, enable_gqa); + 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; +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 3a7396b0b..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 = 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 From 9a15759e046d10850f10461abda9d221177cf412 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 17:08:00 +0800 Subject: [PATCH 10/22] refactor(ops): internalize custom softmax and sampling kernels --- ...al_softmax.h => internal_causal_softmax.h} | 8 +- ...ed_softmax.h => internal_scaled_softmax.h} | 10 +- src/base/internal_top_k_top_p_sample.h | 79 ++++++++++++++++ src/base/top_k_top_p_sampler.h | 78 --------------- .../kernel.h | 13 +-- .../kernel.h | 52 +++++----- .../internal_causal_softmax.h} | 15 +-- .../internal_top_k_top_p_sample.h} | 41 ++++---- .../cuda/iluvatar/ops/causal_softmax/kernel.h | 21 ----- .../ops/internal_causal_softmax/kernel.h | 22 +++++ .../cuda/metax/ops/causal_softmax/kernel.h | 21 ----- .../ops/internal_causal_softmax/kernel.h | 22 +++++ .../kernel.h | 19 ++-- .../cuda/nvidia/ops/causal_softmax/kernel.h | 21 ----- .../ops/internal_causal_softmax/kernel.h | 22 +++++ .../kernel.cuh | 8 +- .../kernel.h | 12 +-- ...max.py => test_internal_causal_softmax.py} | 12 ++- ...max.py => test_internal_scaled_softmax.py} | 8 +- tests/test_internal_top_k_top_p_sample.py | 94 +++++++++++++++++++ tests/test_top_k_top_p_sampler.py | 78 --------------- 21 files changed, 346 insertions(+), 310 deletions(-) rename src/base/{causal_softmax.h => internal_causal_softmax.h} (87%) rename src/base/{scaled_softmax.h => internal_scaled_softmax.h} (86%) create mode 100644 src/base/internal_top_k_top_p_sample.h delete mode 100644 src/base/top_k_top_p_sampler.h rename src/native/ascend/ops/{scaled_softmax => internal_scaled_softmax}/kernel.h (90%) rename src/native/ascend/ops/{top_k_top_p_sampler => internal_top_k_top_p_sample}/kernel.h (85%) rename src/native/cpu/ops/{causal_softmax/causal_softmax.h => internal_causal_softmax/internal_causal_softmax.h} (83%) rename src/native/cpu/ops/{top_k_top_p_sampler/top_k_top_p_sampler.h => internal_top_k_top_p_sample/internal_top_k_top_p_sample.h} (77%) delete mode 100644 src/native/cuda/iluvatar/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h delete mode 100644 src/native/cuda/metax/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/metax/ops/internal_causal_softmax/kernel.h rename src/native/cuda/moore/ops/{causal_softmax => internal_causal_softmax}/kernel.h (50%) delete mode 100644 src/native/cuda/nvidia/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h rename src/native/cuda/ops/{causal_softmax => internal_causal_softmax}/kernel.cuh (94%) rename src/native/cuda/ops/{causal_softmax => internal_causal_softmax}/kernel.h (87%) rename tests/{test_causal_softmax.py => test_internal_causal_softmax.py} (82%) rename tests/{test_scaled_softmax.py => test_internal_scaled_softmax.py} (87%) create mode 100644 tests/test_internal_top_k_top_p_sample.py delete mode 100644 tests/test_top_k_top_p_sampler.py diff --git a/src/base/causal_softmax.h b/src/base/internal_causal_softmax.h similarity index 87% rename from src/base/causal_softmax.h rename to src/base/internal_causal_softmax.h index b8393d829..4b9ae1855 100644 --- a/src/base/causal_softmax.h +++ b/src/base/internal_causal_softmax.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_CAUSAL_SOFTMAX_H_ -#define INFINI_OPS_BASE_CAUSAL_SOFTMAX_H_ +#ifndef INFINI_OPS_BASE_INTERNAL_CAUSAL_SOFTMAX_H_ +#define INFINI_OPS_BASE_INTERNAL_CAUSAL_SOFTMAX_H_ #include #include @@ -7,7 +7,7 @@ #include "operator.h" #include "tensor.h" -namespace infini::ops { +namespace infini::ops::internal { class CausalSoftmax : public Operator { public: @@ -47,6 +47,6 @@ class CausalSoftmax : public Operator { Tensor::Strides out_strides_; }; -} // namespace infini::ops +} // namespace infini::ops::internal #endif diff --git a/src/base/scaled_softmax.h b/src/base/internal_scaled_softmax.h similarity index 86% rename from src/base/scaled_softmax.h rename to src/base/internal_scaled_softmax.h index 707bf1ffc..3d97bba9d 100644 --- a/src/base/scaled_softmax.h +++ b/src/base/internal_scaled_softmax.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_SCALED_SOFTMAX_H_ -#define INFINI_OPS_BASE_SCALED_SOFTMAX_H_ +#ifndef INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ +#define INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ #include #include @@ -9,7 +9,7 @@ #include "operator.h" #include "tensor.h" -namespace infini::ops { +namespace infini::ops::internal { class ScaledSoftmax : public Operator { public: @@ -50,6 +50,6 @@ class ScaledSoftmax : public Operator { Tensor::Strides out_strides_; }; -} // namespace infini::ops +} // namespace infini::ops::internal -#endif // INFINI_OPS_BASE_SCALED_SOFTMAX_H_ +#endif // INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ diff --git a/src/base/internal_top_k_top_p_sample.h b/src/base/internal_top_k_top_p_sample.h new file mode 100644 index 000000000..6e74b46fa --- /dev/null +++ b/src/base/internal_top_k_top_p_sample.h @@ -0,0 +1,79 @@ +#ifndef INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#define INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ + +#include +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops::internal { + +class TopKTopPSample : public Operator { + public: + TopKTopPSample(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 && + "`TopKTopPSample` requires 2D `[batch_size, vocab_size]` logits"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && + "`TopKTopPSample` requires floating-point logits"); + assert(out.ndim() == 1 && + "`TopKTopPSample` requires 1D `[batch_size]` output"); + assert(out.size(0) == batch_size_ && + "`TopKTopPSample` requires output batch size to match logits"); + assert(out.dtype() == DataType::kInt32 && + "`TopKTopPSample` 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 && + "`TopKTopPSample` requires `k` to be 1D when provided"); + assert((k->size(0) == 1 || k->size(0) == batch_size_) && + "`TopKTopPSample` requires `k` shape [1] or [batch_size]"); + assert((k->dtype() == DataType::kInt32 || k->dtype() == DataType::kInt64) && + "`TopKTopPSample` requires int32 or int64 `k`"); + } + + void ValidateP(std::optional p) const { + if (!p.has_value()) return; + + assert(p->ndim() == 1 && + "`TopKTopPSample` requires `p` to be 1D when provided"); + assert((p->size(0) == 1 || p->size(0) == batch_size_) && + "`TopKTopPSample` requires `p` shape [1] or [batch_size]"); + assert((p->dtype() == DataType::kFloat16 || + p->dtype() == DataType::kBFloat16 || + p->dtype() == DataType::kFloat32 || + p->dtype() == DataType::kFloat64) && + "`TopKTopPSample` requires floating-point `p`"); + } + + Tensor::Size batch_size_{0}; + + Tensor::Size vocab_size_{0}; + + DataType dtype_; +}; + +} // namespace infini::ops::internal + +#endif // INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ diff --git a/src/base/top_k_top_p_sampler.h b/src/base/top_k_top_p_sampler.h deleted file mode 100644 index 085a2b937..000000000 --- a/src/base/top_k_top_p_sampler.h +++ /dev/null @@ -1,78 +0,0 @@ -#ifndef INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLER_H_ -#define INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLER_H_ - -#include -#include - -#include "data_type.h" -#include "operator.h" -#include "tensor.h" - -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. -// The optional `k` and `p` tensors may be shaped as `[1]` or `[batch_size]`. -class TopKTopPSampler : public Operator { - public: - TopKTopPSampler(const Tensor logits, std::optional k, - std::optional p, Tensor out) - : batch_size_{logits.size(0)}, - vocab_size_{logits.size(1)}, - dtype_{logits.dtype()} { - assert(logits.ndim() == 2 && - "`TopKTopPSampler` requires 2D `[batch_size, vocab_size]` logits"); - assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || - dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && - "`TopKTopPSampler` requires floating-point logits"); - assert(out.ndim() == 1 && - "`TopKTopPSampler` requires 1D `[batch_size]` output"); - assert(out.size(0) == batch_size_ && - "`TopKTopPSampler` requires output batch size to match logits"); - assert(out.dtype() == DataType::kInt32 && - "`TopKTopPSampler` requires int32 output"); - - ValidateK(k); - ValidateP(p); - } - - virtual void operator()(const Tensor logits, std::optional k, - std::optional p, Tensor out) const = 0; - - protected: - void ValidateK(std::optional k) const { - if (!k.has_value()) return; - - assert(k->ndim() == 1 && - "`TopKTopPSampler` requires `k` to be 1D when provided"); - assert((k->size(0) == 1 || k->size(0) == batch_size_) && - "`TopKTopPSampler` requires `k` shape [1] or [batch_size]"); - assert((k->dtype() == DataType::kInt32 || k->dtype() == DataType::kInt64) && - "`TopKTopPSampler` requires int32 or int64 `k`"); - } - - void ValidateP(std::optional p) const { - if (!p.has_value()) return; - - assert(p->ndim() == 1 && - "`TopKTopPSampler` requires `p` to be 1D when provided"); - assert((p->size(0) == 1 || p->size(0) == batch_size_) && - "`TopKTopPSampler` requires `p` shape [1] or [batch_size]"); - assert((p->dtype() == DataType::kFloat16 || - p->dtype() == DataType::kBFloat16 || - p->dtype() == DataType::kFloat32 || - p->dtype() == DataType::kFloat64) && - "`TopKTopPSampler` 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_SAMPLER_H_ diff --git a/src/native/ascend/ops/scaled_softmax/kernel.h b/src/native/ascend/ops/internal_scaled_softmax/kernel.h similarity index 90% rename from src/native/ascend/ops/scaled_softmax/kernel.h rename to src/native/ascend/ops/internal_scaled_softmax/kernel.h index c6c0df6f7..291fdb204 100644 --- a/src/native/ascend/ops/scaled_softmax/kernel.h +++ b/src/native/ascend/ops/internal_scaled_softmax/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_ASCEND_SCALED_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_ASCEND_SCALED_SOFTMAX_KERNEL_H_ +#ifndef INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ #include @@ -7,7 +7,7 @@ #include "aclnn/aclnn_base.h" #include "aclnn_mul.h" #include "aclnn_softmax.h" -#include "base/scaled_softmax.h" +#include "base/internal_scaled_softmax.h" #include "data_type.h" #include "native/ascend/common.h" #include "native/ascend/workspace_pool_.h" @@ -16,10 +16,11 @@ namespace infini::ops { template <> -class Operator : public ScaledSoftmax { +class Operator + : public internal::ScaledSoftmax { public: Operator(const Tensor input, double scale, Tensor out) - : ScaledSoftmax(input, scale, out), + : internal::ScaledSoftmax(input, scale, out), in_cache_(input), out_cache_(out), temp_cache_(input), @@ -121,4 +122,4 @@ class Operator : public ScaledSoftmax { } // namespace infini::ops -#endif // INFINI_OPS_ASCEND_SCALED_SOFTMAX_KERNEL_H_ +#endif // INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ diff --git a/src/native/ascend/ops/top_k_top_p_sampler/kernel.h b/src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h similarity index 85% rename from src/native/ascend/ops/top_k_top_p_sampler/kernel.h rename to src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h index 6da146118..0a48701ad 100644 --- a/src/native/ascend/ops/top_k_top_p_sampler/kernel.h +++ b/src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLER_KERNEL_H_ -#define INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLER_KERNEL_H_ +#ifndef INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ +#define INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ #include #include @@ -13,7 +13,7 @@ #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_sampler.h" +#include "base/internal_top_k_top_p_sample.h" #include "data_type.h" #include "native/ascend/common.h" #include "native/ascend/workspace_pool_.h" @@ -23,19 +23,19 @@ namespace infini::ops { template <> -class Operator - : public TopKTopPSampler { +class Operator + : public internal::TopKTopPSample { public: Operator(const Tensor logits, std::optional k, - std::optional p, Tensor out) - : TopKTopPSampler(logits, k, p, out) { + std::optional p, uint64_t seed, uint64_t offset, Tensor out) + : internal::TopKTopPSample(logits, k, p, seed, offset, out) { assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16) && - "`TopKTopPSampler` Ascend ACLNN path requires float16 or bfloat16 " + "`TopKTopPSample` Ascend ACLNN path requires float16 or bfloat16 " "logits"); assert(logits.IsContiguous() && - "`TopKTopPSampler` Ascend ACLNN path requires contiguous logits"); + "`TopKTopPSample` Ascend ACLNN path requires contiguous logits"); assert(out.IsContiguous() && - "`TopKTopPSampler` Ascend ACLNN path requires contiguous output"); + "`TopKTopPSample` Ascend ACLNN path requires contiguous output"); ValidateHostTensor(k); ValidateHostTensor(p); @@ -68,11 +68,12 @@ class Operator } void operator()(const Tensor logits, std::optional k, - std::optional p, Tensor out) const override { + std::optional p, uint64_t seed, uint64_t offset, + Tensor out) const override { assert(logits.IsContiguous() && - "`TopKTopPSampler` Ascend ACLNN path requires contiguous logits"); + "`TopKTopPSample` Ascend ACLNN path requires contiguous logits"); assert(out.IsContiguous() && - "`TopKTopPSampler` Ascend ACLNN path requires contiguous output"); + "`TopKTopPSample` Ascend ACLNN path requires contiguous output"); auto stream = static_cast(stream_); auto top_k_bytes = batch_size_ * kDataTypeToSize.at(DataType::kInt32); @@ -84,7 +85,7 @@ class Operator auto selected_logits_bytes = batch_size_ * vocab_size_ * kDataTypeToSize.at(DataType::kFloat32); - FillParams(k, p); + FillParams(k, p, seed, offset); auto& top_k_arena = ascend::GetWorkspacePool().Ensure( stream, top_k_bytes, "top_k_top_p_sample_top_k"); @@ -95,15 +96,15 @@ class Operator 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 && - "`TopKTopPSampler`: copying `top_k` to Ascend failed"); + "`TopKTopPSample`: 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 && - "`TopKTopPSampler`: copying `top_p` to Ascend failed"); + "`TopKTopPSample`: 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 && - "`TopKTopPSampler`: copying `q` to Ascend failed"); + "`TopKTopPSample`: copying `q` to Ascend failed"); auto& selected_idx_arena = ascend::GetWorkspacePool().Ensure( stream, selected_idx_bytes, "top_k_top_p_sample_idx"); @@ -152,10 +153,10 @@ class Operator if (!tensor.has_value()) return; assert(tensor->device().type() == Device::Type::kCpu && - "`TopKTopPSampler` Ascend path currently requires host-side " + "`TopKTopPSample` Ascend path currently requires host-side " "`k`/`p` tensors"); assert(tensor->IsContiguous() && - "`TopKTopPSampler` Ascend path requires contiguous `k`/`p` " + "`TopKTopPSample` Ascend path requires contiguous `k`/`p` " "tensors"); } @@ -180,11 +181,14 @@ class Operator assert(ret == ACL_SUCCESS && "`aclnnCast` failed"); } - void FillParams(std::optional k, std::optional p) const { + 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)); @@ -203,7 +207,7 @@ class Operator // CANN samples by taking argmax of softmax probabilities divided by // exponential noise. for (auto& value : q_host_) { - value = dist(rng_); + value = dist(rng); } } @@ -241,7 +245,7 @@ class Operator value = static_cast(p->data())[offset]; break; default: - assert(false && "`TopKTopPSampler` has unsupported `p` dtype"); + assert(false && "`TopKTopPSample` has unsupported `p` dtype"); } if (value <= 0.0 || value > 1.0) return 1.0; @@ -268,8 +272,6 @@ class Operator mutable std::vector q_host_; - mutable std::mt19937 rng_{std::random_device{}()}; - mutable aclOpExecutor* sample_exec_ = nullptr; mutable uint64_t sample_ws_size_ = 0; @@ -281,4 +283,4 @@ class Operator } // namespace infini::ops -#endif // INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLER_KERNEL_H_ +#endif // INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ diff --git a/src/native/cpu/ops/causal_softmax/causal_softmax.h b/src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h similarity index 83% rename from src/native/cpu/ops/causal_softmax/causal_softmax.h rename to src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h index 8b00c6535..4e7f62524 100644 --- a/src/native/cpu/ops/causal_softmax/causal_softmax.h +++ b/src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h @@ -1,9 +1,9 @@ -#ifndef INFINI_OPS_CPU_CAUSAL_SOFTMAX_H_ -#define INFINI_OPS_CPU_CAUSAL_SOFTMAX_H_ +#ifndef INFINI_OPS_CPU_INTERNAL_CAUSAL_SOFTMAX_H_ +#define INFINI_OPS_CPU_INTERNAL_CAUSAL_SOFTMAX_H_ #include -#include "base/causal_softmax.h" +#include "base/internal_causal_softmax.h" #include "common/generic_utils.h" #include "data_type.h" #include "native/cpu/caster_.h" @@ -12,10 +12,11 @@ namespace infini::ops { template <> -class Operator : public CausalSoftmax, - Caster { +class Operator + : public internal::CausalSoftmax, Caster { public: - Operator(const Tensor input, Tensor out) : CausalSoftmax{input, out} {} + Operator(const Tensor input, Tensor out) + : internal::CausalSoftmax{input, out} {} void operator()(const Tensor input, Tensor out) const override { DispatchFunc( @@ -24,7 +25,7 @@ class Operator : public CausalSoftmax, using T = typename decltype(tag)::type; Compute(input, out); }, - "`Operator::operator()`"); + "`Operator::operator()`"); } private: diff --git a/src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h b/src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h similarity index 77% rename from src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h rename to src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h index b0246fad2..f5d4fcb0b 100644 --- a/src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h +++ b/src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ -#define INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ +#ifndef INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#define INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ #include #include @@ -11,7 +11,7 @@ #include #include -#include "base/top_k_top_p_sampler.h" +#include "base/internal_top_k_top_p_sample.h" #include "data_type.h" #include "native/cpu/caster_.h" #include "operator.h" @@ -20,42 +20,45 @@ namespace infini::ops { template <> -class Operator - : public TopKTopPSampler, Caster { +class Operator + : public internal::TopKTopPSample, Caster { public: Operator(const Tensor logits, std::optional k, - std::optional p, Tensor out) - : TopKTopPSampler(logits, k, p, out) {} + std::optional p, uint64_t seed, uint64_t offset, Tensor out) + : internal::TopKTopPSample(logits, k, p, seed, offset, out) {} void operator()(const Tensor logits, std::optional k, - std::optional p, Tensor out) const override { + 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, out); + Compute(logits, k, p, seed, offset, out); }, - "`Operator::operator()`"); + "`Operator::operator()`"); } private: template void Compute(const Tensor logits, std::optional k, - std::optional p, Tensor out) const { + 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); + 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) const { + 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) { @@ -84,7 +87,9 @@ class Operator std::discrete_distribution dist(weights.begin(), weights.end()); - return static_cast(indices[dist(rng_)]); + std::mt19937_64 rng(seed); + rng.discard(offset); + return static_cast(indices[dist(rng)]); } template @@ -147,14 +152,12 @@ class Operator case DataType::kFloat64: return static_cast(p->data())[offset]; default: - assert(false && "`TopKTopPSampler` has unsupported `p` dtype."); + assert(false && "`TopKTopPSample` has unsupported `p` dtype."); return 1.0; } } - - mutable std::mt19937 rng_{std::random_device{}()}; }; } // namespace infini::ops -#endif // INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ +#endif // INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ diff --git a/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h b/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h deleted file mode 100644 index 8c35c87ce..000000000 --- a/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/iluvatar/caster.cuh" -#include "native/cuda/iluvatar/runtime_.h" -#include "native/cuda/ops/causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaCausalSoftmax> { - public: - using CudaCausalSoftmax>::CudaCausalSoftmax; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h b/src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h new file mode 100644 index 000000000..0f6bf9739 --- /dev/null +++ b/src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_ILUVATAR_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/internal_causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public internal::CudaCausalSoftmax> { + public: + using internal::CudaCausalSoftmax< + Runtime>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/causal_softmax/kernel.h b/src/native/cuda/metax/ops/causal_softmax/kernel.h deleted file mode 100644 index 426465984..000000000 --- a/src/native/cuda/metax/ops/causal_softmax/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_METAX_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_METAX_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/metax/caster.cuh" -#include "native/cuda/metax/runtime_.h" -#include "native/cuda/ops/causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaCausalSoftmax> { - public: - using CudaCausalSoftmax>::CudaCausalSoftmax; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/metax/ops/internal_causal_softmax/kernel.h b/src/native/cuda/metax/ops/internal_causal_softmax/kernel.h new file mode 100644 index 000000000..167653961 --- /dev/null +++ b/src/native/cuda/metax/ops/internal_causal_softmax/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_METAX_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_METAX_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/internal_causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public internal::CudaCausalSoftmax> { + public: + using internal::CudaCausalSoftmax< + Runtime>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/causal_softmax/kernel.h b/src/native/cuda/moore/ops/internal_causal_softmax/kernel.h similarity index 50% rename from src/native/cuda/moore/ops/causal_softmax/kernel.h rename to src/native/cuda/moore/ops/internal_causal_softmax/kernel.h index e13118763..d608784a7 100644 --- a/src/native/cuda/moore/ops/causal_softmax/kernel.h +++ b/src/native/cuda/moore/ops/internal_causal_softmax/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_MOORE_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_MOORE_CAUSAL_SOFTMAX_KERNEL_H_ +#ifndef INFINI_OPS_MOORE_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_MOORE_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ // clang-format off #include @@ -11,20 +11,25 @@ #include "native/cuda/moore/caster.cuh" #include "native/cuda/moore/runtime_.h" -#include "native/cuda/ops/causal_softmax/kernel.h" +#include "native/cuda/ops/internal_causal_softmax/kernel.h" -namespace infini::ops { +namespace infini::ops::internal { struct MooreCausalSoftmaxBackend : Runtime { // Moore's causal softmax kernel should not dispatch block sizes above 1024. static constexpr int max_block_size = 1024; }; +} // namespace infini::ops::internal + +namespace infini::ops { + template <> -class Operator - : public CudaCausalSoftmax { +class Operator + : public internal::CudaCausalSoftmax { public: - using CudaCausalSoftmax::CudaCausalSoftmax; + using internal::CudaCausalSoftmax< + internal::MooreCausalSoftmaxBackend>::CudaCausalSoftmax; }; } // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/causal_softmax/kernel.h b/src/native/cuda/nvidia/ops/causal_softmax/kernel.h deleted file mode 100644 index 5a7e18a94..000000000 --- a/src/native/cuda/nvidia/ops/causal_softmax/kernel.h +++ /dev/null @@ -1,21 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/nvidia/caster.cuh" -#include "native/cuda/nvidia/runtime_.h" -#include "native/cuda/ops/causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public CudaCausalSoftmax> { - public: - using CudaCausalSoftmax>::CudaCausalSoftmax; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h b/src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h new file mode 100644 index 000000000..d7bf94dd4 --- /dev/null +++ b/src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_NVIDIA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_NVIDIA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/internal_causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public internal::CudaCausalSoftmax> { + public: + using internal::CudaCausalSoftmax< + Runtime>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/causal_softmax/kernel.cuh b/src/native/cuda/ops/internal_causal_softmax/kernel.cuh similarity index 94% rename from src/native/cuda/ops/causal_softmax/kernel.cuh rename to src/native/cuda/ops/internal_causal_softmax/kernel.cuh index 31f685e84..4079a508f 100644 --- a/src/native/cuda/ops/causal_softmax/kernel.cuh +++ b/src/native/cuda/ops/internal_causal_softmax/kernel.cuh @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_CUH_ -#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_CUH_ +#ifndef INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_CUH_ +#define INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_CUH_ #include #include @@ -8,7 +8,7 @@ #include "native/cuda/caster.cuh" #include "native/cuda/kernel_commons.cuh" -namespace infini::ops { +namespace infini::ops::internal { namespace { @@ -101,6 +101,6 @@ __global__ void CausalSoftmaxKernel( } } -} // namespace infini::ops +} // namespace infini::ops::internal #endif diff --git a/src/native/cuda/ops/causal_softmax/kernel.h b/src/native/cuda/ops/internal_causal_softmax/kernel.h similarity index 87% rename from src/native/cuda/ops/causal_softmax/kernel.h rename to src/native/cuda/ops/internal_causal_softmax/kernel.h index c77d97387..6af110771 100644 --- a/src/native/cuda/ops/causal_softmax/kernel.h +++ b/src/native/cuda/ops/internal_causal_softmax/kernel.h @@ -1,18 +1,18 @@ -#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_H_ +#ifndef INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ #include #include #include -#include "base/causal_softmax.h" +#include "base/internal_causal_softmax.h" #include "data_type.h" #include "dispatcher.h" #include "native/cuda/kernel_commons.cuh" -#include "native/cuda/ops/causal_softmax/kernel.cuh" +#include "native/cuda/ops/internal_causal_softmax/kernel.cuh" #include "native/cuda/runtime_utils.h" -namespace infini::ops { +namespace infini::ops::internal { template class CudaCausalSoftmax : public CausalSoftmax { @@ -58,6 +58,6 @@ class CudaCausalSoftmax : public CausalSoftmax { } }; -} // namespace infini::ops +} // namespace infini::ops::internal #endif diff --git a/tests/test_causal_softmax.py b/tests/test_internal_causal_softmax.py similarity index 82% rename from tests/test_causal_softmax.py rename to tests/test_internal_causal_softmax.py index ef48601f7..e3b55a8e1 100644 --- a/tests/test_causal_softmax.py +++ b/tests/test_internal_causal_softmax.py @@ -25,12 +25,14 @@ (torch.bfloat16, 1e-2, 1e-2), ), ) -def test_causal_softmax(shape, input_strides, out_strides, dtype, device, rtol, atol): +def test_internal_causal_softmax( + 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, + _internal_causal_softmax, _torch_causal_softmax, (input_tensor, out), {}, @@ -39,8 +41,10 @@ def test_causal_softmax(shape, input_strides, out_strides, dtype, device, rtol, ) -def _causal_softmax(input, out): - infini.ops.causal_softmax(input, out, stream=get_stream(input.device)) +def _internal_causal_softmax(input, out): + infini.ops.internal_causal_softmax( + input, out, stream=get_stream(input.device) + ) return out diff --git a/tests/test_scaled_softmax.py b/tests/test_internal_scaled_softmax.py similarity index 87% rename from tests/test_scaled_softmax.py rename to tests/test_internal_scaled_softmax.py index ec860b34b..5f88d0b93 100644 --- a/tests/test_scaled_softmax.py +++ b/tests/test_internal_scaled_softmax.py @@ -23,7 +23,7 @@ (torch.bfloat16, 1e-2, 1e-2), ), ) -def test_scaled_softmax( +def test_internal_scaled_softmax( shape, scale, dtype, @@ -36,7 +36,7 @@ def test_scaled_softmax( out = empty_strided(shape, None, dtype=dtype, device=device) return Payload( - _scaled_softmax, + _internal_scaled_softmax, _torch_scaled_softmax, (input_tensor, out), {"scale": scale, "implementation_index": implementation_index}, @@ -45,8 +45,8 @@ def test_scaled_softmax( ) -def _scaled_softmax(input_tensor, out, *, scale, implementation_index): - infini.ops.scaled_softmax( +def _internal_scaled_softmax(input_tensor, out, *, scale, implementation_index): + infini.ops.internal_scaled_softmax( input_tensor, scale, out, diff --git a/tests/test_internal_top_k_top_p_sample.py b/tests/test_internal_top_k_top_p_sample.py new file mode 100644 index 000000000..353cea7d8 --- /dev/null +++ b/tests/test_internal_top_k_top_p_sample.py @@ -0,0 +1,94 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_internal_top_k_top_p_sample_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) + + _internal_top_k_top_p_sample( + logits, None, None, 1234, 9, first, implementation_index + ) + _internal_top_k_top_p_sample( + logits, None, None, 1234, 9, second, implementation_index + ) + _internal_top_k_top_p_sample( + 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_internal_top_k_top_p_sample_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) + + _internal_top_k_top_p_sample( + 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 _internal_top_k_top_p_sample( + logits, + k, + p, + seed, + offset, + out, + implementation_index, +): + infini.ops.internal_top_k_top_p_sample( + logits, + k, + p, + seed, + offset, + out, + stream=get_stream(logits.device), + implementation_index=implementation_index, + ) diff --git a/tests/test_top_k_top_p_sampler.py b/tests/test_top_k_top_p_sampler.py deleted file mode 100644 index 8d898e0e6..000000000 --- a/tests/test_top_k_top_p_sampler.py +++ /dev/null @@ -1,78 +0,0 @@ -import infini.ops -import pytest -import torch - -from tests.utils import Payload, get_stream - - -@pytest.mark.auto_act_and_assert -@pytest.mark.parametrize("shape", ((1, 8), (3, 16))) -@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_top_k_top_p_sampler( - shape, - dtype, - device, - implementation_index, -): - batch_size, vocab_size = shape - logits = torch.full(shape, -10.0, dtype=dtype, device=device) - - for i in range(batch_size): - logits[i, i % vocab_size] = 10.0 - - k = torch.ones((batch_size,), dtype=torch.int64, device="cpu") - p = torch.ones((batch_size,), dtype=torch.float32, device="cpu") - out = torch.empty((batch_size,), dtype=torch.int32, device=device) - - return Payload( - _top_k_top_p_sampler, - _torch_argmax, - (logits, k, p, out), - {"implementation_index": implementation_index}, - ) - - -@pytest.mark.auto_act_and_assert -@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_top_k_top_p_sampler_optional_p( - dtype, - device, - implementation_index, -): - shape = (3, 16) - batch_size, vocab_size = shape - logits = torch.full(shape, -10.0, dtype=dtype, device=device) - - for i in range(batch_size): - logits[i, (i + 1) % vocab_size] = 10.0 - - k = torch.ones((1,), dtype=torch.int64, device="cpu") - out = torch.empty((batch_size,), dtype=torch.int32, device=device) - - return Payload( - _top_k_top_p_sampler, - _torch_argmax, - (logits, k, None, out), - {"implementation_index": implementation_index}, - ) - - -def _top_k_top_p_sampler(logits, k, p, out, *, implementation_index): - infini.ops.top_k_top_p_sampler( - logits, - k, - p, - out, - stream=get_stream(logits.device), - implementation_index=implementation_index, - ) - - return out - - -def _torch_argmax(logits, k, p, out, *, implementation_index): - del k, p, implementation_index - - out.copy_(torch.argmax(logits, dim=-1).to(torch.int32)) - - return out From d95494b35c50249ef08313f7725c01778a462259 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 17:09:11 +0800 Subject: [PATCH 11/22] test(smoke): update renamed operator coverage --- src/CMakeLists.txt | 3 ++- tests/conftest.py | 15 +++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..78e3affad 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 silu_and_mul + internal_causal_softmax abs clamp exp) set(_infini_ops_smoke_torch_ops abs clamp exp) if(INFINI_OPS_SMOKE_BUILD) diff --git a/tests/conftest.py b/tests/conftest.py index 230801748..51c626807 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -192,13 +192,13 @@ def _is_smoke_item(item): "tests/test_add.py": _is_smoke_add_case, "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_internal_causal_softmax.py": _is_smoke_internal_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_swiglu.py": _is_smoke_swiglu_case, + "tests/test_silu_and_mul.py": _is_smoke_silu_and_mul_case, "tests/test_torch_ops.py": _is_smoke_torch_op_case, } matcher = matchers.get(module_path) @@ -378,26 +378,25 @@ def _is_smoke_rms_norm_case(params): ) -def _is_smoke_swiglu_case(params): +def _is_smoke_silu_and_mul_case(params): cases = { - ((13, 4), None, None, None), - ((16, 5632), None, None, None), + ((13, 4), None, None), + ((16, 5632), None, None), } return ( _is_float32(params) and _shape_case( params, - "shape", + "out_shape", "input_strides", - "gate_strides", "out_strides", ) in cases ) -def _is_smoke_causal_softmax_case(params): +def _is_smoke_internal_causal_softmax_case(params): cases = { ((3, 3), None, None), ((32, 512), None, None), From 5057edaf6486e0dc07226e09ca747511f7ce26aa Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 14 Jul 2026 18:26:31 +0800 Subject: [PATCH 12/22] test(rotary): remove unused test variable --- tests/test_rotary_embedding.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_rotary_embedding.py b/tests/test_rotary_embedding.py index a714769a1..682ac3cd4 100644 --- a/tests/test_rotary_embedding.py +++ b/tests/test_rotary_embedding.py @@ -39,7 +39,6 @@ def test_rotary_embedding( rtol, atol, ): - num_tokens = 4 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( From 8330e719c2995dc3813be602cdb646352180e11b Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 09:09:09 +0800 Subject: [PATCH 13/22] fix(rotary): preserve writable optional key --- src/base/rotary_embedding.h | 6 +++--- src/native/cuda/ops/rotary_embedding/kernel.h | 2 +- src/torch/ops/rotary_embedding/rotary_embedding.cc | 6 +++--- src/torch/ops/rotary_embedding/rotary_embedding.h | 7 +++---- 4 files changed, 10 insertions(+), 11 deletions(-) diff --git a/src/base/rotary_embedding.h b/src/base/rotary_embedding.h index 0f11d2f45..f405faa43 100644 --- a/src/base/rotary_embedding.h +++ b/src/base/rotary_embedding.h @@ -11,7 +11,7 @@ namespace infini::ops { class RotaryEmbedding : public Operator { public: RotaryEmbedding(const Tensor positions, Tensor query, - const std::optional key, int64_t head_size, + 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()}, @@ -28,7 +28,7 @@ class RotaryEmbedding : public Operator { num_tokens_{positions.numel()}, positions_ndim_{positions.ndim()}, head_size_{head_size}, - rot_dim_{cos_sin_cache.size(1)}, + rot_dim_{static_cast(cos_sin_cache.size(1))}, rope_dim_offset_{rope_dim_offset}, is_neox_{is_neox}, inverse_{inverse}, @@ -117,7 +117,7 @@ class RotaryEmbedding : public Operator { } virtual void operator()(const Tensor positions, Tensor query, - const std::optional key, int64_t head_size, + 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; diff --git a/src/native/cuda/ops/rotary_embedding/kernel.h b/src/native/cuda/ops/rotary_embedding/kernel.h index 8566303a5..f10936909 100644 --- a/src/native/cuda/ops/rotary_embedding/kernel.h +++ b/src/native/cuda/ops/rotary_embedding/kernel.h @@ -18,7 +18,7 @@ class CudaRotaryEmbedding : public RotaryEmbedding { using RotaryEmbedding::RotaryEmbedding; void operator()(const Tensor positions, Tensor query, - const std::optional key, int64_t, + std::optional key, int64_t, const Tensor cos_sin_cache, bool, int64_t, bool) const override { if (num_tokens_ == 0) { diff --git a/src/torch/ops/rotary_embedding/rotary_embedding.cc b/src/torch/ops/rotary_embedding/rotary_embedding.cc index b30f8c33b..c95360e3b 100644 --- a/src/torch/ops/rotary_embedding/rotary_embedding.cc +++ b/src/torch/ops/rotary_embedding/rotary_embedding.cc @@ -39,7 +39,7 @@ void ApplyRotaryEmbedding(at::Tensor data, const at::Tensor& positions, template Operator::Operator( - const Tensor positions, Tensor query, const std::optional key, + 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, @@ -48,8 +48,8 @@ Operator::Operator( template void Operator::operator()( - const Tensor positions, Tensor query, const std::optional key, - int64_t, const Tensor cos_sin_cache, bool, int64_t, bool) const { + 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; } diff --git a/src/torch/ops/rotary_embedding/rotary_embedding.h b/src/torch/ops/rotary_embedding/rotary_embedding.h index ec1d6811a..916db723c 100644 --- a/src/torch/ops/rotary_embedding/rotary_embedding.h +++ b/src/torch/ops/rotary_embedding/rotary_embedding.h @@ -10,13 +10,12 @@ namespace infini::ops { template class Operator : public RotaryEmbedding { public: - Operator(const Tensor positions, Tensor query, - const std::optional key, int64_t head_size, - const Tensor cos_sin_cache, bool is_neox, + 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, - const std::optional key, int64_t head_size, + 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; From f3fa89075fb0f0b064a2ee43e5f191897213659e Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 10:12:20 +0800 Subject: [PATCH 14/22] fix(torch): remove retired device instantiations --- src/torch/ops/reshape_and_cache/reshape_and_cache.cc | 2 -- src/torch/ops/rotary_embedding/rotary_embedding.cc | 2 -- .../scaled_dot_product_attention.cc | 2 -- 3 files changed, 6 deletions(-) diff --git a/src/torch/ops/reshape_and_cache/reshape_and_cache.cc b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc index c91c05363..c1d5f45f2 100644 --- a/src/torch/ops/reshape_and_cache/reshape_and_cache.cc +++ b/src/torch/ops/reshape_and_cache/reshape_and_cache.cc @@ -84,8 +84,6 @@ 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.cc b/src/torch/ops/rotary_embedding/rotary_embedding.cc index c95360e3b..96a9c53ef 100644 --- a/src/torch/ops/rotary_embedding/rotary_embedding.cc +++ b/src/torch/ops/rotary_embedding/rotary_embedding.cc @@ -82,8 +82,6 @@ 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.cc b/src/torch/ops/scaled_dot_product_attention/scaled_dot_product_attention.cc index 169c1eacd..789d47f85 100644 --- 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 @@ -60,8 +60,6 @@ template class Operator; template class Operator; template class Operator; template class Operator; -template class Operator; template class Operator; -template class Operator; } // namespace infini::ops From d82d59e2b72ba7284341c911c8af68fb0795cc5e Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 10:54:35 +0800 Subject: [PATCH 15/22] chore(tests): apply ruff formatting --- tests/test_add.py | 4 +--- tests/test_embedding.py | 4 +--- tests/test_fused_add_rms_norm.py | 4 +--- tests/test_internal_causal_softmax.py | 4 +--- tests/test_internal_top_k_top_p_sample.py | 4 +--- tests/test_linear.py | 8 ++------ tests/test_reshape_and_cache.py | 12 +++--------- 7 files changed, 10 insertions(+), 30 deletions(-) diff --git a/tests/test_add.py b/tests/test_add.py index f54073e28..ae826cc30 100644 --- a/tests/test_add.py +++ b/tests/test_add.py @@ -120,9 +120,7 @@ def _torch_add(input, other, out): ), ) @pytest.mark.parametrize("alpha", (0.0, 0.5, 1.0, 2.0)) -@pytest.mark.parametrize( - "dtype", (torch.float32, torch.float16, torch.bfloat16) -) +@pytest.mark.parametrize("dtype", (torch.float32, torch.float16, torch.bfloat16)) def test_add_alpha_and_broadcast( input_shape, other_shape, diff --git a/tests/test_embedding.py b/tests/test_embedding.py index 499a5776b..7bd22a3bd 100644 --- a/tests/test_embedding.py +++ b/tests/test_embedding.py @@ -111,9 +111,7 @@ def test_embedding_parameters( 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 - ) + 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( diff --git a/tests/test_fused_add_rms_norm.py b/tests/test_fused_add_rms_norm.py index 377a7bd85..77684bf5a 100644 --- a/tests/test_fused_add_rms_norm.py +++ b/tests/test_fused_add_rms_norm.py @@ -44,9 +44,7 @@ def test_fused_add_rms_norm( ): 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 - ) + weight = torch.randn(shape[-1], dtype=dtype, device=device) if has_weight else None expected_input = clone_strided(input) expected_residual = clone_strided(residual) diff --git a/tests/test_internal_causal_softmax.py b/tests/test_internal_causal_softmax.py index e3b55a8e1..33b110f77 100644 --- a/tests/test_internal_causal_softmax.py +++ b/tests/test_internal_causal_softmax.py @@ -42,9 +42,7 @@ def test_internal_causal_softmax( def _internal_causal_softmax(input, out): - infini.ops.internal_causal_softmax( - input, out, stream=get_stream(input.device) - ) + infini.ops.internal_causal_softmax(input, out, stream=get_stream(input.device)) return out diff --git a/tests/test_internal_top_k_top_p_sample.py b/tests/test_internal_top_k_top_p_sample.py index 353cea7d8..3f151bf1d 100644 --- a/tests/test_internal_top_k_top_p_sample.py +++ b/tests/test_internal_top_k_top_p_sample.py @@ -65,9 +65,7 @@ def test_internal_top_k_top_p_sample_filters( ) out = torch.empty((32,), dtype=torch.int32, device=device) - _internal_top_k_top_p_sample( - logits, k, p, 1234, 0, out, implementation_index - ) + _internal_top_k_top_p_sample(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)) diff --git a/tests/test_linear.py b/tests/test_linear.py index b51b6414d..ba692d323 100644 --- a/tests/test_linear.py +++ b/tests/test_linear.py @@ -58,9 +58,7 @@ def test_linear( def _linear(input, weight, bias, out): - infini.ops.linear( - input, weight, bias, out, stream=get_stream(input.device) - ) + infini.ops.linear(input, weight, bias, out, stream=get_stream(input.device)) return out @@ -76,9 +74,7 @@ def _torch_linear(input, weight, bias, out): @pytest.mark.auto_act_and_assert -@pytest.mark.parametrize( - "dtype", (torch.float32, torch.float16, torch.bfloat16) -) +@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) diff --git a/tests/test_reshape_and_cache.py b/tests/test_reshape_and_cache.py index b24d32813..773b20ac7 100644 --- a/tests/test_reshape_and_cache.py +++ b/tests/test_reshape_and_cache.py @@ -87,12 +87,8 @@ def test_reshape_and_cache( ) 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 - ) + 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( @@ -118,9 +114,7 @@ def _torch_reshape_and_cache( if kv_cache_dtype != "auto": fp8_dtype = ( - torch.float8_e5m2 - if kv_cache_dtype == "fp8_e5m2" - else torch.float8_e4m3fn + 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) From d3e888d1e37031b82ecb7362788eda7741425705 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 11:45:40 +0800 Subject: [PATCH 16/22] fix(attention): support older PyTorch toolchains --- src/base/scaled_dot_product_attention.h | 17 +++++++------ .../scaled_dot_product_attention.cc | 25 ++++++++++++++++--- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/src/base/scaled_dot_product_attention.h b/src/base/scaled_dot_product_attention.h index c6d24d4c7..2dc37919e 100644 --- a/src/base/scaled_dot_product_attention.h +++ b/src/base/scaled_dot_product_attention.h @@ -8,6 +8,15 @@ namespace infini::ops { class ScaledDotProductAttention : public Operator { + protected: + template + static auto ReturnShape(const TensorLike& query, const TensorLike& value) { + auto shape = query.shape(); + shape.back() = value.size(-1); + + return shape; + } + public: ScaledDotProductAttention(const Tensor query, const Tensor key, const Tensor value, @@ -86,14 +95,6 @@ class ScaledDotProductAttention : public Operator { } protected: - template - static auto ReturnShape(const TensorLike& query, const TensorLike& value) { - auto shape = query.shape(); - shape.back() = value.size(-1); - - return shape; - } - Tensor::Shape query_shape_; Tensor::Shape key_shape_; 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 index 789d47f85..81489cf8c 100644 --- 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 @@ -37,7 +37,7 @@ void Operator::operator()( auto at_out = ToAtenTensor(out.data(), out_shape_, out_strides_, query_type_, device_index_); - std::optional at_attn_mask; + c10::optional at_attn_mask; if (attn_mask.has_value()) { const auto dtype_override = attn_mask_type_ == DataType::kUInt8 ? std::optional{at::kBool} @@ -47,9 +47,26 @@ void Operator::operator()( attn_mask_strides_, attn_mask_type_, device_index_, dtype_override)); } - auto result = - at::scaled_dot_product_attention(at_query, at_key, at_value, at_attn_mask, - dropout_p, is_causal, scale, enable_gqa); + 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); } From 115ff3d59010874bf0c80cc6874c0f4dbac5facc Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 16:51:49 +0800 Subject: [PATCH 17/22] refactor(ops): suffix custom kernels with Infinilm --- src/CMakeLists.txt | 2 +- ...al_softmax.h => causal_softmax_infinilm.h} | 20 ++++----- ...ed_softmax.h => scaled_softmax_infinilm.h} | 27 ++++++------ ...sample.h => top_k_top_p_sample_infinilm.h} | 42 ++++++++++--------- .../kernel.h | 22 +++++----- .../kernel.h | 41 ++++++++++-------- .../causal_softmax_infinilm.h} | 14 +++---- .../top_k_top_p_sample_infinilm.h} | 18 ++++---- .../ops/causal_softmax_infinilm/kernel.h | 22 ++++++++++ .../ops/internal_causal_softmax/kernel.h | 22 ---------- .../ops/causal_softmax_infinilm/kernel.h | 22 ++++++++++ .../ops/internal_causal_softmax/kernel.h | 22 ---------- .../ops/causal_softmax_infinilm/kernel.h | 33 +++++++++++++++ .../ops/internal_causal_softmax/kernel.h | 37 ---------------- .../ops/causal_softmax_infinilm/kernel.h | 22 ++++++++++ .../ops/internal_causal_softmax/kernel.h | 22 ---------- .../kernel.cuh | 10 ++--- .../kernel.h | 20 ++++----- tests/conftest.py | 4 +- ...max.py => test_causal_softmax_infinilm.py} | 8 ++-- ...max.py => test_scaled_softmax_infinilm.py} | 8 ++-- ...py => test_top_k_top_p_sample_infinilm.py} | 18 ++++---- 22 files changed, 232 insertions(+), 224 deletions(-) rename src/base/{internal_causal_softmax.h => causal_softmax_infinilm.h} (59%) rename src/base/{internal_scaled_softmax.h => scaled_softmax_infinilm.h} (54%) rename src/base/{internal_top_k_top_p_sample.h => top_k_top_p_sample_infinilm.h} (54%) rename src/native/ascend/ops/{internal_scaled_softmax => scaled_softmax_infinilm}/kernel.h (82%) rename src/native/ascend/ops/{internal_top_k_top_p_sample => top_k_top_p_sample_infinilm}/kernel.h (87%) rename src/native/cpu/ops/{internal_causal_softmax/internal_causal_softmax.h => causal_softmax_infinilm/causal_softmax_infinilm.h} (85%) rename src/native/cpu/ops/{internal_top_k_top_p_sample/internal_top_k_top_p_sample.h => top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h} (89%) create mode 100644 src/native/cuda/iluvatar/ops/causal_softmax_infinilm/kernel.h delete mode 100644 src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h create mode 100644 src/native/cuda/metax/ops/causal_softmax_infinilm/kernel.h delete mode 100644 src/native/cuda/metax/ops/internal_causal_softmax/kernel.h create mode 100644 src/native/cuda/moore/ops/causal_softmax_infinilm/kernel.h delete mode 100644 src/native/cuda/moore/ops/internal_causal_softmax/kernel.h create mode 100644 src/native/cuda/nvidia/ops/causal_softmax_infinilm/kernel.h delete mode 100644 src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h rename src/native/cuda/ops/{internal_causal_softmax => causal_softmax_infinilm}/kernel.cuh (93%) rename src/native/cuda/ops/{internal_causal_softmax => causal_softmax_infinilm}/kernel.h (77%) rename tests/{test_internal_causal_softmax.py => test_causal_softmax_infinilm.py} (89%) rename tests/{test_internal_scaled_softmax.py => test_scaled_softmax_infinilm.py} (88%) rename tests/{test_internal_top_k_top_p_sample.py => test_top_k_top_p_sample_infinilm.py} (84%) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 78e3affad..0e3567308 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -420,7 +420,7 @@ 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 silu_and_mul - internal_causal_softmax abs clamp exp) + 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/internal_causal_softmax.h b/src/base/causal_softmax_infinilm.h similarity index 59% rename from src/base/internal_causal_softmax.h rename to src/base/causal_softmax_infinilm.h index 4b9ae1855..31bb206d1 100644 --- a/src/base/internal_causal_softmax.h +++ b/src/base/causal_softmax_infinilm.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_INTERNAL_CAUSAL_SOFTMAX_H_ -#define INFINI_OPS_BASE_INTERNAL_CAUSAL_SOFTMAX_H_ +#ifndef INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_BASE_CAUSAL_SOFTMAX_INFINILM_H_ #include #include @@ -7,11 +7,11 @@ #include "operator.h" #include "tensor.h" -namespace infini::ops::internal { +namespace infini::ops { -class CausalSoftmax : public Operator { +class CausalSoftmaxInfinilm : public Operator { public: - CausalSoftmax(const Tensor input, Tensor out) + CausalSoftmaxInfinilm(const Tensor input, Tensor out) : dtype_{input.dtype()}, ndim_{out.ndim()}, batch_size_{ndim_ == 2 ? 1 : out.size(-3)}, @@ -20,13 +20,13 @@ class CausalSoftmax : public Operator { input_strides_{input.strides()}, out_strides_{out.strides()} { assert(input.shape() == out.shape() && - "`CausalSoftmax` requires `input` and `out` same shape"); + "`CausalSoftmaxInfinilm` requires `input` and `out` same shape"); assert(input.dtype() == out.dtype() && - "`CausalSoftmax` requires `input` and `out` same dtype"); + "`CausalSoftmaxInfinilm` requires `input` and `out` same dtype"); assert((ndim_ == 2 || ndim_ == 3) && - "`CausalSoftmax` requires 2D or 3D tensor"); + "`CausalSoftmaxInfinilm` requires 2D or 3D tensor"); assert(seq_len_ <= total_seq_len_ && - "`CausalSoftmax` requires shape[-2] <= shape[-1]"); + "`CausalSoftmaxInfinilm` requires shape[-2] <= shape[-1]"); } virtual void operator()(const Tensor input, Tensor out) const = 0; @@ -47,6 +47,6 @@ class CausalSoftmax : public Operator { Tensor::Strides out_strides_; }; -} // namespace infini::ops::internal +} // namespace infini::ops #endif diff --git a/src/base/internal_scaled_softmax.h b/src/base/scaled_softmax_infinilm.h similarity index 54% rename from src/base/internal_scaled_softmax.h rename to src/base/scaled_softmax_infinilm.h index 3d97bba9d..1ca081be1 100644 --- a/src/base/internal_scaled_softmax.h +++ b/src/base/scaled_softmax_infinilm.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ -#define INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ +#ifndef INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ #include #include @@ -9,11 +9,11 @@ #include "operator.h" #include "tensor.h" -namespace infini::ops::internal { +namespace infini::ops { -class ScaledSoftmax : public Operator { +class ScaledSoftmaxInfinilm : public Operator { public: - ScaledSoftmax(const Tensor input, double scale, Tensor out) + ScaledSoftmaxInfinilm(const Tensor input, double scale, Tensor out) : scale_{scale}, batch_size_{input.size(0)}, vocab_size_{input.size(1)}, @@ -21,16 +21,19 @@ class ScaledSoftmax : public Operator { input_strides_{input.strides()}, out_strides_{out.strides()} { assert(input.ndim() == 2 && - "`ScaledSoftmax` currently supports 2D `[batch, vocab]` input"); + "`ScaledSoftmaxInfinilm` currently supports 2D `[batch, vocab]` " + "input"); assert(input.shape() == out.shape() && - "`ScaledSoftmax` requires `input` and `out` to have the same shape"); + "`ScaledSoftmaxInfinilm` requires `input` and `out` to have the " + "same shape"); assert(input.dtype() == out.dtype() && - "`ScaledSoftmax` requires `input` and `out` to have the same dtype"); + "`ScaledSoftmaxInfinilm` requires `input` and `out` to have the " + "same dtype"); assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && - "`ScaledSoftmax` requires a floating point dtype"); + "`ScaledSoftmaxInfinilm` requires a floating point dtype"); assert(std::isfinite(scale_) && - "`ScaledSoftmax` requires a finite `scale`"); + "`ScaledSoftmaxInfinilm` requires a finite `scale`"); } virtual void operator()(const Tensor input, double scale, @@ -50,6 +53,6 @@ class ScaledSoftmax : public Operator { Tensor::Strides out_strides_; }; -} // namespace infini::ops::internal +} // namespace infini::ops -#endif // INFINI_OPS_BASE_INTERNAL_SCALED_SOFTMAX_H_ +#endif // INFINI_OPS_BASE_SCALED_SOFTMAX_INFINILM_H_ diff --git a/src/base/internal_top_k_top_p_sample.h b/src/base/top_k_top_p_sample_infinilm.h similarity index 54% rename from src/base/internal_top_k_top_p_sample.h rename to src/base/top_k_top_p_sample_infinilm.h index 6e74b46fa..2d9b092d9 100644 --- a/src/base/internal_top_k_top_p_sample.h +++ b/src/base/top_k_top_p_sample_infinilm.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ -#define INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#ifndef INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ +#define INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ #include #include @@ -9,29 +9,31 @@ #include "operator.h" #include "tensor.h" -namespace infini::ops::internal { +namespace infini::ops { -class TopKTopPSample : public Operator { +class TopKTopPSampleInfinilm : public Operator { public: - TopKTopPSample(const Tensor logits, std::optional k, - std::optional p, uint64_t seed, uint64_t offset, - Tensor out) + 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 && - "`TopKTopPSample` requires 2D `[batch_size, vocab_size]` logits"); + "`TopKTopPSampleInfinilm` requires 2D `[batch_size, vocab_size]` " + "logits"); assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && - "`TopKTopPSample` requires floating-point logits"); + "`TopKTopPSampleInfinilm` requires floating-point logits"); assert(out.ndim() == 1 && - "`TopKTopPSample` requires 1D `[batch_size]` output"); + "`TopKTopPSampleInfinilm` requires 1D `[batch_size]` output"); assert(out.size(0) == batch_size_ && - "`TopKTopPSample` requires output batch size to match logits"); + "`TopKTopPSampleInfinilm` requires output batch size to match " + "logits"); assert(out.dtype() == DataType::kInt32 && - "`TopKTopPSample` requires int32 output"); + "`TopKTopPSampleInfinilm` requires int32 output"); ValidateK(k); ValidateP(p); @@ -46,25 +48,25 @@ class TopKTopPSample : public Operator { if (!k.has_value()) return; assert(k->ndim() == 1 && - "`TopKTopPSample` requires `k` to be 1D when provided"); + "`TopKTopPSampleInfinilm` requires `k` to be 1D when provided"); assert((k->size(0) == 1 || k->size(0) == batch_size_) && - "`TopKTopPSample` requires `k` shape [1] or [batch_size]"); + "`TopKTopPSampleInfinilm` requires `k` shape [1] or [batch_size]"); assert((k->dtype() == DataType::kInt32 || k->dtype() == DataType::kInt64) && - "`TopKTopPSample` requires int32 or int64 `k`"); + "`TopKTopPSampleInfinilm` requires int32 or int64 `k`"); } void ValidateP(std::optional p) const { if (!p.has_value()) return; assert(p->ndim() == 1 && - "`TopKTopPSample` requires `p` to be 1D when provided"); + "`TopKTopPSampleInfinilm` requires `p` to be 1D when provided"); assert((p->size(0) == 1 || p->size(0) == batch_size_) && - "`TopKTopPSample` requires `p` shape [1] or [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) && - "`TopKTopPSample` requires floating-point `p`"); + "`TopKTopPSampleInfinilm` requires floating-point `p`"); } Tensor::Size batch_size_{0}; @@ -74,6 +76,6 @@ class TopKTopPSample : public Operator { DataType dtype_; }; -} // namespace infini::ops::internal +} // namespace infini::ops -#endif // INFINI_OPS_BASE_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#endif // INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLE_INFINILM_H_ diff --git a/src/native/ascend/ops/internal_scaled_softmax/kernel.h b/src/native/ascend/ops/scaled_softmax_infinilm/kernel.h similarity index 82% rename from src/native/ascend/ops/internal_scaled_softmax/kernel.h rename to src/native/ascend/ops/scaled_softmax_infinilm/kernel.h index 291fdb204..c4d11f516 100644 --- a/src/native/ascend/ops/internal_scaled_softmax/kernel.h +++ b/src/native/ascend/ops/scaled_softmax_infinilm/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ +#ifndef INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ #include @@ -7,7 +7,7 @@ #include "aclnn/aclnn_base.h" #include "aclnn_mul.h" #include "aclnn_softmax.h" -#include "base/internal_scaled_softmax.h" +#include "base/scaled_softmax_infinilm.h" #include "data_type.h" #include "native/ascend/common.h" #include "native/ascend/workspace_pool_.h" @@ -16,11 +16,11 @@ namespace infini::ops { template <> -class Operator - : public internal::ScaledSoftmax { +class Operator + : public ScaledSoftmaxInfinilm { public: Operator(const Tensor input, double scale, Tensor out) - : internal::ScaledSoftmax(input, scale, out), + : ScaledSoftmaxInfinilm(input, scale, out), in_cache_(input), out_cache_(out), temp_cache_(input), @@ -28,12 +28,12 @@ class Operator needs_scale_(std::fabs(scale - 1.0) > 1e-6) { assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || dtype_ == DataType::kFloat32) && - "`ScaledSoftmax` Ascend path requires float16, bfloat16, or " + "`ScaledSoftmaxInfinilm` Ascend path requires float16, bfloat16, or " "float32 input"); assert(input.IsContiguous() && - "`ScaledSoftmax` Ascend path requires contiguous input"); + "`ScaledSoftmaxInfinilm` Ascend path requires contiguous input"); assert(out.IsContiguous() && - "`ScaledSoftmax` Ascend path requires contiguous output"); + "`ScaledSoftmaxInfinilm` Ascend path requires contiguous output"); temp_size_ = input.numel() * kDataTypeToSize.at(dtype_); scale_scalar_ = aclCreateScalar(&scale_storage_, ACL_FLOAT); @@ -51,7 +51,7 @@ class Operator void operator()(const Tensor input, double scale, Tensor out) const override { assert(scale == scale_ && - "`ScaledSoftmax` scale changed after descriptor creation"); + "`ScaledSoftmaxInfinilm` scale changed after descriptor creation"); auto stream = static_cast(stream_); auto t_in = in_cache_.get(const_cast(input.data())); @@ -122,4 +122,4 @@ class Operator } // namespace infini::ops -#endif // INFINI_OPS_ASCEND_INTERNAL_SCALED_SOFTMAX_KERNEL_H_ +#endif // INFINI_OPS_ASCEND_SCALED_SOFTMAX_INFINILM_KERNEL_H_ diff --git a/src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h b/src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h similarity index 87% rename from src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h rename to src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h index 0a48701ad..bf820a7ca 100644 --- a/src/native/ascend/ops/internal_top_k_top_p_sample/kernel.h +++ b/src/native/ascend/ops/top_k_top_p_sample_infinilm/kernel.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ -#define INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ +#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 @@ -13,7 +13,7 @@ #include "aclnn/aclnn_base.h" #include "aclnnop/aclnn_cast.h" #include "aclnnop/aclnn_top_k_top_p_sample.h" -#include "base/internal_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" @@ -23,19 +23,22 @@ namespace infini::ops { template <> -class Operator - : public internal::TopKTopPSample { +class Operator + : public TopKTopPSampleInfinilm { public: Operator(const Tensor logits, std::optional k, std::optional p, uint64_t seed, uint64_t offset, Tensor out) - : internal::TopKTopPSample(logits, k, p, seed, offset, out) { + : TopKTopPSampleInfinilm(logits, k, p, seed, offset, out) { assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16) && - "`TopKTopPSample` Ascend ACLNN path requires float16 or bfloat16 " + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires float16 or " + "bfloat16 " "logits"); assert(logits.IsContiguous() && - "`TopKTopPSample` Ascend ACLNN path requires contiguous logits"); + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "logits"); assert(out.IsContiguous() && - "`TopKTopPSample` Ascend ACLNN path requires contiguous output"); + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "output"); ValidateHostTensor(k); ValidateHostTensor(p); @@ -71,9 +74,11 @@ class Operator std::optional p, uint64_t seed, uint64_t offset, Tensor out) const override { assert(logits.IsContiguous() && - "`TopKTopPSample` Ascend ACLNN path requires contiguous logits"); + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "logits"); assert(out.IsContiguous() && - "`TopKTopPSample` Ascend ACLNN path requires contiguous output"); + "`TopKTopPSampleInfinilm` Ascend ACLNN path requires contiguous " + "output"); auto stream = static_cast(stream_); auto top_k_bytes = batch_size_ * kDataTypeToSize.at(DataType::kInt32); @@ -96,15 +101,15 @@ class Operator 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 && - "`TopKTopPSample`: copying `top_k` to Ascend failed"); + "`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 && - "`TopKTopPSample`: copying `top_p` to Ascend failed"); + "`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 && - "`TopKTopPSample`: copying `q` to Ascend failed"); + "`TopKTopPSampleInfinilm`: copying `q` to Ascend failed"); auto& selected_idx_arena = ascend::GetWorkspacePool().Ensure( stream, selected_idx_bytes, "top_k_top_p_sample_idx"); @@ -153,10 +158,10 @@ class Operator if (!tensor.has_value()) return; assert(tensor->device().type() == Device::Type::kCpu && - "`TopKTopPSample` Ascend path currently requires host-side " + "`TopKTopPSampleInfinilm` Ascend path currently requires host-side " "`k`/`p` tensors"); assert(tensor->IsContiguous() && - "`TopKTopPSample` Ascend path requires contiguous `k`/`p` " + "`TopKTopPSampleInfinilm` Ascend path requires contiguous `k`/`p` " "tensors"); } @@ -245,7 +250,7 @@ class Operator value = static_cast(p->data())[offset]; break; default: - assert(false && "`TopKTopPSample` has unsupported `p` dtype"); + assert(false && "`TopKTopPSampleInfinilm` has unsupported `p` dtype"); } if (value <= 0.0 || value > 1.0) return 1.0; @@ -283,4 +288,4 @@ class Operator } // namespace infini::ops -#endif // INFINI_OPS_ASCEND_INTERNAL_TOP_K_TOP_P_SAMPLE_KERNEL_H_ +#endif // INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLE_INFINILM_KERNEL_H_ diff --git a/src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h b/src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h similarity index 85% rename from src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h rename to src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h index 4e7f62524..94d8709b5 100644 --- a/src/native/cpu/ops/internal_causal_softmax/internal_causal_softmax.h +++ b/src/native/cpu/ops/causal_softmax_infinilm/causal_softmax_infinilm.h @@ -1,9 +1,9 @@ -#ifndef INFINI_OPS_CPU_INTERNAL_CAUSAL_SOFTMAX_H_ -#define INFINI_OPS_CPU_INTERNAL_CAUSAL_SOFTMAX_H_ +#ifndef INFINI_OPS_CPU_CAUSAL_SOFTMAX_INFINILM_H_ +#define INFINI_OPS_CPU_CAUSAL_SOFTMAX_INFINILM_H_ #include -#include "base/internal_causal_softmax.h" +#include "base/causal_softmax_infinilm.h" #include "common/generic_utils.h" #include "data_type.h" #include "native/cpu/caster_.h" @@ -12,11 +12,11 @@ namespace infini::ops { template <> -class Operator - : public internal::CausalSoftmax, Caster { +class Operator + : public CausalSoftmaxInfinilm, Caster { public: Operator(const Tensor input, Tensor out) - : internal::CausalSoftmax{input, out} {} + : CausalSoftmaxInfinilm{input, out} {} void operator()(const Tensor input, Tensor out) const override { DispatchFunc( @@ -25,7 +25,7 @@ class Operator using T = typename decltype(tag)::type; Compute(input, out); }, - "`Operator::operator()`"); + "`Operator::operator()`"); } private: diff --git a/src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h b/src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h similarity index 89% rename from src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h rename to src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h index f5d4fcb0b..409fc960b 100644 --- a/src/native/cpu/ops/internal_top_k_top_p_sample/internal_top_k_top_p_sample.h +++ b/src/native/cpu/ops/top_k_top_p_sample_infinilm/top_k_top_p_sample_infinilm.h @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ -#define INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#ifndef INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLE_INFINILM_H_ +#define INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLE_INFINILM_H_ #include #include @@ -11,7 +11,7 @@ #include #include -#include "base/internal_top_k_top_p_sample.h" +#include "base/top_k_top_p_sample_infinilm.h" #include "data_type.h" #include "native/cpu/caster_.h" #include "operator.h" @@ -20,12 +20,12 @@ namespace infini::ops { template <> -class Operator - : public internal::TopKTopPSample, Caster { +class Operator + : public TopKTopPSampleInfinilm, Caster { public: Operator(const Tensor logits, std::optional k, std::optional p, uint64_t seed, uint64_t offset, Tensor out) - : internal::TopKTopPSample(logits, k, p, seed, offset, 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, @@ -36,7 +36,7 @@ class Operator using T = typename decltype(tag)::type; Compute(logits, k, p, seed, offset, out); }, - "`Operator::operator()`"); + "`Operator::operator()`"); } private: @@ -152,7 +152,7 @@ class Operator case DataType::kFloat64: return static_cast(p->data())[offset]; default: - assert(false && "`TopKTopPSample` has unsupported `p` dtype."); + assert(false && "`TopKTopPSampleInfinilm` has unsupported `p` dtype."); return 1.0; } } @@ -160,4 +160,4 @@ class Operator } // namespace infini::ops -#endif // INFINI_OPS_CPU_INTERNAL_TOP_K_TOP_P_SAMPLE_H_ +#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/internal_causal_softmax/kernel.h b/src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h deleted file mode 100644 index 0f6bf9739..000000000 --- a/src/native/cuda/iluvatar/ops/internal_causal_softmax/kernel.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef INFINI_OPS_ILUVATAR_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_ILUVATAR_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/iluvatar/caster.cuh" -#include "native/cuda/iluvatar/runtime_.h" -#include "native/cuda/ops/internal_causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public internal::CudaCausalSoftmax> { - public: - using internal::CudaCausalSoftmax< - Runtime>::CudaCausalSoftmax; -}; - -} // 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/internal_causal_softmax/kernel.h b/src/native/cuda/metax/ops/internal_causal_softmax/kernel.h deleted file mode 100644 index 167653961..000000000 --- a/src/native/cuda/metax/ops/internal_causal_softmax/kernel.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef INFINI_OPS_METAX_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_METAX_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/metax/caster.cuh" -#include "native/cuda/metax/runtime_.h" -#include "native/cuda/ops/internal_causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public internal::CudaCausalSoftmax> { - public: - using internal::CudaCausalSoftmax< - Runtime>::CudaCausalSoftmax; -}; - -} // 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/internal_causal_softmax/kernel.h b/src/native/cuda/moore/ops/internal_causal_softmax/kernel.h deleted file mode 100644 index d608784a7..000000000 --- a/src/native/cuda/moore/ops/internal_causal_softmax/kernel.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef INFINI_OPS_MOORE_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_MOORE_INTERNAL_CAUSAL_SOFTMAX_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/internal_causal_softmax/kernel.h" - -namespace infini::ops::internal { - -struct MooreCausalSoftmaxBackend : Runtime { - // Moore's causal softmax kernel should not dispatch block sizes above 1024. - static constexpr int max_block_size = 1024; -}; - -} // namespace infini::ops::internal - -namespace infini::ops { - -template <> -class Operator - : public internal::CudaCausalSoftmax { - public: - using internal::CudaCausalSoftmax< - internal::MooreCausalSoftmaxBackend>::CudaCausalSoftmax; -}; - -} // 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/internal_causal_softmax/kernel.h b/src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h deleted file mode 100644 index d7bf94dd4..000000000 --- a/src/native/cuda/nvidia/ops/internal_causal_softmax/kernel.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef INFINI_OPS_NVIDIA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_NVIDIA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ - -#include - -#include "native/cuda/nvidia/caster.cuh" -#include "native/cuda/nvidia/runtime_.h" -#include "native/cuda/ops/internal_causal_softmax/kernel.h" - -namespace infini::ops { - -template <> -class Operator - : public internal::CudaCausalSoftmax> { - public: - using internal::CudaCausalSoftmax< - Runtime>::CudaCausalSoftmax; -}; - -} // namespace infini::ops - -#endif diff --git a/src/native/cuda/ops/internal_causal_softmax/kernel.cuh b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh similarity index 93% rename from src/native/cuda/ops/internal_causal_softmax/kernel.cuh rename to src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh index 4079a508f..5a061c662 100644 --- a/src/native/cuda/ops/internal_causal_softmax/kernel.cuh +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh @@ -1,5 +1,5 @@ -#ifndef INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_CUH_ -#define INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_CUH_ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_CUH_ #include #include @@ -8,7 +8,7 @@ #include "native/cuda/caster.cuh" #include "native/cuda/kernel_commons.cuh" -namespace infini::ops::internal { +namespace infini::ops { namespace { @@ -54,7 +54,7 @@ __device__ __forceinline__ Compute BlockSum(const Data* data_ptr, template -__global__ void CausalSoftmaxKernel( +__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, @@ -101,6 +101,6 @@ __global__ void CausalSoftmaxKernel( } } -} // namespace infini::ops::internal +} // namespace infini::ops #endif diff --git a/src/native/cuda/ops/internal_causal_softmax/kernel.h b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h similarity index 77% rename from src/native/cuda/ops/internal_causal_softmax/kernel.h rename to src/native/cuda/ops/causal_softmax_infinilm/kernel.h index 6af110771..d2f6bf882 100644 --- a/src/native/cuda/ops/internal_causal_softmax/kernel.h +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h @@ -1,23 +1,23 @@ -#ifndef INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ -#define INFINI_OPS_CUDA_INTERNAL_CAUSAL_SOFTMAX_KERNEL_H_ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_INFINILM_KERNEL_H_ #include #include #include -#include "base/internal_causal_softmax.h" +#include "base/causal_softmax_infinilm.h" #include "data_type.h" #include "dispatcher.h" #include "native/cuda/kernel_commons.cuh" -#include "native/cuda/ops/internal_causal_softmax/kernel.cuh" +#include "native/cuda/ops/causal_softmax_infinilm/kernel.cuh" #include "native/cuda/runtime_utils.h" -namespace infini::ops::internal { +namespace infini::ops { template -class CudaCausalSoftmax : public CausalSoftmax { +class CudaCausalSoftmaxInfinilm : public CausalSoftmaxInfinilm { public: - using CausalSoftmax::CausalSoftmax; + using CausalSoftmaxInfinilm::CausalSoftmaxInfinilm; void operator()(const Tensor input, Tensor out) const override { auto cuda_stream = @@ -47,17 +47,17 @@ class CudaCausalSoftmax : public CausalSoftmax { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); - CausalSoftmaxKernel + 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); }, - "CudaCausalSoftmax::operator()"); + "CudaCausalSoftmaxInfinilm::operator()"); } }; -} // namespace infini::ops::internal +} // namespace infini::ops #endif diff --git a/tests/conftest.py b/tests/conftest.py index 51c626807..c5f7c6a7d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -192,7 +192,7 @@ def _is_smoke_item(item): "tests/test_add.py": _is_smoke_add_case, "tests/test_cast.py": _is_smoke_cast_case, "tests/test_cat.py": _is_smoke_cat_case, - "tests/test_internal_causal_softmax.py": _is_smoke_internal_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, @@ -396,7 +396,7 @@ def _is_smoke_silu_and_mul_case(params): ) -def _is_smoke_internal_causal_softmax_case(params): +def _is_smoke_causal_softmax_case(params): cases = { ((3, 3), None, None), ((32, 512), None, None), diff --git a/tests/test_internal_causal_softmax.py b/tests/test_causal_softmax_infinilm.py similarity index 89% rename from tests/test_internal_causal_softmax.py rename to tests/test_causal_softmax_infinilm.py index 33b110f77..15cfc0d39 100644 --- a/tests/test_internal_causal_softmax.py +++ b/tests/test_causal_softmax_infinilm.py @@ -25,14 +25,14 @@ (torch.bfloat16, 1e-2, 1e-2), ), ) -def test_internal_causal_softmax( +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( - _internal_causal_softmax, + _causal_softmax_infinilm, _torch_causal_softmax, (input_tensor, out), {}, @@ -41,8 +41,8 @@ def test_internal_causal_softmax( ) -def _internal_causal_softmax(input, out): - infini.ops.internal_causal_softmax(input, out, stream=get_stream(input.device)) +def _causal_softmax_infinilm(input, out): + infini.ops.causal_softmax_infinilm(input, out, stream=get_stream(input.device)) return out diff --git a/tests/test_internal_scaled_softmax.py b/tests/test_scaled_softmax_infinilm.py similarity index 88% rename from tests/test_internal_scaled_softmax.py rename to tests/test_scaled_softmax_infinilm.py index 5f88d0b93..cfb19700b 100644 --- a/tests/test_internal_scaled_softmax.py +++ b/tests/test_scaled_softmax_infinilm.py @@ -23,7 +23,7 @@ (torch.bfloat16, 1e-2, 1e-2), ), ) -def test_internal_scaled_softmax( +def test_scaled_softmax_infinilm( shape, scale, dtype, @@ -36,7 +36,7 @@ def test_internal_scaled_softmax( out = empty_strided(shape, None, dtype=dtype, device=device) return Payload( - _internal_scaled_softmax, + _scaled_softmax_infinilm, _torch_scaled_softmax, (input_tensor, out), {"scale": scale, "implementation_index": implementation_index}, @@ -45,8 +45,8 @@ def test_internal_scaled_softmax( ) -def _internal_scaled_softmax(input_tensor, out, *, scale, implementation_index): - infini.ops.internal_scaled_softmax( +def _scaled_softmax_infinilm(input_tensor, out, *, scale, implementation_index): + infini.ops.scaled_softmax_infinilm( input_tensor, scale, out, diff --git a/tests/test_internal_top_k_top_p_sample.py b/tests/test_top_k_top_p_sample_infinilm.py similarity index 84% rename from tests/test_internal_top_k_top_p_sample.py rename to tests/test_top_k_top_p_sample_infinilm.py index 3f151bf1d..fb8d66ba7 100644 --- a/tests/test_internal_top_k_top_p_sample.py +++ b/tests/test_top_k_top_p_sample_infinilm.py @@ -6,7 +6,7 @@ @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_internal_top_k_top_p_sample_reproducible( +def test_top_k_top_p_sample_infinilm_reproducible( dtype, device, implementation_index, @@ -16,13 +16,13 @@ def test_internal_top_k_top_p_sample_reproducible( second = torch.empty_like(first) different_seed = torch.empty_like(first) - _internal_top_k_top_p_sample( + _top_k_top_p_sample_infinilm( logits, None, None, 1234, 9, first, implementation_index ) - _internal_top_k_top_p_sample( + _top_k_top_p_sample_infinilm( logits, None, None, 1234, 9, second, implementation_index ) - _internal_top_k_top_p_sample( + _top_k_top_p_sample_infinilm( logits, None, None, 5678, 9, different_seed, implementation_index ) @@ -41,7 +41,7 @@ def test_internal_top_k_top_p_sample_reproducible( ), ) @pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) -def test_internal_top_k_top_p_sample_filters( +def test_top_k_top_p_sample_infinilm_filters( k_value, p_value, allowed, @@ -65,13 +65,15 @@ def test_internal_top_k_top_p_sample_filters( ) out = torch.empty((32,), dtype=torch.int32, device=device) - _internal_top_k_top_p_sample(logits, k, p, 1234, 0, out, implementation_index) + _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 _internal_top_k_top_p_sample( +def _top_k_top_p_sample_infinilm( logits, k, p, @@ -80,7 +82,7 @@ def _internal_top_k_top_p_sample( out, implementation_index, ): - infini.ops.internal_top_k_top_p_sample( + infini.ops.top_k_top_p_sample_infinilm( logits, k, p, From 5a6b396e81cd229cb54384b6dd0bbddc29e0240e Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 17:12:10 +0800 Subject: [PATCH 18/22] feat(ops): retain deprecated legacy interfaces --- src/CMakeLists.txt | 4 +- src/base/add_rms_norm.h | 105 +++++++ src/base/causal_softmax.h | 54 ++++ src/base/flash_attention.h | 110 +++++++ src/base/scaled_softmax.h | 57 ++++ src/base/swiglu.h | 70 +++++ src/base/top_k_top_p_sampler.h | 79 +++++ src/native/ascend/ops/scaled_softmax/kernel.h | 124 ++++++++ .../ascend/ops/top_k_top_p_sampler/kernel.h | 284 ++++++++++++++++++ .../cpu/ops/causal_softmax/causal_softmax.h | 83 +++++ src/native/cpu/ops/swiglu/swiglu.h | 65 ++++ .../top_k_top_p_sampler/top_k_top_p_sampler.h | 160 ++++++++++ .../cuda/iluvatar/ops/add_rms_norm/kernel.h | 21 ++ .../cuda/iluvatar/ops/causal_softmax/kernel.h | 21 ++ src/native/cuda/iluvatar/ops/swiglu/kernel.h | 21 ++ .../cuda/metax/ops/add_rms_norm/kernel.h | 21 ++ .../cuda/metax/ops/causal_softmax/kernel.h | 21 ++ src/native/cuda/metax/ops/swiglu/kernel.h | 21 ++ .../cuda/moore/ops/add_rms_norm/kernel.h | 25 ++ .../cuda/moore/ops/causal_softmax/kernel.h | 32 ++ src/native/cuda/moore/ops/swiglu/kernel.h | 26 ++ .../cuda/nvidia/ops/add_rms_norm/kernel.h | 21 ++ .../cuda/nvidia/ops/causal_softmax/kernel.h | 21 ++ src/native/cuda/nvidia/ops/swiglu/kernel.h | 21 ++ src/native/cuda/ops/add_rms_norm/kernel.cuh | 80 +++++ src/native/cuda/ops/add_rms_norm/kernel.h | 84 ++++++ src/native/cuda/ops/causal_softmax/kernel.cuh | 106 +++++++ src/native/cuda/ops/causal_softmax/kernel.h | 63 ++++ src/native/cuda/ops/swiglu/kernel.cuh | 87 ++++++ src/native/cuda/ops/swiglu/kernel.h | 106 +++++++ tests/conftest.py | 21 ++ tests/test_add_rms_norm.py | 199 ++++++++++++ tests/test_causal_softmax.py | 55 ++++ tests/test_scaled_softmax.py | 66 ++++ tests/test_swiglu.py | 72 +++++ tests/test_top_k_top_p_sampler.py | 78 +++++ 36 files changed, 2482 insertions(+), 2 deletions(-) create mode 100644 src/base/add_rms_norm.h create mode 100644 src/base/causal_softmax.h create mode 100644 src/base/flash_attention.h create mode 100644 src/base/scaled_softmax.h create mode 100644 src/base/swiglu.h create mode 100644 src/base/top_k_top_p_sampler.h create mode 100644 src/native/ascend/ops/scaled_softmax/kernel.h create mode 100644 src/native/ascend/ops/top_k_top_p_sampler/kernel.h create mode 100644 src/native/cpu/ops/causal_softmax/causal_softmax.h create mode 100644 src/native/cpu/ops/swiglu/swiglu.h create mode 100644 src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h create mode 100644 src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/iluvatar/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/iluvatar/ops/swiglu/kernel.h create mode 100644 src/native/cuda/metax/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/metax/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/metax/ops/swiglu/kernel.h create mode 100644 src/native/cuda/moore/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/moore/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/moore/ops/swiglu/kernel.h create mode 100644 src/native/cuda/nvidia/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/nvidia/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/nvidia/ops/swiglu/kernel.h create mode 100644 src/native/cuda/ops/add_rms_norm/kernel.cuh create mode 100644 src/native/cuda/ops/add_rms_norm/kernel.h create mode 100644 src/native/cuda/ops/causal_softmax/kernel.cuh create mode 100644 src/native/cuda/ops/causal_softmax/kernel.h create mode 100644 src/native/cuda/ops/swiglu/kernel.cuh create mode 100644 src/native/cuda/ops/swiglu/kernel.h create mode 100644 tests/test_add_rms_norm.py create mode 100644 tests/test_causal_softmax.py create mode 100644 tests/test_scaled_softmax.py create mode 100644 tests/test_swiglu.py create mode 100644 tests/test_top_k_top_p_sampler.py diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 0e3567308..ea512b3de 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -419,8 +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 silu_and_mul - causal_softmax_infinilm 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_rms_norm.h b/src/base/add_rms_norm.h new file mode 100644 index 000000000..5c8e4fa6a --- /dev/null +++ b/src/base/add_rms_norm.h @@ -0,0 +1,105 @@ +#ifndef INFINI_OPS_BASE_ADD_RMS_NORM_H_ +#define INFINI_OPS_BASE_ADD_RMS_NORM_H_ + +#include +#include + +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +// 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, + std::optional eps, Tensor out, Tensor residual_out) + : input_shape_{input.shape()}, + out_shape_{out.shape()}, + input_strides_{input.strides()}, + residual_strides_{residual.strides()}, + out_strides_{out.strides()}, + residual_out_strides_{residual_out.strides()}, + eps_{eps.value_or(1e-6f)}, + dim_{out.size(-1)}, + ndim_{out.ndim()}, + batch_size_{ndim_ == 2 ? out.size(-2) : out.size(-3)}, + nhead_{ndim_ == 2 ? 1 : out.size(-2)} { + assert((ndim_ == 2 || ndim_ == 3) && + "`AddRmsNorm` supports 2D or 3D tensors only"); + assert(input.shape() == out.shape() && + "`AddRmsNorm` requires `input` and `out` to have the same shape"); + assert(input.shape() == residual.shape() && + "`AddRmsNorm` requires `input` and `residual` to have the same " + "shape"); + assert(input.shape() == residual_out.shape() && + "`AddRmsNorm` requires `input` and `residual_out` to have the " + "same shape"); + assert(weight.ndim() == 1 && weight.size(-1) == dim_ && + "`AddRmsNorm` requires 1D `weight` with size equal to the " + "normalized dimension"); + assert(input.dtype() == out.dtype() && + "`AddRmsNorm` requires `input` and `out` to have the same dtype"); + assert(input.dtype() == residual.dtype() && + "`AddRmsNorm` requires `input` and `residual` to have the same " + "dtype"); + assert(input.dtype() == residual_out.dtype() && + "`AddRmsNorm` requires `input` and `residual_out` to have the same " + "dtype"); + // The CUDA kernel indexes the normalized dimension with stride 1. + assert(input.stride(-1) == 1 && + "`AddRmsNorm` requires the last dimension of `input` to be " + "contiguous"); + assert(residual.stride(-1) == 1 && + "`AddRmsNorm` requires the last dimension of `residual` to be " + "contiguous"); + assert(out.stride(-1) == 1 && + "`AddRmsNorm` requires the last dimension of `out` to be " + "contiguous"); + assert(residual_out.stride(-1) == 1 && + "`AddRmsNorm` requires the last dimension of `residual_out` to be " + "contiguous"); + assert(weight.stride(-1) == 1 && + "`AddRmsNorm` requires the last dimension of `weight` to be " + "contiguous"); + } + + virtual void operator()(const Tensor input, const Tensor residual, + const Tensor weight, std::optional eps, + Tensor out, Tensor residual_out) const = 0; + + virtual void operator()(const Tensor input, const Tensor residual, + const Tensor weight, Tensor out, + Tensor residual_out) const { + return operator()(input, residual, weight, std::nullopt, out, residual_out); + } + + protected: + Tensor::Shape input_shape_; + + Tensor::Shape out_shape_; + + Tensor::Strides input_strides_; + + Tensor::Strides residual_strides_; + + Tensor::Strides out_strides_; + + Tensor::Strides residual_out_strides_; + + float eps_{1e-6f}; + + Tensor::Size dim_{0}; + + Tensor::Size ndim_{0}; + + Tensor::Size batch_size_{0}; + + Tensor::Size nhead_{1}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/causal_softmax.h b/src/base/causal_softmax.h new file mode 100644 index 000000000..e03cfb2df --- /dev/null +++ b/src/base/causal_softmax.h @@ -0,0 +1,54 @@ +#ifndef INFINI_OPS_BASE_CAUSAL_SOFTMAX_H_ +#define INFINI_OPS_BASE_CAUSAL_SOFTMAX_H_ + +#include +#include + +#include "operator.h" +#include "tensor.h" + +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) + : 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() && + "`CausalSoftmax` requires `input` and `out` same shape"); + assert(input.dtype() == out.dtype() && + "`CausalSoftmax` requires `input` and `out` same dtype"); + assert((ndim_ == 2 || ndim_ == 3) && + "`CausalSoftmax` requires 2D or 3D tensor"); + assert(seq_len_ <= total_seq_len_ && + "`CausalSoftmax` 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/flash_attention.h b/src/base/flash_attention.h new file mode 100644 index 000000000..9674ea789 --- /dev/null +++ b/src/base/flash_attention.h @@ -0,0 +1,110 @@ +#ifndef INFINI_OPS_BASE_FLASH_ATTENTION_H_ +#define INFINI_OPS_BASE_FLASH_ATTENTION_H_ + +#include +#include +#include + +#include "operator.h" + +namespace infini::ops { + +// 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, + std::optional cu_seqlens_q, + std::optional cu_seqlens_kv, + std::optional block_table, int64_t num_heads, + int64_t num_kv_heads, int64_t head_size, double scale, + bool causal, int64_t window_left, int64_t window_right, + int64_t block_size, Tensor output) + : num_tokens_{query.size(0)}, + num_heads_{num_heads}, + num_kv_heads_{num_kv_heads}, + head_size_{head_size}, + scale_{scale}, + causal_{causal}, + window_left_{window_left}, + window_right_{window_right}, + block_size_{block_size}, + dtype_{query.dtype()}, + query_shape_{query.shape()}, + key_shape_{key.shape()}, + value_shape_{value.shape()}, + output_shape_{output.shape()}, + query_strides_{query.strides()}, + key_strides_{key.strides()}, + value_strides_{value.strides()}, + output_strides_{output.strides()}, + has_cu_seqlens_q_{cu_seqlens_q.has_value()}, + has_cu_seqlens_kv_{cu_seqlens_kv.has_value()}, + has_block_table_{block_table.has_value()} { + assert(num_heads % num_kv_heads == 0 && + "`FlashAttention` requires num_heads divisible by num_kv_heads"); + assert(query.ndim() == 3 && + "`FlashAttention` requires query to be 3D [T, N, D]"); + } + + virtual void operator()(const Tensor query, const Tensor key, + const Tensor value, + std::optional cu_seqlens_q, + std::optional cu_seqlens_kv, + std::optional block_table, int64_t num_heads, + int64_t num_kv_heads, int64_t head_size, double scale, + bool causal, int64_t window_left, + int64_t window_right, int64_t block_size, + Tensor output) const = 0; + + protected: + Tensor::Size num_tokens_{0}; + + int64_t num_heads_{0}; + + int64_t num_kv_heads_{0}; + + int64_t head_size_{0}; + + double scale_{0.0}; + + bool causal_{false}; + + int64_t window_left_{-1}; + + int64_t window_right_{-1}; + + int64_t block_size_{0}; + + const DataType dtype_; + + Tensor::Shape query_shape_; + + Tensor::Shape key_shape_; + + Tensor::Shape value_shape_; + + Tensor::Shape output_shape_; + + Tensor::Strides query_strides_; + + Tensor::Strides key_strides_; + + Tensor::Strides value_strides_; + + Tensor::Strides output_strides_; + + bool has_cu_seqlens_q_{false}; + + bool has_cu_seqlens_kv_{false}; + + bool has_block_table_{false}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/scaled_softmax.h b/src/base/scaled_softmax.h new file mode 100644 index 000000000..671d052ee --- /dev/null +++ b/src/base/scaled_softmax.h @@ -0,0 +1,57 @@ +#ifndef INFINI_OPS_BASE_SCALED_SOFTMAX_H_ +#define INFINI_OPS_BASE_SCALED_SOFTMAX_H_ + +#include +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +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) + : 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 && + "`ScaledSoftmax` currently supports 2D `[batch, vocab]` input"); + assert(input.shape() == out.shape() && + "`ScaledSoftmax` requires `input` and `out` to have the same shape"); + assert(input.dtype() == out.dtype() && + "`ScaledSoftmax` requires `input` and `out` to have the same dtype"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && + "`ScaledSoftmax` requires a floating point dtype"); + assert(std::isfinite(scale_) && + "`ScaledSoftmax` 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_H_ diff --git a/src/base/swiglu.h b/src/base/swiglu.h new file mode 100644 index 000000000..145748ced --- /dev/null +++ b/src/base/swiglu.h @@ -0,0 +1,70 @@ +#ifndef INFINI_OPS_BASE_SWIGLU_H_ +#define INFINI_OPS_BASE_SWIGLU_H_ + +#include + +#include "operator.h" + +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) + : ndim_{out.ndim()}, + output_size_{out.numel()}, + input_type_{input.dtype()}, + gate_type_{gate.dtype()}, + out_type_{out.dtype()}, + input_shape_{input.shape()}, + gate_shape_{gate.shape()}, + out_shape_{out.shape()}, + input_strides_{input.strides()}, + gate_strides_{gate.strides()}, + out_strides_{out.strides()}, + is_input_contiguous_{input.IsContiguous()}, + is_gate_contiguous_{gate.IsContiguous()}, + is_out_contiguous_{out.IsContiguous()} { + assert( + input_type_ == gate_type_ && gate_type_ == out_type_ && + "operator `Swiglu` requires all input and output tensors to have the " + "same dtype"); + } + + virtual void operator()(const Tensor input, const Tensor gate, + Tensor out) const = 0; + + protected: + Tensor::Size ndim_{0}; + + Tensor::Size output_size_{0}; + + const DataType input_type_; + + const DataType gate_type_; + + const DataType out_type_; + + Tensor::Shape input_shape_; + + Tensor::Shape gate_shape_; + + Tensor::Shape out_shape_; + + Tensor::Strides input_strides_; + + Tensor::Strides gate_strides_; + + Tensor::Strides out_strides_; + + bool is_input_contiguous_{false}; + + bool is_gate_contiguous_{false}; + + bool is_out_contiguous_{false}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/base/top_k_top_p_sampler.h b/src/base/top_k_top_p_sampler.h new file mode 100644 index 000000000..55e4fa0bb --- /dev/null +++ b/src/base/top_k_top_p_sampler.h @@ -0,0 +1,79 @@ +#ifndef INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLER_H_ +#define INFINI_OPS_BASE_TOP_K_TOP_P_SAMPLER_H_ + +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +// 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, + std::optional p, Tensor out) + : batch_size_{logits.size(0)}, + vocab_size_{logits.size(1)}, + dtype_{logits.dtype()} { + assert(logits.ndim() == 2 && + "`TopKTopPSampler` requires 2D `[batch_size, vocab_size]` logits"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16 || + dtype_ == DataType::kFloat32 || dtype_ == DataType::kFloat64) && + "`TopKTopPSampler` requires floating-point logits"); + assert(out.ndim() == 1 && + "`TopKTopPSampler` requires 1D `[batch_size]` output"); + assert(out.size(0) == batch_size_ && + "`TopKTopPSampler` requires output batch size to match logits"); + assert(out.dtype() == DataType::kInt32 && + "`TopKTopPSampler` requires int32 output"); + + ValidateK(k); + ValidateP(p); + } + + virtual void operator()(const Tensor logits, std::optional k, + std::optional p, Tensor out) const = 0; + + protected: + void ValidateK(std::optional k) const { + if (!k.has_value()) return; + + assert(k->ndim() == 1 && + "`TopKTopPSampler` requires `k` to be 1D when provided"); + assert((k->size(0) == 1 || k->size(0) == batch_size_) && + "`TopKTopPSampler` requires `k` shape [1] or [batch_size]"); + assert((k->dtype() == DataType::kInt32 || k->dtype() == DataType::kInt64) && + "`TopKTopPSampler` requires int32 or int64 `k`"); + } + + void ValidateP(std::optional p) const { + if (!p.has_value()) return; + + assert(p->ndim() == 1 && + "`TopKTopPSampler` requires `p` to be 1D when provided"); + assert((p->size(0) == 1 || p->size(0) == batch_size_) && + "`TopKTopPSampler` requires `p` shape [1] or [batch_size]"); + assert((p->dtype() == DataType::kFloat16 || + p->dtype() == DataType::kBFloat16 || + p->dtype() == DataType::kFloat32 || + p->dtype() == DataType::kFloat64) && + "`TopKTopPSampler` 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_SAMPLER_H_ diff --git a/src/native/ascend/ops/scaled_softmax/kernel.h b/src/native/ascend/ops/scaled_softmax/kernel.h new file mode 100644 index 000000000..c6c0df6f7 --- /dev/null +++ b/src/native/ascend/ops/scaled_softmax/kernel.h @@ -0,0 +1,124 @@ +#ifndef INFINI_OPS_ASCEND_SCALED_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_ASCEND_SCALED_SOFTMAX_KERNEL_H_ + +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_mul.h" +#include "aclnn_softmax.h" +#include "base/scaled_softmax.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 ScaledSoftmax { + public: + Operator(const Tensor input, double scale, Tensor out) + : ScaledSoftmax(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) && + "`ScaledSoftmax` Ascend path requires float16, bfloat16, or " + "float32 input"); + assert(input.IsContiguous() && + "`ScaledSoftmax` Ascend path requires contiguous input"); + assert(out.IsContiguous() && + "`ScaledSoftmax` 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_ && + "`ScaledSoftmax` 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_KERNEL_H_ diff --git a/src/native/ascend/ops/top_k_top_p_sampler/kernel.h b/src/native/ascend/ops/top_k_top_p_sampler/kernel.h new file mode 100644 index 000000000..6da146118 --- /dev/null +++ b/src/native/ascend/ops/top_k_top_p_sampler/kernel.h @@ -0,0 +1,284 @@ +#ifndef INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLER_KERNEL_H_ +#define INFINI_OPS_ASCEND_TOP_K_TOP_P_SAMPLER_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_sampler.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 TopKTopPSampler { + public: + Operator(const Tensor logits, std::optional k, + std::optional p, Tensor out) + : TopKTopPSampler(logits, k, p, out) { + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16) && + "`TopKTopPSampler` Ascend ACLNN path requires float16 or bfloat16 " + "logits"); + assert(logits.IsContiguous() && + "`TopKTopPSampler` Ascend ACLNN path requires contiguous logits"); + assert(out.IsContiguous() && + "`TopKTopPSampler` 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, Tensor out) const override { + assert(logits.IsContiguous() && + "`TopKTopPSampler` Ascend ACLNN path requires contiguous logits"); + assert(out.IsContiguous() && + "`TopKTopPSampler` 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); + + 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 && + "`TopKTopPSampler`: 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 && + "`TopKTopPSampler`: 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 && + "`TopKTopPSampler`: 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 && + "`TopKTopPSampler` Ascend path currently requires host-side " + "`k`/`p` tensors"); + assert(tensor->IsContiguous() && + "`TopKTopPSampler` 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) 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); + + 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 && "`TopKTopPSampler` 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 std::mt19937 rng_{std::random_device{}()}; + + 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_SAMPLER_KERNEL_H_ diff --git a/src/native/cpu/ops/causal_softmax/causal_softmax.h b/src/native/cpu/ops/causal_softmax/causal_softmax.h new file mode 100644 index 000000000..8b00c6535 --- /dev/null +++ b/src/native/cpu/ops/causal_softmax/causal_softmax.h @@ -0,0 +1,83 @@ +#ifndef INFINI_OPS_CPU_CAUSAL_SOFTMAX_H_ +#define INFINI_OPS_CPU_CAUSAL_SOFTMAX_H_ + +#include + +#include "base/causal_softmax.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 CausalSoftmax, + Caster { + public: + Operator(const Tensor input, Tensor out) : CausalSoftmax{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/swiglu/swiglu.h b/src/native/cpu/ops/swiglu/swiglu.h new file mode 100644 index 000000000..37ae46767 --- /dev/null +++ b/src/native/cpu/ops/swiglu/swiglu.h @@ -0,0 +1,65 @@ +#ifndef INFINI_OPS_CPU_SWIGLU_SWIGLU_H_ +#define INFINI_OPS_CPU_SWIGLU_SWIGLU_H_ + +#include + +#include "base/swiglu.h" +#include "common/generic_utils.h" +#include "native/cpu/caster_.h" + +namespace infini::ops { + +template <> +class Operator : public Swiglu, + Caster { + public: + using Swiglu::Swiglu; + + void operator()(const Tensor input, const Tensor gate, + Tensor out) const override { + DispatchFunc( + out_type_, + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(input, gate, out); + }, + "Operator::operator()"); + } + + private: + template + void Compute(const Tensor input, const Tensor gate, Tensor out) const { + using ComputeType = std::conditional_t || + IsFP16, + float, T>; + + const auto* input_ptr = static_cast(input.data()); + const auto* gate_ptr = static_cast(gate.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) { + auto input_idx = get_idx(i, is_input_contiguous_, input_shape_.data(), + input_strides_.data()); + auto gate_idx = get_idx(i, is_gate_contiguous_, gate_shape_.data(), + gate_strides_.data()); + auto out_idx = get_idx(i, is_out_contiguous_, out_shape_.data(), + out_strides_.data()); + const ComputeType gate_val = Cast(gate_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(Cast(input_ptr[input_idx]) * swish_gate); + } + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h b/src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h new file mode 100644 index 000000000..b0246fad2 --- /dev/null +++ b/src/native/cpu/ops/top_k_top_p_sampler/top_k_top_p_sampler.h @@ -0,0 +1,160 @@ +#ifndef INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ +#define INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "base/top_k_top_p_sampler.h" +#include "data_type.h" +#include "native/cpu/caster_.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public TopKTopPSampler, Caster { + public: + Operator(const Tensor logits, std::optional k, + std::optional p, Tensor out) + : TopKTopPSampler(logits, k, p, out) {} + + void operator()(const Tensor logits, std::optional k, + std::optional p, Tensor out) const override { + DispatchFunc( + logits.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(logits, k, p, out); + }, + "`Operator::operator()`"); + } + + private: + template + void Compute(const Tensor logits, std::optional k, + std::optional p, 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); + } + } + + template + int32_t SampleRow(const T* row, Tensor::Stride stride, int64_t top_k, + double top_p) 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()); + 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 && "`TopKTopPSampler` has unsupported `p` dtype."); + return 1.0; + } + } + + mutable std::mt19937 rng_{std::random_device{}()}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CPU_TOP_K_TOP_P_SAMPLER_H_ diff --git a/src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h b/src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h new file mode 100644 index 000000000..828b93bbc --- /dev/null +++ b/src/native/cuda/iluvatar/ops/add_rms_norm/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_ILUVATAR_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaAddRmsNorm> { + public: + using CudaAddRmsNorm>::CudaAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h b/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..8c35c87ce --- /dev/null +++ b/src/native/cuda/iluvatar/ops/causal_softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmax> { + public: + using CudaCausalSoftmax>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/iluvatar/ops/swiglu/kernel.h b/src/native/cuda/iluvatar/ops/swiglu/kernel.h new file mode 100644 index 000000000..b03ef19c9 --- /dev/null +++ b/src/native/cuda/iluvatar/ops/swiglu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_ILUVATAR_SWIGLU_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_SWIGLU_KERNEL_H_ + +#include + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/swiglu/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSwiglu> { + public: + using CudaSwiglu>::CudaSwiglu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/add_rms_norm/kernel.h b/src/native/cuda/metax/ops/add_rms_norm/kernel.h new file mode 100644 index 000000000..564ceba61 --- /dev/null +++ b/src/native/cuda/metax/ops/add_rms_norm/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_METAX_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaAddRmsNorm> { + public: + using CudaAddRmsNorm>::CudaAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/causal_softmax/kernel.h b/src/native/cuda/metax/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..426465984 --- /dev/null +++ b/src/native/cuda/metax/ops/causal_softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_METAX_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmax> { + public: + using CudaCausalSoftmax>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/swiglu/kernel.h b/src/native/cuda/metax/ops/swiglu/kernel.h new file mode 100644 index 000000000..0bcfd3198 --- /dev/null +++ b/src/native/cuda/metax/ops/swiglu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_SWIGLU_KERNEL_H_ +#define INFINI_OPS_METAX_SWIGLU_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/swiglu/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSwiglu> { + public: + using CudaSwiglu>::CudaSwiglu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/add_rms_norm/kernel.h b/src/native/cuda/moore/ops/add_rms_norm/kernel.h new file mode 100644 index 000000000..8d3f5cc79 --- /dev/null +++ b/src/native/cuda/moore/ops/add_rms_norm/kernel.h @@ -0,0 +1,25 @@ +#ifndef INFINI_OPS_MOORE_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_MOORE_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/add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaAddRmsNorm> { + public: + using CudaAddRmsNorm>::CudaAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/causal_softmax/kernel.h b/src/native/cuda/moore/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..e13118763 --- /dev/null +++ b/src/native/cuda/moore/ops/causal_softmax/kernel.h @@ -0,0 +1,32 @@ +#ifndef INFINI_OPS_MOORE_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_MOORE_CAUSAL_SOFTMAX_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/kernel.h" + +namespace infini::ops { + +struct MooreCausalSoftmaxBackend : Runtime { + // Moore's causal softmax kernel should not dispatch block sizes above 1024. + static constexpr int max_block_size = 1024; +}; + +template <> +class Operator + : public CudaCausalSoftmax { + public: + using CudaCausalSoftmax::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/moore/ops/swiglu/kernel.h b/src/native/cuda/moore/ops/swiglu/kernel.h new file mode 100644 index 000000000..ac54bff50 --- /dev/null +++ b/src/native/cuda/moore/ops/swiglu/kernel.h @@ -0,0 +1,26 @@ +#ifndef INFINI_OPS_MOORE_SWIGLU_KERNEL_H_ +#define INFINI_OPS_MOORE_SWIGLU_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/swiglu/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSwiglu> { + public: + using CudaSwiglu>::CudaSwiglu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/add_rms_norm/kernel.h b/src/native/cuda/nvidia/ops/add_rms_norm/kernel.h new file mode 100644 index 000000000..2bb6f6051 --- /dev/null +++ b/src/native/cuda/nvidia/ops/add_rms_norm/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_NVIDIA_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_ADD_RMS_NORM_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/add_rms_norm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaAddRmsNorm> { + public: + using CudaAddRmsNorm>::CudaAddRmsNorm; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/causal_softmax/kernel.h b/src/native/cuda/nvidia/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..5a7e18a94 --- /dev/null +++ b/src/native/cuda/nvidia/ops/causal_softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_NVIDIA_CAUSAL_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/causal_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaCausalSoftmax> { + public: + using CudaCausalSoftmax>::CudaCausalSoftmax; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/nvidia/ops/swiglu/kernel.h b/src/native/cuda/nvidia/ops/swiglu/kernel.h new file mode 100644 index 000000000..fd58f29ef --- /dev/null +++ b/src/native/cuda/nvidia/ops/swiglu/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_NVIDIA_SWIGLU_KERNEL_H_ +#define INFINI_OPS_NVIDIA_SWIGLU_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/swiglu/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaSwiglu> { + public: + using CudaSwiglu>::CudaSwiglu; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/add_rms_norm/kernel.cuh b/src/native/cuda/ops/add_rms_norm/kernel.cuh new file mode 100644 index 000000000..7fd8215e1 --- /dev/null +++ b/src/native/cuda/ops/add_rms_norm/kernel.cuh @@ -0,0 +1,80 @@ +#ifndef INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_CUH_ + +#include +#include +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { +namespace 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 add_rms_norm_detail + +template +__global__ void AddRmsNormKernel( + TData* __restrict__ y, int64_t stride_y_batch, int64_t stride_y_nhead, + TData* __restrict__ residual_out, int64_t stride_residual_out_batch, + int64_t stride_residual_out_nhead, const TData* __restrict__ input, + int64_t stride_input_batch, int64_t stride_input_nhead, + const TData* __restrict__ residual, int64_t stride_residual_batch, + int64_t stride_residual_nhead, const TWeight* __restrict__ w, size_t nhead, + size_t dim, float epsilon) { + size_t batch_idx = blockIdx.x / nhead; + size_t head_idx = blockIdx.x % nhead; + + auto y_ptr = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead; + auto input_ptr = + input + batch_idx * stride_input_batch + head_idx * stride_input_nhead; + auto residual_ptr = residual + batch_idx * stride_residual_batch + + head_idx * stride_residual_nhead; + auto w_ptr = w; + auto residual_out_ptr = residual_out + batch_idx * stride_residual_out_batch + + head_idx * stride_residual_out_nhead; + + 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_out_ptr[i] = Caster::template Cast(sum_val); + } + + TCompute sum_squared = + add_rms_norm_detail::SumSquared( + residual_out_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 sum_val = + Caster::template Cast(residual_out_ptr[i]); + y_ptr[i] = Caster::template Cast( + sum_val * Caster::template Cast(w_ptr[i]) * rms); + } +} + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/add_rms_norm/kernel.h b/src/native/cuda/ops/add_rms_norm/kernel.h new file mode 100644 index 000000000..a2c0c7cbf --- /dev/null +++ b/src/native/cuda/ops/add_rms_norm/kernel.h @@ -0,0 +1,84 @@ +#ifndef INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_H_ +#define INFINI_OPS_CUDA_ADD_RMS_NORM_KERNEL_H_ + +#include +#include +#include + +#include "base/add_rms_norm.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/add_rms_norm/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaAddRmsNorm : public AddRmsNorm { + public: + using AddRmsNorm::AddRmsNorm; + + void operator()(const Tensor input, const Tensor residual, + const Tensor weight, std::optional eps, Tensor out, + Tensor residual_out) const override { + const float kernel_eps = eps.value_or(eps_); + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + + auto stride_input_batch = input_strides_.size() > 1 ? input_strides_[0] : 0; + auto stride_input_nhead = + input_strides_.size() > 1 ? input_strides_[1] : input_strides_[0]; + auto stride_residual_batch = + residual_strides_.size() > 1 ? residual_strides_[0] : 0; + auto stride_residual_nhead = residual_strides_.size() > 1 + ? residual_strides_[1] + : residual_strides_[0]; + auto stride_out_batch = out_strides_.size() > 1 ? out_strides_[0] : 0; + auto stride_out_nhead = + out_strides_.size() > 1 ? out_strides_[1] : out_strides_[0]; + auto stride_residual_out_batch = + residual_out_strides_.size() > 1 ? residual_out_strides_[0] : 0; + auto stride_residual_out_nhead = residual_out_strides_.size() > 1 + ? residual_out_strides_[1] + : residual_out_strides_[0]; + + uint32_t num_blocks = static_cast(batch_size_ * nhead_); + + assert(out.dtype() == input.dtype() && out.dtype() == residual.dtype() && + out.dtype() == residual_out.dtype() && + "`CudaAddRmsNorm` requires input, residual, out and residual_out to " + "have the same dtype"); + + int block_size = RuntimeUtils::GetOptimalBlockSize(); + + DispatchFunc, ReducedFloatTypes>, + ConcatType, ReducedFloatTypes>, + AllCudaBlockSizes>( + {static_cast(out.dtype()), + static_cast(weight.dtype()), block_size}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + using TWeight = + TypeMapType(list_tag)>; + constexpr int kBlockSize = ListGet<2>(list_tag); + + AddRmsNormKernel + <<>>( + reinterpret_cast(out.data()), stride_out_batch, + stride_out_nhead, reinterpret_cast(residual_out.data()), + stride_residual_out_batch, stride_residual_out_nhead, + reinterpret_cast(input.data()), stride_input_batch, + stride_input_nhead, + reinterpret_cast(residual.data()), + stride_residual_batch, stride_residual_nhead, + reinterpret_cast(weight.data()), nhead_, dim_, + kernel_eps); + }, + "CudaAddRmsNorm::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/causal_softmax/kernel.cuh b/src/native/cuda/ops/causal_softmax/kernel.cuh new file mode 100644 index 000000000..31f685e84 --- /dev/null +++ b/src/native/cuda/ops/causal_softmax/kernel.cuh @@ -0,0 +1,106 @@ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_CUH_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_CUH_ + +#include +#include +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { + +namespace { + +template +__device__ __forceinline__ Data ExpAndCast(Compute x) { + return Caster::template Cast( + expf(Caster::template Cast(x))); +} + +struct BlockMaxOp { + template + __device__ __forceinline__ T operator()(const T& a, const T& b) const { + return (a > b) ? a : b; + } +}; + +template +__device__ __forceinline__ Data BlockMax(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, BlockMaxOp()); +} + +template +__device__ __forceinline__ Compute BlockSum(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 CausalSoftmaxKernel( + 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 = BlockMax(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] = ExpAndCast(diff); + } else { + out_row[col] = Caster::template Cast(0.0f); + } + } + __syncthreads(); + + __shared__ Compute sum_val; + Compute block_sum = + BlockSum(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/kernel.h b/src/native/cuda/ops/causal_softmax/kernel.h new file mode 100644 index 000000000..c77d97387 --- /dev/null +++ b/src/native/cuda/ops/causal_softmax/kernel.h @@ -0,0 +1,63 @@ +#ifndef INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_CUDA_CAUSAL_SOFTMAX_KERNEL_H_ + +#include +#include +#include + +#include "base/causal_softmax.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/causal_softmax/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaCausalSoftmax : public CausalSoftmax { + public: + using CausalSoftmax::CausalSoftmax; + + 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); + + CausalSoftmaxKernel + <<>>( + 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); + }, + "CudaCausalSoftmax::operator()"); + } +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/ops/swiglu/kernel.cuh b/src/native/cuda/ops/swiglu/kernel.cuh new file mode 100644 index 000000000..da4bfb053 --- /dev/null +++ b/src/native/cuda/ops/swiglu/kernel.cuh @@ -0,0 +1,87 @@ +#ifndef INFINI_OPS_CUDA_SWIGLU_KERNEL_CUH_ +#define INFINI_OPS_CUDA_SWIGLU_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 + +// SwiGLU(x, gate) = Swish(x) * gate = (x * sigmoid(x)) * gate. +template +__global__ void SwigluKernel(T* __restrict__ out, const T* __restrict__ a, + const T* __restrict__ b, + const size_t* __restrict__ out_shape, + const size_t* __restrict__ input_shape, + const size_t* __restrict__ gate_shape, + const ptrdiff_t* __restrict__ out_strides, + const ptrdiff_t* __restrict__ input_strides, + const ptrdiff_t* __restrict__ gate_strides, + size_t output_size, size_t ndim, + bool out_contiguous, bool input_contiguous, + bool gate_contiguous) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + + if (idx < output_size) { + size_t out_idx, input_idx, gate_idx; + + if (out_contiguous) { + out_idx = idx; + } else { + out_idx = IndexToOffset(idx, ndim, out_shape, out_strides); + } + + if (input_contiguous) { + input_idx = idx; + } else { + input_idx = IndexToOffset(idx, ndim, input_shape, input_strides); + } + + if (gate_contiguous) { + gate_idx = idx; + } else { + gate_idx = IndexToOffset(idx, ndim, gate_shape, gate_strides); + } + + T up = a[input_idx]; + T gate = b[gate_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/swiglu/kernel.h b/src/native/cuda/ops/swiglu/kernel.h new file mode 100644 index 000000000..486617f43 --- /dev/null +++ b/src/native/cuda/ops/swiglu/kernel.h @@ -0,0 +1,106 @@ +#ifndef INFINI_OPS_CUDA_SWIGLU_KERNEL_H_ +#define INFINI_OPS_CUDA_SWIGLU_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/swiglu.h" +#include "common/generic_utils.h" +#include "native/cuda/ops/swiglu/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +template +class CudaSwiglu : public Swiglu { + public: + CudaSwiglu(const Tensor input, const Tensor gate, Tensor out) + : Swiglu{input, gate, 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); + 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_gate_shape_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, gate_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_gate_strides_ = reinterpret_cast(d_metadata_ + offset); + std::memcpy(metadata.data() + offset, gate_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); + } + + ~CudaSwiglu() { Backend::Free(d_metadata_); } + + void operator()(const Tensor input, const Tensor gate, + 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()); + const T* d_gate = reinterpret_cast(gate.data()); + + SwigluKernel + <<>>( + d_out, d_input, d_gate, d_out_shape_, d_input_shape_, + d_gate_shape_, d_out_strides_, d_input_strides_, + d_gate_strides_, output_size_, ndim_, is_out_contiguous_, + is_input_contiguous_, is_gate_contiguous_); + }, + "CudaSwiglu::operator()"); + } + + private: + std::byte* d_metadata_{nullptr}; + + Tensor::Size* d_input_shape_{nullptr}; + + Tensor::Size* d_gate_shape_{nullptr}; + + Tensor::Size* d_out_shape_{nullptr}; + + Tensor::Stride* d_input_strides_{nullptr}; + + Tensor::Stride* d_gate_strides_{nullptr}; + + Tensor::Stride* d_out_strides_{nullptr}; +}; + +} // namespace infini::ops + +#endif diff --git a/tests/conftest.py b/tests/conftest.py index c5f7c6a7d..c1552d93a 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -192,6 +192,7 @@ def _is_smoke_item(item): "tests/test_add.py": _is_smoke_add_case, "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, @@ -199,6 +200,7 @@ def _is_smoke_item(item): "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, } matcher = matchers.get(module_path) @@ -396,6 +398,25 @@ def _is_smoke_silu_and_mul_case(params): ) +def _is_smoke_swiglu_case(params): + cases = { + ((13, 4), None, None, None), + ((16, 5632), None, None, None), + } + + return ( + _is_float32(params) + and _shape_case( + params, + "shape", + "input_strides", + "gate_strides", + "out_strides", + ) + in cases + ) + + def _is_smoke_causal_softmax_case(params): cases = { ((3, 3), None, None), diff --git a/tests/test_add_rms_norm.py b/tests/test_add_rms_norm.py new file mode 100644 index 000000000..1d94f7455 --- /dev/null +++ b/tests/test_add_rms_norm.py @@ -0,0 +1,199 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randn_strided + + +# Format: (input_shape, weight_shape, input_strides, residual_strides, weight_strides, out_strides); +# input/residual/residual_out share input_shape. +_TEST_CASES = ( + ((1, 4), (4,), None, None, None, None), + ((2, 4), (4,), None, None, None, None), + ((2, 2, 4), (4,), None, None, None, None), + ((2, 2, 4), (4,), (12, 8, 1), (12, 8, 1), None, (12, 8, 1)), + ((16, 2048), (2048,), None, None, None, None), + ((16, 2048), (2048,), (4096, 1), (4096, 1), None, (4096, 1)), + ((15, 3584), (3584,), None, None, None, None), + ((4, 4, 2048), (2048,), None, None, None, None), + ((4, 4, 2048), (2048,), (2048, 8192, 1), (2048, 8192, 1), None, (2048, 8192, 1)), + ( + (4, 4, 2048), + (2048,), + (16384, 4096, 1), + (16384, 4096, 1), + None, + (16384, 4096, 1), + ), +) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("check_output", ("out", "residual_out")) +@pytest.mark.parametrize( + "input_shape, weight_shape, input_strides, residual_strides, weight_strides, out_strides", + _TEST_CASES, +) +@pytest.mark.parametrize("eps", (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_add_rms_norm( + check_output, + input_shape, + weight_shape, + input_strides, + residual_strides, + weight_strides, + out_strides, + eps, + implementation_index, + dtype, + device, + rtol, + atol, +): + input = randn_strided(input_shape, input_strides, dtype=dtype, device=device) + residual = randn_strided(input_shape, residual_strides, dtype=dtype, device=device) + weight = randn_strided(weight_shape, weight_strides, dtype=dtype, device=device) + residual_out = empty_strided(input_shape, input_strides, dtype=dtype, device=device) + out = empty_strided(input_shape, out_strides, dtype=dtype, device=device) + + if check_output == "out": + func = _add_rms_norm_out + ref = _torch_add_rms_norm_out + else: + func = _add_rms_norm_residual_out + ref = _torch_add_rms_norm_residual_out + + return Payload( + lambda *args, **kwargs: func( + *args, **kwargs, implementation_index=implementation_index + ), + ref, + (input, residual, weight), + {"eps": eps, "out": out, "residual_out": residual_out}, + rtol=rtol, + atol=atol, + ) + + +def _add_rms_norm( + input, + residual, + weight, + *, + eps=1e-6, + out=None, + residual_out=None, + implementation_index=0, +): + infini.ops.add_rms_norm( + input, + residual, + weight, + eps, + out, + residual_out, + implementation_index=implementation_index, + stream=get_stream(input.device), + ) + + +def _add_rms_norm_out( + input, + residual, + weight, + *, + eps=1e-6, + out=None, + residual_out=None, + implementation_index=0, +): + _add_rms_norm( + input, + residual, + weight, + eps=eps, + out=out, + residual_out=residual_out, + implementation_index=implementation_index, + ) + + return out + + +def _add_rms_norm_residual_out( + input, + residual, + weight, + *, + eps=1e-6, + out=None, + residual_out=None, + implementation_index=0, +): + _add_rms_norm( + input, + residual, + weight, + eps=eps, + out=out, + residual_out=residual_out, + implementation_index=implementation_index, + ) + + return residual_out + + +def _torch_add_rms_norm( + input, residual, weight, *, eps=1e-6, out=None, residual_out=None +): + """Reference aligned with vLLM `fused_add_rms_norm` (ignoring `variance_size`).""" + orig_dtype = input.dtype + x = input.to(torch.float32) + x = x + residual.to(torch.float32) + add_result = x.to(orig_dtype) + + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + if weight is not None: + x = x.to(weight.dtype) * weight + normalized_result = x.to(orig_dtype) + + if out is not None: + out.copy_(normalized_result) + else: + out = normalized_result + + if residual_out is not None: + residual_out.copy_(add_result) + else: + residual_out = add_result + + return out, residual_out + + +def _torch_add_rms_norm_out( + input, residual, weight, *, eps=1e-6, out=None, residual_out=None +): + out, _ = _torch_add_rms_norm( + input, residual, weight, eps=eps, out=out, residual_out=residual_out + ) + + return out + + +def _torch_add_rms_norm_residual_out( + input, residual, weight, *, eps=1e-6, out=None, residual_out=None +): + _, residual_out = _torch_add_rms_norm( + input, residual, weight, eps=eps, out=out, residual_out=residual_out + ) + + return residual_out diff --git a/tests/test_causal_softmax.py b/tests/test_causal_softmax.py new file mode 100644 index 000000000..ef48601f7 --- /dev/null +++ b/tests/test_causal_softmax.py @@ -0,0 +1,55 @@ +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(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, + _torch_causal_softmax, + (input_tensor, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _causal_softmax(input, out): + infini.ops.causal_softmax(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_scaled_softmax.py b/tests/test_scaled_softmax.py new file mode 100644 index 000000000..ec860b34b --- /dev/null +++ b/tests/test_scaled_softmax.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( + 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, + _torch_scaled_softmax, + (input_tensor, out), + {"scale": scale, "implementation_index": implementation_index}, + rtol=rtol, + atol=atol, + ) + + +def _scaled_softmax(input_tensor, out, *, scale, implementation_index): + infini.ops.scaled_softmax( + 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_swiglu.py b/tests/test_swiglu.py new file mode 100644 index 000000000..f159742ca --- /dev/null +++ b/tests/test_swiglu.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( + "shape, input_strides, gate_strides, out_strides", + ( + ((13, 4), None, None, None), + ((13, 4), (10, 1), (10, 1), (10, 1)), + ((13, 4, 4), None, None, None), + ((13, 4, 4), (20, 4, 1), (20, 4, 1), (20, 4, 1)), + ((16, 5632), None, None, None), + ((16, 5632), (13312, 1), (13312, 1), (13312, 1)), + ((4, 4, 5632), None, None, None), + ((4, 4, 5632), (45056, 5632, 1), (45056, 5632, 1), (45056, 5632, 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_swiglu( + shape, + input_strides, + gate_strides, + out_strides, + implementation_index, + dtype, + device, + rtol, + atol, +): + input = rand_strided(shape, input_strides, dtype=dtype, device=device) + gate = rand_strided(shape, gate_strides, dtype=dtype, device=device) + out = empty_strided(shape, out_strides, dtype=dtype, device=device) + + return Payload( + lambda *args, **kwargs: _swiglu( + *args, **kwargs, implementation_index=implementation_index + ), + _torch_swiglu, + (input, gate, out), + {}, + rtol=rtol, + atol=atol, + ) + + +def _swiglu(input, gate, out, implementation_index=0): + infini.ops.swiglu( + input, + gate, + out, + implementation_index=implementation_index, + stream=get_stream(input.device), + ) + + return out + + +def _torch_swiglu(input, gate, out): + swish_x = gate * torch.sigmoid(gate) + + return torch.mul(input, swish_x, out=out) diff --git a/tests/test_top_k_top_p_sampler.py b/tests/test_top_k_top_p_sampler.py new file mode 100644 index 000000000..8d898e0e6 --- /dev/null +++ b/tests/test_top_k_top_p_sampler.py @@ -0,0 +1,78 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, get_stream + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("shape", ((1, 8), (3, 16))) +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_top_k_top_p_sampler( + shape, + dtype, + device, + implementation_index, +): + batch_size, vocab_size = shape + logits = torch.full(shape, -10.0, dtype=dtype, device=device) + + for i in range(batch_size): + logits[i, i % vocab_size] = 10.0 + + k = torch.ones((batch_size,), dtype=torch.int64, device="cpu") + p = torch.ones((batch_size,), dtype=torch.float32, device="cpu") + out = torch.empty((batch_size,), dtype=torch.int32, device=device) + + return Payload( + _top_k_top_p_sampler, + _torch_argmax, + (logits, k, p, out), + {"implementation_index": implementation_index}, + ) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("dtype", (torch.float16, torch.bfloat16)) +def test_top_k_top_p_sampler_optional_p( + dtype, + device, + implementation_index, +): + shape = (3, 16) + batch_size, vocab_size = shape + logits = torch.full(shape, -10.0, dtype=dtype, device=device) + + for i in range(batch_size): + logits[i, (i + 1) % vocab_size] = 10.0 + + k = torch.ones((1,), dtype=torch.int64, device="cpu") + out = torch.empty((batch_size,), dtype=torch.int32, device=device) + + return Payload( + _top_k_top_p_sampler, + _torch_argmax, + (logits, k, None, out), + {"implementation_index": implementation_index}, + ) + + +def _top_k_top_p_sampler(logits, k, p, out, *, implementation_index): + infini.ops.top_k_top_p_sampler( + logits, + k, + p, + out, + stream=get_stream(logits.device), + implementation_index=implementation_index, + ) + + return out + + +def _torch_argmax(logits, k, p, out, *, implementation_index): + del k, p, implementation_index + + out.copy_(torch.argmax(logits, dim=-1).to(torch.int32)) + + return out From 5135130e658f3aa8eae8f2b9156ba62ee070debf Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 17:14:22 +0800 Subject: [PATCH 19/22] docs(torch): clarify generated backend allowlist --- scripts/torch_ops.yaml | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/scripts/torch_ops.yaml b/scripts/torch_ops.yaml index 3fc72844c..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 From e3ce1c9c97fae50c7283866d14ea1eb68336acdf Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 17:16:20 +0800 Subject: [PATCH 20/22] style(attention): order access specifiers --- src/base/scaled_dot_product_attention.h | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/base/scaled_dot_product_attention.h b/src/base/scaled_dot_product_attention.h index 2dc37919e..ace5d5cde 100644 --- a/src/base/scaled_dot_product_attention.h +++ b/src/base/scaled_dot_product_attention.h @@ -8,15 +8,6 @@ namespace infini::ops { class ScaledDotProductAttention : public Operator { - protected: - template - static auto ReturnShape(const TensorLike& query, const TensorLike& value) { - auto shape = query.shape(); - shape.back() = value.size(-1); - - return shape; - } - public: ScaledDotProductAttention(const Tensor query, const Tensor key, const Tensor value, @@ -52,7 +43,9 @@ class ScaledDotProductAttention : public Operator { "dtypes"); assert(query.size(-1) == key.size(-1) && key.size(-2) == value.size(-2) && "`ScaledDotProductAttention` input dimensions are incompatible"); - assert(out.shape() == ReturnShape(query, value) && + 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]"); @@ -90,8 +83,9 @@ class ScaledDotProductAttention : public Operator { (void)scale; (void)enable_gqa; - return TensorLike::Empty(ReturnShape(query, value), query.dtype(), - query.device()); + auto out_shape = query.shape(); + out_shape.back() = value.size(-1); + return TensorLike::Empty(out_shape, query.dtype(), query.device()); } protected: From 590506e405cc6d3aa851ac571823cb5c3e9fb004 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 21:41:52 +0800 Subject: [PATCH 21/22] fix(cuda): isolate Infinilm softmax helpers --- .../ops/causal_softmax_infinilm/kernel.cuh | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh index 5a061c662..8ffb079ca 100644 --- a/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.cuh @@ -13,12 +13,12 @@ namespace infini::ops { namespace { template -__device__ __forceinline__ Data ExpAndCast(Compute x) { +__device__ __forceinline__ Data ExpAndCastInfinilm(Compute x) { return Caster::template Cast( expf(Caster::template Cast(x))); } -struct BlockMaxOp { +struct BlockMaxInfinilmOp { template __device__ __forceinline__ T operator()(const T& a, const T& b) const { return (a > b) ? a : b; @@ -26,7 +26,8 @@ struct BlockMaxOp { }; template -__device__ __forceinline__ Data BlockMax(const Data* data_ptr, size_t count) { +__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]; @@ -34,13 +35,13 @@ __device__ __forceinline__ Data BlockMax(const Data* data_ptr, size_t count) { } using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage temp_storage; - return BlockReduce(temp_storage).Reduce(thread_max, BlockMaxOp()); + return BlockReduce(temp_storage).Reduce(thread_max, BlockMaxInfinilmOp()); } template -__device__ __forceinline__ Compute BlockSum(const Data* data_ptr, - size_t count) { +__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]); @@ -70,7 +71,7 @@ __global__ void CausalSoftmaxInfinilmKernel( size_t valid_len = total_seq_len - seq_len + row_idx + 1; __shared__ Data max_val; - Data block_max = BlockMax(input_row, valid_len); + Data block_max = BlockMaxInfinilm(input_row, valid_len); if (threadIdx.x == 0) { max_val = block_max; } @@ -80,7 +81,7 @@ __global__ void CausalSoftmaxInfinilmKernel( if (col < valid_len) { Compute diff = Caster::template Cast(input_row[col]) - Caster::template Cast(max_val); - out_row[col] = ExpAndCast(diff); + out_row[col] = ExpAndCastInfinilm(diff); } else { out_row[col] = Caster::template Cast(0.0f); } @@ -89,7 +90,7 @@ __global__ void CausalSoftmaxInfinilmKernel( __shared__ Compute sum_val; Compute block_sum = - BlockSum(out_row, total_seq_len); + BlockSumInfinilm(out_row, total_seq_len); if (threadIdx.x == 0) { sum_val = block_sum; } From 3cf536b2ec89eb4047dc77092a269bda0b63cdde Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 15 Jul 2026 21:42:36 +0800 Subject: [PATCH 22/22] chore(format): format review follow-ups --- src/native/cuda/ops/causal_softmax_infinilm/kernel.h | 3 ++- tests/test_top_k_top_p_sample_infinilm.py | 4 +--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/native/cuda/ops/causal_softmax_infinilm/kernel.h b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h index d2f6bf882..4e1a97924 100644 --- a/src/native/cuda/ops/causal_softmax_infinilm/kernel.h +++ b/src/native/cuda/ops/causal_softmax_infinilm/kernel.h @@ -47,7 +47,8 @@ class CudaCausalSoftmaxInfinilm : public CausalSoftmaxInfinilm { using T = TypeMapType(list_tag)>; constexpr int kBlockSize = ListGet<1>(list_tag); - CausalSoftmaxInfinilmKernel + CausalSoftmaxInfinilmKernel <<>>( reinterpret_cast(out.data()), reinterpret_cast(input.data()), batch_size_, diff --git a/tests/test_top_k_top_p_sample_infinilm.py b/tests/test_top_k_top_p_sample_infinilm.py index fb8d66ba7..f9a4e2705 100644 --- a/tests/test_top_k_top_p_sample_infinilm.py +++ b/tests/test_top_k_top_p_sample_infinilm.py @@ -65,9 +65,7 @@ def test_top_k_top_p_sample_infinilm_filters( ) out = torch.empty((32,), dtype=torch.int32, device=device) - _top_k_top_p_sample_infinilm( - logits, k, p, 1234, 0, out, implementation_index - ) + _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))