diff --git a/src/base/lightning_attention_infinilm.h b/src/base/lightning_attention_infinilm.h new file mode 100644 index 000000000..3cedb496f --- /dev/null +++ b/src/base/lightning_attention_infinilm.h @@ -0,0 +1,171 @@ +#ifndef INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ +#define INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ + +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +/// Lightning attention with an indexed recurrent-state pool. +/// +/// The operator evaluates the recurrent form of MiniMax-style linear attention +/// with a per-head ALiBi-style decay, one token at a time: +/// +/// ratio[h] = exp(-slope[h]) +/// S = ratio[h] * S + outer(k_t[h], v_t[h]) +/// out_t[h] = q_t[h] @ S +/// +/// The recurrent state of request `b` is read from +/// `initial_state[initial_state_indices[b]]` and the final state is written to +/// `initial_state[final_state_indices[b]]`. The row referenced by +/// `initial_state_indices` is left untouched, so a caller may keep using it. +/// +/// Requests are independent and may execute concurrently, so a row used as the +/// destination of one request must not be the source row of another request in +/// the same call. +/// +/// This operator is InfiniLM-specific and is therefore classified as custom +/// rather than aligned to an open-source operator. The closest public +/// reference implementation is Flash-Linear-Attention's +/// `fused_recurrent_lightning_attn`; `dexp` there is the per-head decay factor +/// `exp(-slope)` used here. +class LightningAttentionInfinilm + : public Operator { + public: + LightningAttentionInfinilm(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) + : data_type_{q.dtype()}, + index_dtype_{initial_state_indices.dtype()}, + batch_size_{q.size(0)}, + seq_len_{q.size(1)}, + num_heads_{q.size(2)}, + head_dim_{q.size(3)}, + state_pool_size_{initial_state.size(0)}, + state_pool_stride_{initial_state.stride(0)}, + state_head_stride_{initial_state.stride(1)}, + state_row_stride_{initial_state.stride(2)}, + state_column_stride_{initial_state.stride(3)}, + q_batch_stride_{q.stride(0)}, + q_seq_stride_{q.stride(1)}, + q_head_stride_{q.stride(2)}, + k_batch_stride_{k.stride(0)}, + k_seq_stride_{k.stride(1)}, + k_head_stride_{k.stride(2)}, + v_batch_stride_{v.stride(0)}, + v_seq_stride_{v.stride(1)}, + v_head_stride_{v.stride(2)}, + out_batch_stride_{out.stride(0)}, + out_seq_stride_{out.stride(1)}, + out_head_stride_{out.stride(2)}, + slope_stride_{slope.stride(0)}, + initial_index_stride_{initial_state_indices.stride(0)}, + final_index_stride_{final_state_indices.stride(0)} { + assert(q.ndim() == 4 && k.ndim() == 4 && v.ndim() == 4 && out.ndim() == 4 && + "`LightningAttentionInfinilm` expects [batch, seq, heads, head_dim] tensors"); + assert(q.dtype() == k.dtype() && k.dtype() == v.dtype() && + v.dtype() == out.dtype() && out.dtype() == data_type_ && + initial_state.dtype() == data_type_ && + "`LightningAttentionInfinilm` requires all data tensors to share one dtype"); + assert((data_type_ == DataType::kFloat32 || + data_type_ == DataType::kFloat16 || + data_type_ == DataType::kBFloat16) && + "`LightningAttentionInfinilm` supports float32, float16 and bfloat16"); + assert(q.shape() == k.shape() && k.shape() == v.shape() && + v.shape() == out.shape() && + "`LightningAttentionInfinilm` requires q, k, v and out to share a shape"); + assert(slope.dtype() == DataType::kFloat32 && slope.ndim() == 1 && + slope.size(0) == num_heads_ && slope_stride_ == 1 && + "`LightningAttentionInfinilm` expects `slope` to be a contiguous float32 tensor of size num_heads"); + assert(initial_state.ndim() == 4 && initial_state.size(1) == num_heads_ && + initial_state.size(2) == head_dim_ && + initial_state.size(3) == head_dim_ && + "`LightningAttentionInfinilm` expects `initial_state` to be [pool, heads, head_dim, head_dim]"); + assert(state_pool_size_ > 0 && state_column_stride_ == 1 && + "`LightningAttentionInfinilm` expects a contiguous state pool on the last dimension"); + assert(initial_state_indices.ndim() == 1 && + final_state_indices.ndim() == 1 && + initial_state_indices.size(0) == batch_size_ && + final_state_indices.size(0) == batch_size_ && + "`LightningAttentionInfinilm` expects one state index per request"); + assert(IsIndexDtype(index_dtype_) && + final_state_indices.dtype() == index_dtype_ && + initial_index_stride_ == 1 && final_index_stride_ == 1 && + "`LightningAttentionInfinilm` expects contiguous int32/int64 state indices"); + assert(q.stride(3) == 1 && k.stride(3) == 1 && v.stride(3) == 1 && + out.stride(3) == 1 && + "`LightningAttentionInfinilm` requires a contiguous last dimension"); + } + + virtual void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, + Tensor out) const = 0; + + protected: + static bool IsIndexDtype(DataType dtype) { + return dtype == DataType::kInt32 || dtype == DataType::kInt64; + } + + DataType data_type_; + + DataType index_dtype_; + + Tensor::Size batch_size_{0}; + + Tensor::Size seq_len_{0}; + + Tensor::Size num_heads_{0}; + + Tensor::Size head_dim_{0}; + + Tensor::Size state_pool_size_{0}; + + Tensor::Stride state_pool_stride_{0}; + + Tensor::Stride state_head_stride_{0}; + + Tensor::Stride state_row_stride_{0}; + + Tensor::Stride state_column_stride_{0}; + + Tensor::Stride q_batch_stride_{0}; + + Tensor::Stride q_seq_stride_{0}; + + Tensor::Stride q_head_stride_{0}; + + Tensor::Stride k_batch_stride_{0}; + + Tensor::Stride k_seq_stride_{0}; + + Tensor::Stride k_head_stride_{0}; + + Tensor::Stride v_batch_stride_{0}; + + Tensor::Stride v_seq_stride_{0}; + + Tensor::Stride v_head_stride_{0}; + + Tensor::Stride out_batch_stride_{0}; + + Tensor::Stride out_seq_stride_{0}; + + Tensor::Stride out_head_stride_{0}; + + Tensor::Stride slope_stride_{0}; + + Tensor::Stride initial_index_stride_{0}; + + Tensor::Stride final_index_stride_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_LIGHTNING_ATTENTION_INFINILM_H_ diff --git a/src/native/ascend/ops/lightning_attention_infinilm/kernel.h b/src/native/ascend/ops/lightning_attention_infinilm/kernel.h new file mode 100644 index 000000000..cd8023e08 --- /dev/null +++ b/src/native/ascend/ops/lightning_attention_infinilm/kernel.h @@ -0,0 +1,389 @@ +#ifndef INFINI_OPS_ASCEND_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ +#define INFINI_OPS_ASCEND_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ + +#include +#include +#include +#include +#include +#include + +#include "acl/acl.h" +#include "aclnn/aclnn_base.h" +#include "aclnn_add.h" +#include "aclnn_mul.h" +#include "aclnnop/aclnn_matmul.h" +#include "base/lightning_attention_infinilm.h" +#include "data_type.h" +#include "native/ascend/common.h" +#include "native/ascend/data_type_.h" +#include "native/ascend/workspace_pool_.h" +#include "operator.h" + +namespace infini::ops { + +// Ascend implementation of the indexed-pool lightning attention. +// +// CANN does not provide a fused lightning-attention operator, so this path +// composes aclnn kernels and walks requests and tokens on the host: +// +// ratio = exp(-slope) (per head, once per call) +// decayed = Mul(state, ratio) (ratio is [H, 1, 1]) +// outer = Matmul(k_t [H, D, 1], v_t [H, 1, D]) +// state = Add(decayed, outer) +// out_t = Matmul(q_t [H, 1, D], state) +// +// Each request is staged in a workspace row. The source state row is copied +// into the workspace first, all tokens are evaluated there, and the result is +// copied back to the destination row last. Therefore the source row is never +// written, even when the source and destination rows are different. +// +// This is a correctness-first fallback. A fused AscendC kernel is the +// follow-up once the composition is validated on hardware. +template <> +class Operator + : public LightningAttentionInfinilm { + public: + Operator(const Tensor q, const Tensor k, const Tensor v, const Tensor slope, + Tensor initial_state, const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) + : LightningAttentionInfinilm(q, k, v, slope, initial_state, + initial_state_indices, final_state_indices, + out), + state_cache_({static_cast(num_heads_), + static_cast(head_dim_), + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr), + decayed_cache_({static_cast(num_heads_), + static_cast(head_dim_), + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr), + outer_cache_({static_cast(num_heads_), + static_cast(head_dim_), + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr), + ratio_cache_({static_cast(num_heads_), 1, 1}, + ascend::ToAclDtype(data_type_), nullptr), + outer_k_cache_({static_cast(num_heads_), + static_cast(head_dim_), 1}, + ascend::ToAclDtype(data_type_), nullptr), + outer_v_cache_({static_cast(num_heads_), 1, + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr), + output_q_cache_({static_cast(num_heads_), 1, + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr), + output_cache_({static_cast(num_heads_), 1, + static_cast(head_dim_)}, + ascend::ToAclDtype(data_type_), nullptr) { + assert((data_type_ == DataType::kFloat16 || + data_type_ == DataType::kBFloat16 || + data_type_ == DataType::kFloat32) && + "`LightningAttentionInfinilm` Ascend path requires float16, " + "bfloat16 or float32"); + assert(q.IsContiguous() && k.IsContiguous() && v.IsContiguous() && + out.IsContiguous() && initial_state.IsContiguous() && + "`LightningAttentionInfinilm` Ascend path requires contiguous " + "tensors"); + assert(num_heads_ > 0 && head_dim_ > 0 && + "`LightningAttentionInfinilm` Ascend path requires non-empty " + "heads and head_dim"); + + add_alpha_storage_ = 1.0f; + add_alpha_ = aclCreateScalar(&add_alpha_storage_, ACL_FLOAT); + assert(add_alpha_ != nullptr && + "`LightningAttentionInfinilm` failed to create the Add alpha " + "scalar"); + } + + ~Operator() { + if (!ascend::IsAclRuntimeAlive()) return; + + // Null cached descriptors: they are still referenced by the Repeatable + // executors. See `AclTensorCache::release()`. + state_cache_.release(); + decayed_cache_.release(); + outer_cache_.release(); + ratio_cache_.release(); + outer_k_cache_.release(); + outer_v_cache_.release(); + output_q_cache_.release(); + output_cache_.release(); + + if (add_alpha_) aclDestroyScalar(add_alpha_); + } + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) const override { + auto stream = static_cast(stream_); + const size_t element_size = kDataTypeToSize.at(data_type_); + const size_t head_elements = head_dim_ * head_dim_; + const size_t state_bytes = num_heads_ * head_elements * element_size; + const size_t ratio_bytes = num_heads_ * element_size; + + // The host reads slope and state indices, and queues device copies below. + // Synchronize once so both reads observe all earlier work on this stream. + CheckAcl(aclrtSynchronizeStream(stream), + "aclrtSynchronizeStream(lightning attention input)"); + + auto& pool = ascend::GetWorkspacePool(); + auto& state_arena = + pool.Ensure(stream, state_bytes, "lightning_attention_state"); + auto& decayed_arena = + pool.Ensure(stream, state_bytes, "lightning_attention_decayed"); + auto& outer_arena = + pool.Ensure(stream, state_bytes, "lightning_attention_outer"); + auto& ratio_arena = + pool.Ensure(stream, ratio_bytes, "lightning_attention_ratio"); + + BuildRatio(slope, ratio_arena.buf, element_size); + + std::vector initial_rows(batch_size_); + std::vector final_rows(batch_size_); + ReadIndices(initial_state_indices, initial_rows); + ReadIndices(final_state_indices, final_rows); + + void* state_buf = state_arena.buf; + void* decayed_buf = decayed_arena.buf; + void* outer_buf = outer_arena.buf; + const void* ratio_buf = ratio_arena.buf; + + for (Tensor::Size b = 0; b < batch_size_; ++b) { + const auto* initial_row_ptr = + static_cast(initial_state.data()) + + initial_rows[b] * state_pool_stride_ * element_size; + auto* final_row_ptr = + static_cast(initial_state.data()) + + final_rows[b] * state_pool_stride_ * element_size; + + // Stage the source row in the workspace. The destination row is not + // touched until the final writeback, so an aliased source row remains + // read-only for the whole request. + CheckAcl(aclrtMemcpyAsync(state_buf, state_bytes, initial_row_ptr, + state_bytes, ACL_MEMCPY_DEVICE_TO_DEVICE, + stream), + "aclrtMemcpyAsync(state staging)"); + + for (Tensor::Size t = 0; t < seq_len_; ++t) { + const auto* q_t = + static_cast(q.data()) + + (b * q_batch_stride_ + t * q_seq_stride_) * element_size; + const auto* k_t = + static_cast(k.data()) + + (b * k_batch_stride_ + t * k_seq_stride_) * element_size; + const auto* v_t = + static_cast(v.data()) + + (b * v_batch_stride_ + t * v_seq_stride_) * element_size; + auto* out_t = + static_cast(out.data()) + + (b * out_batch_stride_ + t * out_seq_stride_) * element_size; + + MulDecay(state_buf, ratio_buf, decayed_buf, stream); + OuterProduct(k_t, v_t, outer_buf, stream); + AddState(decayed_buf, outer_buf, state_buf, stream); + WriteOutput(q_t, state_buf, out_t, stream); + } + + CheckAcl(aclrtMemcpyAsync(final_row_ptr, state_bytes, state_buf, + state_bytes, ACL_MEMCPY_DEVICE_TO_DEVICE, + stream), + "aclrtMemcpyAsync(state writeback)"); + } + } + + private: + void CheckAcl(aclError error, const char* what) const { + (void)what; + assert(error == ACL_SUCCESS); + } + + // Materialize exp(-slope) once per forward call. The ratio buffer is + // [heads, 1, 1] so that aclnnMul broadcasts it over the [heads, head_dim, + // head_dim] state. + void BuildRatio(const Tensor& slope, void* ratio_buf, + size_t element_size) const { + std::vector host_slope(num_heads_); + CheckAcl(aclrtMemcpy(host_slope.data(), num_heads_ * sizeof(float), + const_cast(slope.data()), + num_heads_ * sizeof(float), + ACL_MEMCPY_DEVICE_TO_HOST), + "aclrtMemcpy(slope to host)"); + + std::vector host_ratio(num_heads_ * element_size); + for (Tensor::Size h = 0; h < num_heads_; ++h) { + const float ratio = std::exp(-host_slope[h]); + auto* destination = host_ratio.data() + h * element_size; + + if (data_type_ == DataType::kFloat32) { + std::memcpy(destination, &ratio, sizeof(float)); + } else if (data_type_ == DataType::kFloat16) { + const auto converted = Float16::FromFloat(ratio); + std::memcpy(destination, &converted, sizeof(converted)); + } else { + const auto converted = BFloat16::FromFloat(ratio); + std::memcpy(destination, &converted, sizeof(converted)); + } + } + + CheckAcl(aclrtMemcpy(ratio_buf, num_heads_ * element_size, + host_ratio.data(), num_heads_ * element_size, + ACL_MEMCPY_HOST_TO_DEVICE), + "aclrtMemcpy(ratio to device)"); + } + + void ReadIndices(const Tensor& indices, std::vector& host) const { + host.resize(batch_size_); + if (index_dtype_ == DataType::kInt64) { + CheckAcl(aclrtMemcpy(host.data(), batch_size_ * sizeof(int64_t), + const_cast(indices.data()), + batch_size_ * sizeof(int64_t), + ACL_MEMCPY_DEVICE_TO_HOST), + "aclrtMemcpy(indices int64)"); + } else { + std::vector narrowed(batch_size_); + CheckAcl(aclrtMemcpy(narrowed.data(), batch_size_ * sizeof(int32_t), + const_cast(indices.data()), + batch_size_ * sizeof(int32_t), + ACL_MEMCPY_DEVICE_TO_HOST), + "aclrtMemcpy(indices int32)"); + for (Tensor::Size b = 0; b < batch_size_; ++b) { + host[b] = static_cast(narrowed[b]); + } + } + } + + void MulDecay(const void* state, const void* ratio, void* decayed, + aclrtStream stream) const { + auto* t_state = state_cache_.get(const_cast(state)); + auto* t_ratio = ratio_cache_.get(const_cast(ratio)); + auto* t_decayed = decayed_cache_.get(decayed); + + if (!mul_exec_) { + CheckAcl(aclnnMulGetWorkspaceSize(t_state, t_ratio, t_decayed, &mul_ws_, + &mul_exec_), + "aclnnMulGetWorkspaceSize"); + aclSetAclOpExecutorRepeatable(mul_exec_); + } else { + aclSetInputTensorAddr(mul_exec_, 0, t_state, const_cast(state)); + aclSetInputTensorAddr(mul_exec_, 1, t_ratio, const_cast(ratio)); + aclSetOutputTensorAddr(mul_exec_, 0, t_decayed, decayed); + } + + auto& arena = ascend::GetWorkspacePool().Ensure( + stream, mul_ws_, "lightning_attention_mul_workspace"); + CheckAcl(aclnnMul(arena.buf, mul_ws_, mul_exec_, stream), "aclnnMul"); + } + + void OuterProduct(const void* k_t, const void* v_t, void* outer, + aclrtStream stream) const { + auto* t_k = outer_k_cache_.get(const_cast(k_t)); + auto* t_v = outer_v_cache_.get(const_cast(v_t)); + auto* t_outer = outer_cache_.get(outer); + + if (!outer_exec_) { + const int8_t cube_math_type = data_type_ == DataType::kFloat32 ? 0 : 1; + CheckAcl(aclnnMatmulGetWorkspaceSize(t_k, t_v, t_outer, cube_math_type, + &outer_ws_, &outer_exec_), + "aclnnMatmulGetWorkspaceSize(outer)"); + aclSetAclOpExecutorRepeatable(outer_exec_); + } else { + aclSetInputTensorAddr(outer_exec_, 0, t_k, const_cast(k_t)); + aclSetInputTensorAddr(outer_exec_, 1, t_v, const_cast(v_t)); + aclSetOutputTensorAddr(outer_exec_, 0, t_outer, outer); + } + + auto& arena = ascend::GetWorkspacePool().Ensure( + stream, outer_ws_, "lightning_attention_outer_workspace"); + CheckAcl(aclnnMatmul(arena.buf, outer_ws_, outer_exec_, stream), + "aclnnMatmul(outer)"); + } + + void AddState(const void* decayed, const void* outer, void* state, + aclrtStream stream) const { + auto* t_decayed = decayed_cache_.get(const_cast(decayed)); + auto* t_outer = outer_cache_.get(const_cast(outer)); + auto* t_state = state_cache_.get(state); + + if (!add_exec_) { + CheckAcl(aclnnAddGetWorkspaceSize(t_decayed, t_outer, add_alpha_, t_state, + &add_ws_, &add_exec_), + "aclnnAddGetWorkspaceSize"); + aclSetAclOpExecutorRepeatable(add_exec_); + } else { + aclSetInputTensorAddr(add_exec_, 0, t_decayed, + const_cast(decayed)); + aclSetInputTensorAddr(add_exec_, 1, t_outer, const_cast(outer)); + aclSetOutputTensorAddr(add_exec_, 0, t_state, state); + } + + auto& arena = ascend::GetWorkspacePool().Ensure( + stream, add_ws_, "lightning_attention_add_workspace"); + CheckAcl(aclnnAdd(arena.buf, add_ws_, add_exec_, stream), "aclnnAdd"); + } + + void WriteOutput(const void* q_t, const void* state, void* out_t, + aclrtStream stream) const { + auto* t_q = output_q_cache_.get(const_cast(q_t)); + auto* t_state = state_cache_.get(const_cast(state)); + auto* t_out = output_cache_.get(out_t); + + if (!output_exec_) { + const int8_t cube_math_type = data_type_ == DataType::kFloat32 ? 0 : 1; + CheckAcl(aclnnMatmulGetWorkspaceSize(t_q, t_state, t_out, cube_math_type, + &output_ws_, &output_exec_), + "aclnnMatmulGetWorkspaceSize(output)"); + aclSetAclOpExecutorRepeatable(output_exec_); + } else { + aclSetInputTensorAddr(output_exec_, 0, t_q, const_cast(q_t)); + aclSetInputTensorAddr(output_exec_, 1, t_state, const_cast(state)); + aclSetOutputTensorAddr(output_exec_, 0, t_out, out_t); + } + + auto& arena = ascend::GetWorkspacePool().Ensure( + stream, output_ws_, "lightning_attention_output_workspace"); + CheckAcl(aclnnMatmul(arena.buf, output_ws_, output_exec_, stream), + "aclnnMatmul(output)"); + } + + mutable ascend::AclTensorCache state_cache_; + + mutable ascend::AclTensorCache decayed_cache_; + + mutable ascend::AclTensorCache outer_cache_; + + mutable ascend::AclTensorCache ratio_cache_; + + mutable ascend::AclTensorCache outer_k_cache_; + + mutable ascend::AclTensorCache outer_v_cache_; + + mutable ascend::AclTensorCache output_q_cache_; + + mutable ascend::AclTensorCache output_cache_; + + float add_alpha_storage_ = 1.0f; + + aclScalar* add_alpha_ = nullptr; + + mutable aclOpExecutor* mul_exec_ = nullptr; + + mutable uint64_t mul_ws_ = 0; + + mutable aclOpExecutor* outer_exec_ = nullptr; + + mutable uint64_t outer_ws_ = 0; + + mutable aclOpExecutor* add_exec_ = nullptr; + + mutable uint64_t add_ws_ = 0; + + mutable aclOpExecutor* output_exec_ = nullptr; + + mutable uint64_t output_ws_ = 0; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ASCEND_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ \ No newline at end of file diff --git a/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h b/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h new file mode 100644 index 000000000..8e8ef1c1d --- /dev/null +++ b/src/native/cpu/ops/lightning_attention_infinilm/lightning_attention_infinilm.h @@ -0,0 +1,139 @@ +#ifndef INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ +#define INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ + +#include +#include +#include + +#include "base/lightning_attention_infinilm.h" +#include "common/generic_utils.h" +#include "data_type.h" +#include "native/cpu/caster_.h" +#include "tensor.h" + +namespace infini::ops { + +template <> +class Operator + : public LightningAttentionInfinilm, Caster { + public: + Operator(const Tensor q, const Tensor k, const Tensor v, const Tensor slope, + Tensor initial_state, const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) + : LightningAttentionInfinilm{q, k, v, slope, initial_state, + initial_state_indices, final_state_indices, + out} {} + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) const override { + DispatchFunc( + out.dtype(), + [&](auto tag) { + using T = typename decltype(tag)::type; + Compute(q, k, v, slope, initial_state, initial_state_indices, + final_state_indices, out); + }, + "`Operator::operator()`"); + } + + private: + template + void Compute(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) const { + const auto* q_ptr = static_cast(q.data()); + const auto* k_ptr = static_cast(k.data()); + const auto* v_ptr = static_cast(v.data()); + const auto* slope_ptr = static_cast(slope.data()); + auto* state_ptr = static_cast(initial_state.data()); + auto* out_ptr = static_cast(out.data()); + + const bool int64_indices = index_dtype_ == DataType::kInt64; + const auto* initial_indices = initial_state_indices.data(); + const auto* final_indices = final_state_indices.data(); + + // The recurrent state of one request, accumulated in float32. + std::vector state(num_heads_ * head_dim_ * head_dim_); + + for (Tensor::Size b = 0; b < batch_size_; ++b) { + Tensor::Size initial_row; + Tensor::Size final_row; + if (int64_indices) { + initial_row = static_cast( + static_cast(initial_indices)[b]); + final_row = + static_cast(static_cast(final_indices)[b]); + } else { + initial_row = static_cast( + static_cast(initial_indices)[b]); + final_row = + static_cast(static_cast(final_indices)[b]); + } + + const T* initial_row_ptr = state_ptr + initial_row * state_pool_stride_; + for (Tensor::Size h = 0; h < num_heads_; ++h) { + const T* head_ptr = initial_row_ptr + h * state_head_stride_; + float* state_head = state.data() + h * head_dim_ * head_dim_; + for (Tensor::Size i = 0; i < head_dim_; ++i) { + for (Tensor::Size j = 0; j < head_dim_; ++j) { + state_head[i * head_dim_ + j] = + Cast(head_ptr[i * state_row_stride_ + + j * state_column_stride_]); + } + } + } + + for (Tensor::Size t = 0; t < seq_len_; ++t) { + for (Tensor::Size h = 0; h < num_heads_; ++h) { + const float ratio = std::exp(-slope_ptr[h * slope_stride_]); + const T* q_row = q_ptr + b * q_batch_stride_ + t * q_seq_stride_ + + h * q_head_stride_; + const T* k_row = k_ptr + b * k_batch_stride_ + t * k_seq_stride_ + + h * k_head_stride_; + const T* v_row = v_ptr + b * v_batch_stride_ + t * v_seq_stride_ + + h * v_head_stride_; + float* state_head = state.data() + h * head_dim_ * head_dim_; + + for (Tensor::Size i = 0; i < head_dim_; ++i) { + const float k_i = Cast(k_row[i]); + for (Tensor::Size j = 0; j < head_dim_; ++j) { + state_head[i * head_dim_ + j] = + ratio * state_head[i * head_dim_ + j] + + k_i * Cast(v_row[j]); + } + } + + T* out_row = out_ptr + b * out_batch_stride_ + t * out_seq_stride_ + + h * out_head_stride_; + for (Tensor::Size j = 0; j < head_dim_; ++j) { + float acc = 0.0f; + for (Tensor::Size i = 0; i < head_dim_; ++i) { + acc += Cast(q_row[i]) * state_head[i * head_dim_ + j]; + } + out_row[j] = Cast(acc); + } + } + } + + // Only the destination row is updated; the initial row stays untouched. + T* final_row_ptr = state_ptr + final_row * state_pool_stride_; + for (Tensor::Size h = 0; h < num_heads_; ++h) { + T* head_ptr = final_row_ptr + h * state_head_stride_; + const float* state_head = state.data() + h * head_dim_ * head_dim_; + for (Tensor::Size i = 0; i < head_dim_; ++i) { + for (Tensor::Size j = 0; j < head_dim_; ++j) { + head_ptr[i * state_row_stride_ + j * state_column_stride_] = + Cast(state_head[i * head_dim_ + j]); + } + } + } + } + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CPU_LIGHTNING_ATTENTION_INFINILM_H_ diff --git a/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h b/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h new file mode 100644 index 000000000..9a21f93f4 --- /dev/null +++ b/src/native/cuda/nvidia/ops/lightning_attention_infinilm/kernel.h @@ -0,0 +1,22 @@ +#ifndef INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ +#define INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ + +#include + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/lightning_attention_infinilm/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaLightningAttentionInfinilm> { + public: + using CudaLightningAttentionInfinilm< + Runtime>::CudaLightningAttentionInfinilm; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ diff --git a/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh b/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh new file mode 100644 index 000000000..35861cf2c --- /dev/null +++ b/src/native/cuda/ops/lightning_attention_infinilm/kernel.cuh @@ -0,0 +1,98 @@ +#ifndef INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ +#define INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ + +#include + +#include "native/cuda/caster.cuh" +#include "native/cuda/kernel_commons.cuh" + +namespace infini::ops { + +/// One block per `(batch, head)` and one thread per state/output column. +/// +/// The recurrent state `[head_dim, head_dim]` of the destination pool row is +/// staged first (the source row is copied into it when the two rows differ, so +/// the source row stays untouched), then updated in place: +/// +/// state[i][j] = ratio * state[i][j] + k[i] * v[j] +/// out[j] = sum_i q[i] * state[i][j] +/// +/// Thread `j` owns column `j`, so the state update needs no cross-thread +/// synchronization; only the shared `k`/`q` rows do. The kernel must be +/// launched with exactly `head_dim` threads so that every thread reaches the +/// barriers. +template +__global__ void LightningAttentionInfinilmKernel( + Data* out, Data* state_pool, const Data* q, const Data* k, const Data* v, + const float* slope, const Index* initial_state_indices, + const Index* final_state_indices, size_t seq_len, size_t head_dim, + ptrdiff_t state_pool_stride, ptrdiff_t state_head_stride, + ptrdiff_t state_row_stride, ptrdiff_t q_batch_stride, ptrdiff_t q_seq_stride, + ptrdiff_t q_head_stride, ptrdiff_t k_batch_stride, ptrdiff_t k_seq_stride, + ptrdiff_t k_head_stride, ptrdiff_t v_batch_stride, ptrdiff_t v_seq_stride, + ptrdiff_t v_head_stride, ptrdiff_t out_batch_stride, ptrdiff_t out_seq_stride, + ptrdiff_t out_head_stride, ptrdiff_t slope_stride) { + const size_t batch = blockIdx.y; + const size_t head = blockIdx.x; + const size_t column = threadIdx.x; + + const size_t initial_row = static_cast(initial_state_indices[batch]); + const size_t final_row = static_cast(final_state_indices[batch]); + + Data* state = state_pool + final_row * state_pool_stride + + head * state_head_stride; + if (initial_row != final_row) { + const Data* source = state_pool + initial_row * state_pool_stride + + head * state_head_stride; + for (size_t index = column; index < head_dim * head_dim; + index += blockDim.x) { + const size_t i = index / head_dim; + const size_t j = index % head_dim; + state[i * state_row_stride + j] = source[i * state_row_stride + j]; + } + } + __syncthreads(); + + extern __shared__ float shared[]; + float* shared_k = shared; + float* shared_q = shared + head_dim; + + const float ratio = expf(-slope[head * slope_stride]); + const Data* q_head = q + batch * q_batch_stride + head * q_head_stride; + const Data* k_head = k + batch * k_batch_stride + head * k_head_stride; + const Data* v_head = v + batch * v_batch_stride + head * v_head_stride; + Data* out_head = out + batch * out_batch_stride + head * out_head_stride; + + for (size_t t = 0; t < seq_len; ++t) { + shared_k[column] = Caster::template Cast( + k_head[t * k_seq_stride + column]); + shared_q[column] = Caster::template Cast( + q_head[t * q_seq_stride + column]); + __syncthreads(); + + const float v_column = Caster::template Cast( + v_head[t * v_seq_stride + column]); + for (size_t i = 0; i < head_dim; ++i) { + Data* element = state + i * state_row_stride + column; + const float updated = ratio * Caster::template Cast(*element) + + shared_k[i] * v_column; + *element = Caster::template Cast(updated); + } + + float accumulator = 0.0f; + for (size_t i = 0; i < head_dim; ++i) { + accumulator += shared_q[i] * Caster::template Cast( + state[i * state_row_stride + column]); + } + out_head[t * out_seq_stride + column] = + Caster::template Cast(accumulator); + + // The next iteration overwrites the shared rows, so all threads must have + // finished reading them. + __syncthreads(); + } +} + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_CUH_ diff --git a/src/native/cuda/ops/lightning_attention_infinilm/kernel.h b/src/native/cuda/ops/lightning_attention_infinilm/kernel.h new file mode 100644 index 000000000..e90c52bba --- /dev/null +++ b/src/native/cuda/ops/lightning_attention_infinilm/kernel.h @@ -0,0 +1,78 @@ +#ifndef INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ +#define INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ + +#include +#include +#include + +#include "base/lightning_attention_infinilm.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/lightning_attention_infinilm/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +using LightningAttentionInfinilmDataTypes = + ConcatType, ReducedFloatTypes>; + +using LightningAttentionInfinilmIndexTypes = + List; + +template +class CudaLightningAttentionInfinilm : public LightningAttentionInfinilm { + public: + using LightningAttentionInfinilm::LightningAttentionInfinilm; + + void operator()(const Tensor q, const Tensor k, const Tensor v, + const Tensor slope, Tensor initial_state, + const Tensor initial_state_indices, + const Tensor final_state_indices, Tensor out) const override { + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + + // One thread per state column, so `head_dim` has to fit into one block. + assert(head_dim_ > 0 && + static_cast(head_dim_) <= BackendMaxBlockSize::value && + "`LightningAttentionInfinilm` requires head_dim to fit one block"); + assert(batch_size_ <= 65535 && + "`LightningAttentionInfinilm` requires batch_size <= 65535"); + + dim3 grid(static_cast(num_heads_), + static_cast(batch_size_)); + dim3 block(static_cast(head_dim_)); + const size_t shared_bytes = 2 * head_dim_ * sizeof(float); + + DispatchFunc( + {static_cast(out.dtype()), static_cast(index_dtype_)}, + [&](auto list_tag) { + using T = TypeMapType(list_tag)>; + using TIndex = + TypeMapType(list_tag)>; + + LightningAttentionInfinilmKernel + <<>>( + reinterpret_cast(out.data()), + reinterpret_cast(initial_state.data()), + reinterpret_cast(q.data()), + reinterpret_cast(k.data()), + reinterpret_cast(v.data()), + static_cast(slope.data()), + reinterpret_cast(initial_state_indices.data()), + reinterpret_cast(final_state_indices.data()), + seq_len_, head_dim_, state_pool_stride_, state_head_stride_, + state_row_stride_, q_batch_stride_, q_seq_stride_, + q_head_stride_, k_batch_stride_, k_seq_stride_, k_head_stride_, + v_batch_stride_, v_seq_stride_, v_head_stride_, + out_batch_stride_, out_seq_stride_, out_head_stride_, + slope_stride_); + }, + "CudaLightningAttentionInfinilm::operator()"); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_LIGHTNING_ATTENTION_INFINILM_KERNEL_H_ diff --git a/tests/test_lightning_attention_infinilm.py b/tests/test_lightning_attention_infinilm.py new file mode 100644 index 000000000..5a458c7bd --- /dev/null +++ b/tests/test_lightning_attention_infinilm.py @@ -0,0 +1,122 @@ +import infini.ops +import pytest +import torch + +from tests.utils import Payload, empty_strided, get_stream, randn_strided + +# (batch, seq_len, num_heads, head_dim, pool_size) +_SHAPES = ( + (1, 1, 2, 4, 1), + (2, 3, 2, 4, 4), + (3, 1, 4, 8, 6), + (1, 5, 4, 8, 2), +) + +# The recurrent state is stored in the tensor dtype, so rounding accumulates +# over the sequence; reduced-precision cases therefore need looser tolerances +# than a single elementwise operator would. +_DTYPE_CASES = ( + (torch.float32, 1e-5, 1e-6), + (torch.float16, 2e-2, 2e-2), + (torch.bfloat16, 4e-2, 4e-2), +) + + +def _make_tensors(shape, dtype, device): + batch, seq_len, num_heads, head_dim, pool_size = shape + + q = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) + k = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) + v = randn_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) + # The decay has to be positive, otherwise `exp(-slope)` would grow. + slope = randn_strided((num_heads,), None, dtype=torch.float32, device=device).abs() * 0.5 + state = randn_strided( + (pool_size, num_heads, head_dim, head_dim), None, dtype=dtype, device=device + ) + + # Deliberately read and write different pool rows so that the test also + # covers the "initial row must stay untouched" contract. Rows are disjoint + # across requests because requests may execute concurrently. + initial_indices = torch.arange(batch, dtype=torch.int32, device=device) % pool_size + final_indices = (initial_indices + batch) % pool_size + + out = empty_strided((batch, seq_len, num_heads, head_dim), None, dtype=dtype, device=device) + + return q, k, v, slope, state, initial_indices, final_indices, out + + +def _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices): + """Recurrent reference: S <- ratio * S + k^T v ; out = q @ S.""" + + q = q.float() + k = k.float() + v = v.float() + state = state.float().clone() + out = torch.empty_like(q) + ratio = torch.exp(-slope.float()) + + for b in range(q.shape[0]): + initial_row = int(initial_indices[b].item()) + final_row = int(final_indices[b].item()) + + current = state[initial_row].clone() # [heads, head_dim, head_dim] + for t in range(q.shape[1]): + current = ratio[:, None, None] * current + k[b, t].unsqueeze(-1) * v[b, t].unsqueeze(-2) + # The implementation stores the recurrent state in the tensor dtype, + # so the reference has to round it before the next read. + current = current.to(state.dtype) + out[b, t] = torch.einsum("hd,hde->he", q[b, t], current.float()) + + state[final_row] = current + + return out, state + + +def _run_lightning_attention(q, k, v, slope, state, initial_indices, final_indices, out): + infini.ops.lightning_attention_infinilm( + q, + k, + v, + slope, + state, + initial_indices, + final_indices, + out, + stream=get_stream(q.device), + ) + return out + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize(("dtype", "rtol", "atol"), _DTYPE_CASES) +def test_lightning_attention_infinilm(shape, dtype, device, rtol, atol): + tensors = _make_tensors(shape, dtype, device) + + return Payload(_run_lightning_attention, _reference_out, tensors, {}, rtol=rtol, atol=atol) + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize(("dtype", "rtol", "atol"), _DTYPE_CASES) +def test_lightning_attention_infinilm_state_pool(shape, dtype, device, rtol, atol): + tensors = _make_tensors(shape, dtype, device) + + return Payload(_run_lightning_attention_state, _reference_state, tensors, {}, rtol=rtol, atol=atol) + + +def _reference_out(q, k, v, slope, state, initial_indices, final_indices, out): + reference, _ = _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices) + out.copy_(reference.to(out.dtype)) + return out + + +def _reference_state(q, k, v, slope, state, initial_indices, final_indices, out): + _, reference = _torch_lightning_attention(q, k, v, slope, state, initial_indices, final_indices) + state.copy_(reference.to(state.dtype)) + return state + + +def _run_lightning_attention_state(q, k, v, slope, state, initial_indices, final_indices, out): + _run_lightning_attention(q, k, v, slope, state, initial_indices, final_indices, out) + return state