From b645a3e73e5e634e1be1096bb5fa4d42a4b6db56 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 7 Aug 2026 02:48:08 +0800 Subject: [PATCH] feat(ops): add vLLM-aligned paged attention v1 --- src/base/paged_attention_v1.h | 202 ++++++++++++++++ .../iluvatar/ops/paged_attention_v1/kernel.h | 20 ++ .../metax/ops/paged_attention_v1/kernel.h | 20 ++ .../moore/ops/paged_attention_v1/kernel.h | 20 ++ .../nvidia/ops/paged_attention_v1/kernel.h | 20 ++ .../ops/paged_attention_infinilm/kernel.cuh | 134 ----------- .../ops/paged_attention_infinilm/kernel.h | 23 +- .../cuda/ops/paged_attention_v1/kernel.cuh | 151 ++++++++++++ .../cuda/ops/paged_attention_v1/kernel.h | 96 ++++++++ tests/test_paged_attention_v1.py | 225 ++++++++++++++++++ 10 files changed, 767 insertions(+), 144 deletions(-) create mode 100644 src/base/paged_attention_v1.h create mode 100644 src/native/cuda/iluvatar/ops/paged_attention_v1/kernel.h create mode 100644 src/native/cuda/metax/ops/paged_attention_v1/kernel.h create mode 100644 src/native/cuda/moore/ops/paged_attention_v1/kernel.h create mode 100644 src/native/cuda/nvidia/ops/paged_attention_v1/kernel.h create mode 100644 src/native/cuda/ops/paged_attention_v1/kernel.cuh create mode 100644 src/native/cuda/ops/paged_attention_v1/kernel.h create mode 100644 tests/test_paged_attention_v1.py diff --git a/src/base/paged_attention_v1.h b/src/base/paged_attention_v1.h new file mode 100644 index 000000000..31f1c58ae --- /dev/null +++ b/src/base/paged_attention_v1.h @@ -0,0 +1,202 @@ +#ifndef INFINI_OPS_BASE_PAGED_ATTENTION_V1_H_ +#define INFINI_OPS_BASE_PAGED_ATTENTION_V1_H_ + +#include +#include +#include +#include +#include + +#include "data_type.h" +#include "operator.h" +#include "tensor.h" + +namespace infini::ops { + +// Aligned with vLLM `_custom_ops.paged_attention_v1`. +class PagedAttentionV1 : public Operator { + public: + PagedAttentionV1(const Tensor query, const Tensor key_cache, + const Tensor value_cache, const Tensor block_tables, + const Tensor seq_lens, + const std::optional alibi_slopes, + const int64_t num_kv_heads, const double scale, + const int64_t block_size, const int64_t max_seq_len, + const std::string kv_cache_dtype, const double k_scale, + const double v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, + const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step, Tensor out) + : dtype_{query.dtype()}, + index_dtype_{block_tables.dtype()}, + num_seqs_{query.size(0)}, + num_heads_{query.size(1)}, + num_kv_heads_{key_cache.size(1)}, + head_size_{query.size(2)}, + block_size_{key_cache.size(3)}, + max_num_blocks_per_seq_{block_tables.size(1)}, + key_cache_x_{key_cache.size(4)}, + query_stride_{query.stride(0)}, + query_head_stride_{query.stride(1)}, + key_cache_block_stride_{key_cache.stride(0)}, + key_cache_head_stride_{key_cache.stride(1)}, + key_cache_dim_stride_{key_cache.stride(2)}, + key_cache_slot_stride_{key_cache.stride(3)}, + key_cache_x_stride_{key_cache.stride(4)}, + value_cache_block_stride_{value_cache.stride(0)}, + value_cache_head_stride_{value_cache.stride(1)}, + value_cache_dim_stride_{value_cache.stride(2)}, + value_cache_slot_stride_{value_cache.stride(3)}, + out_stride_{out.stride(0)}, + out_head_stride_{out.stride(1)}, + block_table_batch_stride_{block_tables.stride(0)}, + seq_lens_stride_{seq_lens.stride(0)}, + scale_{scale}, + device_index_{query.device().index()} { + assert(query.ndim() == 3 && out.ndim() == 3 && + "`PagedAttentionV1` requires 3D `query` and `out`"); + assert(key_cache.ndim() == 5 && value_cache.ndim() == 4 && + "`PagedAttentionV1` requires a 5D key cache and 4D value cache"); + assert(block_tables.ndim() == 2 && seq_lens.ndim() == 1 && + "`PagedAttentionV1` requires 2D block tables and 1D sequence " + "lengths"); + assert((dtype_ == DataType::kFloat16 || dtype_ == DataType::kBFloat16) && + "`PagedAttentionV1` supports float16 and bfloat16"); + assert(out.dtype() == dtype_ && key_cache.dtype() == dtype_ && + value_cache.dtype() == dtype_ && + "`PagedAttentionV1` requires matching data dtypes"); + assert(IsIndexDtype(index_dtype_) && seq_lens.dtype() == index_dtype_ && + "`PagedAttentionV1` requires matching integer metadata dtypes"); + assert(query.shape() == out.shape() && + "`PagedAttentionV1` requires `out` to match `query`"); + assert(num_kv_heads > 0 && + static_cast(num_kv_heads) == num_kv_heads_ && + num_heads_ % num_kv_heads_ == 0 && + "`PagedAttentionV1` received incompatible KV heads"); + assert(key_cache.size(0) == value_cache.size(0) && + value_cache.size(1) == num_kv_heads_ && + key_cache.size(2) * key_cache_x_ == head_size_ && + value_cache.size(2) == head_size_ && + "`PagedAttentionV1` cache shapes do not match `query`"); + assert(block_size > 0 && + static_cast(block_size) == block_size_ && + value_cache.size(3) == block_size_ && + "`PagedAttentionV1` received an incompatible block size"); + assert(key_cache_x_ == + static_cast(16 / query.element_size()) && + "`PagedAttentionV1` requires the vLLM vectorized key layout"); + assert(max_seq_len > 0 && + max_seq_len <= + static_cast(max_num_blocks_per_seq_ * block_size_) && + "`PagedAttentionV1` maximum sequence length exceeds block table " + "capacity"); + assert((head_size_ == 64 || head_size_ == 128) && + "`PagedAttentionV1` supports head sizes 64 and 128"); + assert(block_tables.size(0) == num_seqs_ && seq_lens.size(0) == num_seqs_ && + "`PagedAttentionV1` metadata batch sizes must match `query`"); + assert(query.stride(2) == 1 && out.stride(2) == 1 && + key_cache_x_stride_ == 1 && + "`PagedAttentionV1` requires contiguous head-vector dimensions"); + assert(block_tables.stride(1) == 1 && seq_lens_stride_ == 1 && + "`PagedAttentionV1` requires contiguous index rows"); + assert(!alibi_slopes.has_value() || + (alibi_slopes->dtype() == DataType::kFloat32 && + alibi_slopes->ndim() == 1 && alibi_slopes->size(0) == num_heads_ && + alibi_slopes->stride(0) == 1) && + "`PagedAttentionV1` received incompatible ALiBi slopes"); + assert(kv_cache_dtype == "auto" && k_scale == 1.0 && v_scale == 1.0 && + "`PagedAttentionV1` currently supports unquantized KV caches"); + assert(blocksparse_local_blocks == 0 && blocksparse_vert_stride == 0 && + blocksparse_head_sliding_step == 0 && + "`PagedAttentionV1` does not yet support block-sparse attention"); + assert(blocksparse_block_size > 0 && + "`PagedAttentionV1` requires a positive block-sparse block size"); + + const auto same_device_as_query = [&](const Tensor tensor) { + return tensor.device() == query.device(); + }; + assert(same_device_as_query(key_cache) && + same_device_as_query(value_cache) && + same_device_as_query(block_tables) && + same_device_as_query(seq_lens) && same_device_as_query(out) && + (!alibi_slopes.has_value() || same_device_as_query(*alibi_slopes)) && + "`PagedAttentionV1` tensors must be on the same device"); + + (void)tp_rank; + } + + virtual void operator()( + const Tensor query, const Tensor key_cache, const Tensor value_cache, + const Tensor block_tables, const Tensor seq_lens, + const std::optional alibi_slopes, const int64_t num_kv_heads, + const double scale, const int64_t block_size, const int64_t max_seq_len, + const std::string kv_cache_dtype, const double k_scale, + const double v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, + const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step, Tensor out) const = 0; + + protected: + static bool IsIndexDtype(DataType dtype) { + return dtype == DataType::kInt32 || dtype == DataType::kInt64 || + dtype == DataType::kUInt32; + } + + DataType dtype_; + + DataType index_dtype_; + + Tensor::Size num_seqs_{0}; + + Tensor::Size num_heads_{0}; + + Tensor::Size num_kv_heads_{0}; + + Tensor::Size head_size_{0}; + + Tensor::Size block_size_{0}; + + Tensor::Size max_num_blocks_per_seq_{0}; + + Tensor::Size key_cache_x_{0}; + + Tensor::Stride query_stride_{0}; + + Tensor::Stride query_head_stride_{0}; + + Tensor::Stride key_cache_block_stride_{0}; + + Tensor::Stride key_cache_head_stride_{0}; + + Tensor::Stride key_cache_dim_stride_{0}; + + Tensor::Stride key_cache_slot_stride_{0}; + + Tensor::Stride key_cache_x_stride_{0}; + + Tensor::Stride value_cache_block_stride_{0}; + + Tensor::Stride value_cache_head_stride_{0}; + + Tensor::Stride value_cache_dim_stride_{0}; + + Tensor::Stride value_cache_slot_stride_{0}; + + Tensor::Stride out_stride_{0}; + + Tensor::Stride out_head_stride_{0}; + + Tensor::Stride block_table_batch_stride_{0}; + + Tensor::Stride seq_lens_stride_{0}; + + double scale_{1.0}; + + int device_index_{0}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_BASE_PAGED_ATTENTION_V1_H_ diff --git a/src/native/cuda/iluvatar/ops/paged_attention_v1/kernel.h b/src/native/cuda/iluvatar/ops/paged_attention_v1/kernel.h new file mode 100644 index 000000000..afd378451 --- /dev/null +++ b/src/native/cuda/iluvatar/ops/paged_attention_v1/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_ILUVATAR_PAGED_ATTENTION_V1_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_PAGED_ATTENTION_V1_KERNEL_H_ + +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/paged_attention_v1/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaPagedAttentionV1> { + public: + using CudaPagedAttentionV1< + Runtime>::CudaPagedAttentionV1; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ILUVATAR_PAGED_ATTENTION_V1_KERNEL_H_ diff --git a/src/native/cuda/metax/ops/paged_attention_v1/kernel.h b/src/native/cuda/metax/ops/paged_attention_v1/kernel.h new file mode 100644 index 000000000..a207ba613 --- /dev/null +++ b/src/native/cuda/metax/ops/paged_attention_v1/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_METAX_PAGED_ATTENTION_V1_KERNEL_H_ +#define INFINI_OPS_METAX_PAGED_ATTENTION_V1_KERNEL_H_ + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/paged_attention_v1/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaPagedAttentionV1> { + public: + using CudaPagedAttentionV1< + Runtime>::CudaPagedAttentionV1; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_METAX_PAGED_ATTENTION_V1_KERNEL_H_ diff --git a/src/native/cuda/moore/ops/paged_attention_v1/kernel.h b/src/native/cuda/moore/ops/paged_attention_v1/kernel.h new file mode 100644 index 000000000..425260324 --- /dev/null +++ b/src/native/cuda/moore/ops/paged_attention_v1/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_MOORE_PAGED_ATTENTION_V1_KERNEL_H_ +#define INFINI_OPS_MOORE_PAGED_ATTENTION_V1_KERNEL_H_ + +#include "native/cuda/moore/caster.cuh" +#include "native/cuda/moore/runtime_.h" +#include "native/cuda/ops/paged_attention_v1/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaPagedAttentionV1> { + public: + using CudaPagedAttentionV1< + Runtime>::CudaPagedAttentionV1; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_MOORE_PAGED_ATTENTION_V1_KERNEL_H_ diff --git a/src/native/cuda/nvidia/ops/paged_attention_v1/kernel.h b/src/native/cuda/nvidia/ops/paged_attention_v1/kernel.h new file mode 100644 index 000000000..8553a45ba --- /dev/null +++ b/src/native/cuda/nvidia/ops/paged_attention_v1/kernel.h @@ -0,0 +1,20 @@ +#ifndef INFINI_OPS_NVIDIA_PAGED_ATTENTION_V1_KERNEL_H_ +#define INFINI_OPS_NVIDIA_PAGED_ATTENTION_V1_KERNEL_H_ + +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/runtime_.h" +#include "native/cuda/ops/paged_attention_v1/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaPagedAttentionV1> { + public: + using CudaPagedAttentionV1< + Runtime>::CudaPagedAttentionV1; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_PAGED_ATTENTION_V1_KERNEL_H_ diff --git a/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh b/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh index 7526ec874..1e8c1a74f 100644 --- a/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh +++ b/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh @@ -2251,140 +2251,6 @@ __global__ void PagedAttentionInfinilmSplitKvCombineKernel( outstride); } -template -__global__ void PagedAttentionInfinilmDecodeWarpKernel( - TData* __restrict__ out, const TData* __restrict__ q, - const TData* __restrict__ k_cache, const TData* __restrict__ v_cache, - const TIndex* __restrict__ block_tables, const TIndex* __restrict__ seqlens, - const float* __restrict__ alibi_slopes, std::size_t num_heads, - std::size_t num_kv_heads, float scale, std::size_t max_num_blocks_per_seq, - std::size_t block_size, std::ptrdiff_t k_cacheblock_stride, - std::ptrdiff_t k_cachehead_stride, std::ptrdiff_t k_cacheslot_stride, - std::ptrdiff_t v_cacheblock_stride, std::ptrdiff_t v_cachehead_stride, - std::ptrdiff_t v_cacheslot_stride, std::ptrdiff_t qstride, - std::ptrdiff_t qhead_stride, std::ptrdiff_t outstride, - std::ptrdiff_t outhead_stride, std::ptrdiff_t block_table_batch_stride, - std::ptrdiff_t seqlens_stride) { - constexpr int kWarpSize = 32; - static_assert(kHeadSize == 64 || kHeadSize == 128, - "PagedAttentionInfinilm decode supports head sizes 64 and 128"); - static_assert(kHeadSize % kWarpSize == 0, - "head size must be divisible by 32"); - - const int seqidx = blockIdx.y; - const int head_idx = blockIdx.x; - const int lane = threadIdx.x; - constexpr int kDimsPerThread = kHeadSize / kWarpSize; - constexpr float kLog2e = 1.4426950408889634f; - - __shared__ float reduce_buf[kWarpSize]; - __shared__ float state_buf[2]; - - const int seqlen = static_cast(seqlens[seqidx * seqlens_stride]); - TData* outptr = out + seqidx * outstride + head_idx * outhead_stride; - if (seqlen <= 0) { -#pragma unroll - for (int i = 0; i < kDimsPerThread; ++i) { - outptr[lane * kDimsPerThread + i] = static_cast(0.0f); - } - return; - } - - const int queries_per_kv = static_cast(num_heads / num_kv_heads); - const int kv_head_idx = head_idx / queries_per_kv; - const float alibi_slope = - alibi_slopes == nullptr ? 0.0f : alibi_slopes[head_idx]; - const float scale_log2 = scale * kLog2e; - const TIndex* block_table = block_tables + seqidx * block_table_batch_stride; - const TData* qptr = q + seqidx * qstride + head_idx * qhead_stride; - - float qreg[kDimsPerThread]; - float acc[kDimsPerThread]; -#pragma unroll - for (int i = 0; i < kDimsPerThread; ++i) { - const int dim = lane * kDimsPerThread + i; - qreg[i] = static_cast(qptr[dim]); - acc[i] = 0.0f; - } - - float m = -FLT_MAX; - float l = 0.0f; - const int page_block_size = static_cast(block_size); - int t_base = 0; - for (int logical_block = 0; - t_base < seqlen && - logical_block < static_cast(max_num_blocks_per_seq); - ++logical_block, t_base += page_block_size) { - const int physical_block = static_cast(block_table[logical_block]); - const TData* k_base = k_cache + physical_block * k_cacheblock_stride + - kv_head_idx * k_cachehead_stride; - const TData* v_base = v_cache + physical_block * v_cacheblock_stride + - kv_head_idx * v_cachehead_stride; - const int token_end = min(page_block_size, seqlen - t_base); - - for (int token_in_block = 0; token_in_block < token_end; ++token_in_block) { - const int token_idx = t_base + token_in_block; - const TData* k_ptr = k_base + token_in_block * k_cacheslot_stride; - const TData* v_ptr = v_base + token_in_block * v_cacheslot_stride; - - float qk = 0.0f; -#pragma unroll - for (int i = 0; i < kDimsPerThread; ++i) { - const int dim = lane * kDimsPerThread + i; - qk += qreg[i] * static_cast(k_ptr[dim]); - } - - reduce_buf[lane] = qk; - __syncthreads(); - for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { - if (lane < offset) { - reduce_buf[lane] += reduce_buf[lane + offset]; - } - __syncthreads(); - } - - float alpha = 1.0f; - float beta = 0.0f; - if (lane == 0) { - float score = reduce_buf[0] * scale_log2; - if (alibi_slope != 0.0f) { - score += - (alibi_slope * static_cast(token_idx - (seqlen - 1))) * - kLog2e; - } - const float m_new = fmaxf(m, score); - alpha = exp2f(m - m_new); - beta = exp2f(score - m_new); - l = l * alpha + beta; - m = m_new; - state_buf[0] = alpha; - state_buf[1] = beta; - } - __syncthreads(); - alpha = state_buf[0]; - beta = state_buf[1]; - -#pragma unroll - for (int i = 0; i < kDimsPerThread; ++i) { - const int dim = lane * kDimsPerThread + i; - acc[i] = acc[i] * alpha + beta * static_cast(v_ptr[dim]); - } - } - } - - if (lane == 0) { - state_buf[0] = 1.0f / (l + 1e-6f); - } - __syncthreads(); - const float inv_l = state_buf[0]; - -#pragma unroll - for (int i = 0; i < kDimsPerThread; ++i) { - const int dim = lane * kDimsPerThread + i; - outptr[dim] = static_cast(acc[i] * inv_l); - } -} - } // namespace infini::ops #endif diff --git a/src/native/cuda/ops/paged_attention_infinilm/kernel.h b/src/native/cuda/ops/paged_attention_infinilm/kernel.h index 65e5f9e1e..6c67b8165 100644 --- a/src/native/cuda/ops/paged_attention_infinilm/kernel.h +++ b/src/native/cuda/ops/paged_attention_infinilm/kernel.h @@ -11,6 +11,7 @@ #include "dispatcher.h" #include "native/cuda/kernel_commons.cuh" #include "native/cuda/ops/paged_attention_infinilm/kernel.cuh" +#include "native/cuda/ops/paged_attention_v1/kernel.cuh" #include "native/cuda/runtime_utils.h" namespace infini::ops { @@ -111,7 +112,7 @@ class CudaPagedAttentionInfinilm : public PagedAttentionInfinilm { reinterpret_cast(out.data()), partial_acc, partial_m, partial_l, num_splits, out_stride_); } else { - PagedAttentionInfinilmDecodeWarpKernel + PagedAttentionDecodeWarpKernel <<>>( reinterpret_cast(out.data()), reinterpret_cast(q.data()), @@ -124,13 +125,14 @@ class CudaPagedAttentionInfinilm : public PagedAttentionInfinilm { : nullptr, num_heads_, num_kv_heads_, scale, max_num_blocks_per_seq_, block_size_, k_cache_block_stride_, k_cache_head_stride_, - k_cache_slot_stride_, v_cache_block_stride_, - v_cache_head_stride_, v_cache_slot_stride_, q_stride_, - q_head_stride_, out_stride_, out_head_stride_, - block_table_batch_stride_, seq_lens_stride_); + k_cache_slot_stride_, 0, 1, static_cast(head_size_), + v_cache_block_stride_, v_cache_head_stride_, + v_cache_slot_stride_, 1, q_stride_, q_head_stride_, + out_stride_, out_head_stride_, block_table_batch_stride_, + seq_lens_stride_); } } else { - PagedAttentionInfinilmDecodeWarpKernel + PagedAttentionDecodeWarpKernel <<>>( reinterpret_cast(out.data()), reinterpret_cast(q.data()), @@ -143,10 +145,11 @@ class CudaPagedAttentionInfinilm : public PagedAttentionInfinilm { : nullptr, num_heads_, num_kv_heads_, scale, max_num_blocks_per_seq_, block_size_, k_cache_block_stride_, k_cache_head_stride_, - k_cache_slot_stride_, v_cache_block_stride_, - v_cache_head_stride_, v_cache_slot_stride_, q_stride_, - q_head_stride_, out_stride_, out_head_stride_, - block_table_batch_stride_, seq_lens_stride_); + k_cache_slot_stride_, 0, 1, static_cast(head_size_), + v_cache_block_stride_, v_cache_head_stride_, + v_cache_slot_stride_, 1, q_stride_, q_head_stride_, + out_stride_, out_head_stride_, block_table_batch_stride_, + seq_lens_stride_); } }, "CudaPagedAttentionInfinilm::operator()"); diff --git a/src/native/cuda/ops/paged_attention_v1/kernel.cuh b/src/native/cuda/ops/paged_attention_v1/kernel.cuh new file mode 100644 index 000000000..6c61d95fe --- /dev/null +++ b/src/native/cuda/ops/paged_attention_v1/kernel.cuh @@ -0,0 +1,151 @@ +#ifndef INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_CUH_ +#define INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_CUH_ + +#include +#include +#include + +namespace infini::ops { + +template +__global__ void PagedAttentionDecodeWarpKernel( + TData* __restrict__ out, const TData* __restrict__ q, + const TData* __restrict__ k_cache, const TData* __restrict__ v_cache, + const TIndex* __restrict__ block_tables, const TIndex* __restrict__ seqlens, + const float* __restrict__ alibi_slopes, std::size_t num_heads, + std::size_t num_kv_heads, float scale, std::size_t max_num_blocks_per_seq, + std::size_t block_size, std::ptrdiff_t k_cacheblock_stride, + std::ptrdiff_t k_cachehead_stride, std::ptrdiff_t k_cacheslot_stride, + std::ptrdiff_t k_cachedim_stride, std::ptrdiff_t k_cachex_stride, + int k_cachex, std::ptrdiff_t v_cacheblock_stride, + std::ptrdiff_t v_cachehead_stride, std::ptrdiff_t v_cacheslot_stride, + std::ptrdiff_t v_cachedim_stride, std::ptrdiff_t qstride, + std::ptrdiff_t qhead_stride, std::ptrdiff_t outstride, + std::ptrdiff_t outhead_stride, std::ptrdiff_t block_table_batch_stride, + std::ptrdiff_t seqlens_stride) { + constexpr int kWarpSize = 32; + static_assert(kHeadSize == 64 || kHeadSize == 128, + "Paged attention decode supports head sizes 64 and 128"); + static_assert(kHeadSize % kWarpSize == 0, + "head size must be divisible by 32"); + + const int seqidx = blockIdx.y; + const int head_idx = blockIdx.x; + const int lane = threadIdx.x; + constexpr int kDimsPerThread = kHeadSize / kWarpSize; + constexpr float kLog2e = 1.4426950408889634f; + + __shared__ float reduce_buf[kWarpSize]; + __shared__ float state_buf[2]; + + const int seqlen = static_cast(seqlens[seqidx * seqlens_stride]); + TData* outptr = out + seqidx * outstride + head_idx * outhead_stride; + if (seqlen <= 0) { +#pragma unroll + for (int i = 0; i < kDimsPerThread; ++i) { + outptr[lane * kDimsPerThread + i] = static_cast(0.0f); + } + return; + } + + const int queries_per_kv = static_cast(num_heads / num_kv_heads); + const int kv_head_idx = head_idx / queries_per_kv; + const float alibi_slope = + alibi_slopes == nullptr ? 0.0f : alibi_slopes[head_idx]; + const float scale_log2 = scale * kLog2e; + const TIndex* block_table = block_tables + seqidx * block_table_batch_stride; + const TData* qptr = q + seqidx * qstride + head_idx * qhead_stride; + + float qreg[kDimsPerThread]; + float acc[kDimsPerThread]; +#pragma unroll + for (int i = 0; i < kDimsPerThread; ++i) { + const int dim = lane * kDimsPerThread + i; + qreg[i] = static_cast(qptr[dim]); + acc[i] = 0.0f; + } + + float m = -FLT_MAX; + float l = 0.0f; + const int page_block_size = static_cast(block_size); + int t_base = 0; + for (int logical_block = 0; + t_base < seqlen && + logical_block < static_cast(max_num_blocks_per_seq); + ++logical_block, t_base += page_block_size) { + const int physical_block = static_cast(block_table[logical_block]); + const TData* k_base = k_cache + physical_block * k_cacheblock_stride + + kv_head_idx * k_cachehead_stride; + const TData* v_base = v_cache + physical_block * v_cacheblock_stride + + kv_head_idx * v_cachehead_stride; + const int token_end = min(page_block_size, seqlen - t_base); + + for (int token_in_block = 0; token_in_block < token_end; ++token_in_block) { + const int token_idx = t_base + token_in_block; + const TData* k_ptr = k_base + token_in_block * k_cacheslot_stride; + const TData* v_ptr = v_base + token_in_block * v_cacheslot_stride; + + float qk = 0.0f; +#pragma unroll + for (int i = 0; i < kDimsPerThread; ++i) { + const int dim = lane * kDimsPerThread + i; + const auto k_offset = (dim / k_cachex) * k_cachedim_stride + + (dim % k_cachex) * k_cachex_stride; + qk += qreg[i] * static_cast(k_ptr[k_offset]); + } + + reduce_buf[lane] = qk; + __syncthreads(); + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + if (lane < offset) { + reduce_buf[lane] += reduce_buf[lane + offset]; + } + __syncthreads(); + } + + float alpha = 1.0f; + float beta = 0.0f; + if (lane == 0) { + float score = reduce_buf[0] * scale_log2; + if (alibi_slope != 0.0f) { + score += + (alibi_slope * static_cast(token_idx - (seqlen - 1))) * + kLog2e; + } + const float m_new = fmaxf(m, score); + alpha = exp2f(m - m_new); + beta = exp2f(score - m_new); + l = l * alpha + beta; + m = m_new; + state_buf[0] = alpha; + state_buf[1] = beta; + } + __syncthreads(); + alpha = state_buf[0]; + beta = state_buf[1]; + +#pragma unroll + for (int i = 0; i < kDimsPerThread; ++i) { + const int dim = lane * kDimsPerThread + i; + acc[i] = acc[i] * alpha + + beta * static_cast(v_ptr[dim * v_cachedim_stride]); + } + } + } + + if (lane == 0) { + state_buf[0] = 1.0f / (l + 1e-6f); + } + __syncthreads(); + const float inv_l = state_buf[0]; + +#pragma unroll + for (int i = 0; i < kDimsPerThread; ++i) { + const int dim = lane * kDimsPerThread + i; + outptr[dim] = static_cast(acc[i] * inv_l); + } +} + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_CUH_ diff --git a/src/native/cuda/ops/paged_attention_v1/kernel.h b/src/native/cuda/ops/paged_attention_v1/kernel.h new file mode 100644 index 000000000..52fdbf6e1 --- /dev/null +++ b/src/native/cuda/ops/paged_attention_v1/kernel.h @@ -0,0 +1,96 @@ +#ifndef INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_H_ +#define INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/paged_attention_v1.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/kernel_commons.cuh" +#include "native/cuda/ops/paged_attention_v1/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +using PagedAttentionV1IndexTypes = + List; + +template +class CudaPagedAttentionV1 : public PagedAttentionV1 { + public: + using PagedAttentionV1::PagedAttentionV1; + + void operator()(const Tensor query, const Tensor key_cache, + const Tensor value_cache, const Tensor block_tables, + const Tensor seq_lens, + const std::optional alibi_slopes, + const int64_t num_kv_heads, const double scale, + const int64_t block_size, const int64_t max_seq_len, + const std::string kv_cache_dtype, const double k_scale, + const double v_scale, const int64_t tp_rank, + const int64_t blocksparse_local_blocks, + const int64_t blocksparse_vert_stride, + const int64_t blocksparse_block_size, + const int64_t blocksparse_head_sliding_step, + Tensor out) const override { + auto cuda_stream = + static_cast(stream_ ? stream_ : 0); + + assert(query.dtype() == dtype_ && key_cache.dtype() == dtype_ && + value_cache.dtype() == dtype_ && out.dtype() == dtype_); + assert(block_tables.dtype() == index_dtype_ && + seq_lens.dtype() == index_dtype_); + assert(num_kv_heads == static_cast(num_kv_heads_)); + assert(scale == scale_); + assert(block_size == static_cast(block_size_)); + assert(max_seq_len > 0); + assert(kv_cache_dtype == "auto" && k_scale == 1.0 && v_scale == 1.0); + assert(blocksparse_local_blocks == 0 && blocksparse_vert_stride == 0 && + blocksparse_head_sliding_step == 0); + assert(blocksparse_block_size > 0); + + (void)tp_rank; + + dim3 grid(static_cast(num_heads_), + static_cast(num_seqs_)); + + DispatchFunc>( + {static_cast(dtype_), static_cast(index_dtype_), + static_cast(head_size_)}, + [&](auto list_tag) { + using TData = TypeMapType(list_tag)>; + using TIndex = + TypeMapType(list_tag)>; + constexpr int kHeadSize = ListGet<2>(list_tag); + + PagedAttentionDecodeWarpKernel + <<>>( + reinterpret_cast(out.data()), + reinterpret_cast(query.data()), + reinterpret_cast(key_cache.data()), + reinterpret_cast(value_cache.data()), + reinterpret_cast(block_tables.data()), + reinterpret_cast(seq_lens.data()), + alibi_slopes.has_value() + ? reinterpret_cast(alibi_slopes->data()) + : nullptr, + num_heads_, num_kv_heads_, static_cast(scale), + max_num_blocks_per_seq_, block_size_, key_cache_block_stride_, + key_cache_head_stride_, key_cache_slot_stride_, + key_cache_dim_stride_, key_cache_x_stride_, + static_cast(key_cache_x_), value_cache_block_stride_, + value_cache_head_stride_, value_cache_slot_stride_, + value_cache_dim_stride_, query_stride_, query_head_stride_, + out_stride_, out_head_stride_, block_table_batch_stride_, + seq_lens_stride_); + }, + "CudaPagedAttentionV1::operator()"); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_PAGED_ATTENTION_V1_KERNEL_H_ diff --git a/tests/test_paged_attention_v1.py b/tests/test_paged_attention_v1.py new file mode 100644 index 000000000..00d97e358 --- /dev/null +++ b/tests/test_paged_attention_v1.py @@ -0,0 +1,225 @@ +import math + +import infini.ops +import pytest +import torch + +from tests.utils import Payload, get_stream + + +def _get_alibi_slopes(num_heads): + closest_power_of_2 = 2 ** math.floor(math.log2(num_heads)) + base = 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))) + powers = [base**i for i in range(1, closest_power_of_2 + 1)] + if num_heads > closest_power_of_2: + powers += [ + base ** (i * 2) + for i in range(1, 2 * (num_heads - closest_power_of_2) + 1, 2) + ] + + return powers[:num_heads] + + +def _reference(query, key_cache, value_cache, block_tables, seq_lens, alibi, scale): + output = torch.empty_like(query) + num_heads = query.shape[1] + num_kv_heads = key_cache.shape[1] + queries_per_kv = num_heads // num_kv_heads + block_size = key_cache.shape[3] + + for seq_id in range(query.shape[0]): + seq_len = seq_lens[seq_id].item() + keys = [] + values = [] + for token_idx in range(seq_len): + block_id = block_tables[seq_id, token_idx // block_size].item() + block_offset = token_idx % block_size + keys.append( + key_cache[block_id, :, :, block_offset, :].reshape(num_kv_heads, -1) + ) + values.append(value_cache[block_id, :, :, block_offset]) + + key = torch.stack(keys) + value = torch.stack(values) + if queries_per_kv > 1: + key = torch.repeat_interleave(key, queries_per_kv, dim=1) + value = torch.repeat_interleave(value, queries_per_kv, dim=1) + + scores = torch.einsum("hd,khd->hk", query[seq_id], key).float() * scale + if alibi is not None: + positions = torch.arange(seq_len, device=query.device) + scores += alibi.view(-1, 1) * (positions - seq_len + 1) + + weights = torch.softmax(scores, dim=-1).to(query.dtype) + output[seq_id] = torch.einsum("hk,khd->hd", weights, value) + + return output + + +@pytest.mark.auto_act_and_assert +@pytest.mark.parametrize( + ( + "num_seqs", + "num_heads", + "num_kv_heads", + "head_size", + "block_size", + "max_seq_len", + "use_alibi", + ), + ( + (1, 1, 1, 64, 16, 128, False), + (3, 8, 2, 128, 16, 256, False), + (2, 4, 2, 64, 8, 64, True), + ), +) +@pytest.mark.parametrize("index_dtype", (torch.int32, torch.int64)) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float16, 1e-2, 1e-3), + (torch.bfloat16, 5e-2, 5e-3), + ), +) +def test_paged_attention_v1( + num_seqs, + num_heads, + num_kv_heads, + head_size, + block_size, + max_seq_len, + use_alibi, + index_dtype, + implementation_index, + dtype, + device, + rtol, + atol, +): + scale = head_size**-0.5 + max_blocks_per_seq = math.ceil(max_seq_len / block_size) + num_blocks = num_seqs * max_blocks_per_seq + key_cache_x = 16 // torch.empty((), dtype=dtype).element_size() + + query = torch.randn((num_seqs, num_heads, head_size), dtype=dtype, device=device) + key_cache = torch.randn( + ( + num_blocks, + num_kv_heads, + head_size // key_cache_x, + block_size, + key_cache_x, + ), + dtype=dtype, + device=device, + ) + value_cache = torch.randn( + (num_blocks, num_kv_heads, head_size, block_size), + dtype=dtype, + device=device, + ) + block_tables = torch.arange(num_blocks, dtype=index_dtype, device=device).view( + num_seqs, max_blocks_per_seq + ) + seq_lens = torch.randint( + 1, max_seq_len + 1, (num_seqs,), dtype=index_dtype, device=device + ) + alibi = ( + torch.tensor(_get_alibi_slopes(num_heads), dtype=torch.float32, device=device) + if use_alibi + else None + ) + out = torch.empty_like(query) + + args = ( + query, + key_cache, + value_cache, + block_tables, + seq_lens, + alibi, + num_kv_heads, + scale, + block_size, + max_seq_len, + "auto", + 1.0, + 1.0, + 0, + 0, + 0, + 64, + 0, + ) + + return Payload( + lambda *call_args, **kwargs: _paged_attention_v1( + *call_args, + **kwargs, + implementation_index=implementation_index, + ), + _torch_paged_attention_v1, + args, + {"out": out}, + rtol=rtol, + atol=atol, + ) + + +def _paged_attention_v1(*args, out, implementation_index): + infini.ops.paged_attention_v1( + *args, + out, + implementation_index=implementation_index, + stream=get_stream(args[0].device), + ) + + return out + + +def _torch_paged_attention_v1( + query, + key_cache, + value_cache, + block_tables, + seq_lens, + alibi, + num_kv_heads, + scale, + block_size, + max_seq_len, + kv_cache_dtype, + k_scale, + v_scale, + tp_rank, + blocksparse_local_blocks, + blocksparse_vert_stride, + blocksparse_block_size, + blocksparse_head_sliding_step, + *, + out, +): + assert num_kv_heads == key_cache.shape[1] + assert block_size == key_cache.shape[3] + assert max_seq_len <= block_tables.shape[1] * block_size + assert kv_cache_dtype == "auto" + assert k_scale == 1.0 and v_scale == 1.0 + assert tp_rank == 0 + assert blocksparse_local_blocks == 0 + assert blocksparse_vert_stride == 0 + assert blocksparse_block_size == 64 + assert blocksparse_head_sliding_step == 0 + + out.copy_( + _reference( + query, + key_cache, + value_cache, + block_tables, + seq_lens, + alibi, + scale, + ) + ) + + return out