issue/847 paged attention prefill一段式接口

include/infinicore/common/hash.hpp

@@ -2,6 +2,7 @@
 
 #include "../tensor.hpp"
 
+#include <optional>
 #include <type_traits>
 
 namespace infinicore {
@@ -24,6 +25,15 @@ inline void hash_combine(size_t &seed, Tensor tensor) {
     }
 }
 
+// Specialization for optional
+template <typename T>
+inline void hash_combine(size_t &seed, const std::optional<T> &opt) {
+    hash_combine(seed, opt.has_value());
+    if (opt) {
+        hash_combine(seed, *opt);
+    }
+}
+
 // Specialization for std::string
 inline void hash_combine(size_t &seed, const std::string &str) {
     hash_combine(seed, std::hash<std::string>{}(str));

include/infinicore/ops.hpp

@@ -6,6 +6,7 @@
 #include "ops/matmul.hpp"
 #include "ops/ones.hpp"
 #include "ops/paged_attention.hpp"
+#include "ops/paged_attention_prefill.hpp"
 #include "ops/paged_caching.hpp"
 #include "ops/random_sample.hpp"
 #include "ops/rearrange.hpp"

include/infinicore/ops/paged_attention_prefill.hpp

+#pragma once
+
+#include "../device.hpp"
+#include "common/op.hpp"
+#include <optional>
+
+namespace infinicore::op {
+
+class PagedAttentionPrefill {
+public:
+    using schema = void (*)(Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, std::optional<Tensor>, float);
+    static void execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float);
+    static common::OpDispatcher<schema> &dispatcher();
+};
+
+Tensor paged_attention_prefill(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale);
+void paged_attention_prefill_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale);
+} // namespace infinicore::op

python/infinicore/__init__.py

@@ -45,6 +45,7 @@ from infinicore.ops.matmul import matmul
 from infinicore.ops.mul import mul
 from infinicore.ops.narrow import narrow
 from infinicore.ops.paged_attention import paged_attention
+from infinicore.ops.paged_attention_prefill import paged_attention_prefill
 from infinicore.ops.paged_caching import paged_caching
 from infinicore.ops.rearrange import rearrange
 from infinicore.ops.squeeze import squeeze
@@ -119,6 +120,7 @@ __all__ = [
     "from_torch",
     "paged_caching",
     "paged_attention",
+    "paged_attention_prefill",
     "ones",
     "strided_empty",
     "strided_from_blob",

python/infinicore/ops/paged_attention_prefill.py

+from infinicore.lib import _infinicore
+from infinicore.tensor import Tensor
+
+
+def paged_attention_prefill(
+    q: Tensor,
+    k_cache: Tensor,
+    v_cache: Tensor,
+    block_tables: Tensor,
+    cache_lens: Tensor,
+    seq_lens: Tensor,
+    seq_offsets: Tensor,
+    alibi_slopes: Tensor | None = None,
+    scale: float = 1.0,
+    *,
+    out: Tensor | None = None,
+):
+    if out is None:
+        return Tensor(
+            _infinicore.paged_attention_prefill(
+                q._underlying,
+                k_cache._underlying,
+                v_cache._underlying,
+                block_tables._underlying,
+                cache_lens._underlying,
+                seq_lens._underlying,
+                seq_offsets._underlying,
+                alibi_slopes._underlying if alibi_slopes is not None else None,
+                scale,
+            )
+        )
+
+    _infinicore.paged_attention_prefill_(
+        out._underlying,
+        q._underlying,
+        k_cache._underlying,
+        v_cache._underlying,
+        block_tables._underlying,
+        cache_lens._underlying,
+        seq_lens._underlying,
+        seq_offsets._underlying,
+        alibi_slopes._underlying if alibi_slopes is not None else None,
+        scale,
+    )
+
+    return out

src/infinicore/ops/paged_attention/paged_attention_infiniop.cc

@@ -16,7 +16,7 @@ thread_local common::OpCache<size_t, infiniopPagedAttentionDescriptor_t> caches(
     });
 
 void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, std::optional<Tensor> alibi_slopes, float scale) {
-    size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens);
+    size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens, alibi_slopes, scale);
 
     auto device = context::getDevice();
     auto &cache = caches.getCache(device);

src/infinicore/ops/paged_attention_prefill/paged_attention_prefill.cc

+#include "infinicore/ops/paged_attention_prefill.hpp"
+
+#include "../../utils.hpp"
+
+namespace infinicore::op {
+
+common::OpDispatcher<PagedAttentionPrefill::schema> &PagedAttentionPrefill::dispatcher() {
+    static common::OpDispatcher<PagedAttentionPrefill::schema> dispatcher_;
+    return dispatcher_;
+};
+
+void PagedAttentionPrefill::execute(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale) {
+    INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, block_tables, cache_lens);
+    infinicore::context::setDevice(out->device());
+    dispatcher().lookup(out->device().getType())(out, q, k_cache, v_cache, block_tables, cache_lens, seq_lens, seq_offsets, alibi_slopes, scale);
+}
+
+Tensor paged_attention_prefill(Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale) {
+    auto out = Tensor::empty(q->shape(), q->dtype(), q->device());
+    paged_attention_prefill_(out, q, k_cache, v_cache, block_tables, cache_lens, seq_lens, seq_offsets, alibi_slopes, scale);
+    return out;
+}
+
+void paged_attention_prefill_(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale) {
+    PagedAttentionPrefill::execute(out, q, k_cache, v_cache, block_tables, cache_lens, seq_lens, seq_offsets, alibi_slopes, scale);
+}
+
+} // namespace infinicore::op

src/infinicore/ops/paged_attention_prefill/paged_attention_prefill_infiniop.cc

+#include "../../utils.hpp"
+#include "infinicore/common/hash.hpp"
+#include "infinicore/ops/common/cache.hpp"
+#include "infinicore/ops/paged_attention_prefill.hpp"
+#include <infiniop.h>
+
+namespace infinicore::op::paged_attention_prefill_impl::infiniop {
+
+thread_local common::OpCache<size_t, infiniopPagedAttentionPrefillDescriptor_t> caches(
+    100, // capacity
+    [](infiniopPagedAttentionPrefillDescriptor_t &desc) {
+        if (desc != nullptr) {
+            INFINICORE_CHECK_ERROR(infiniopDestroyPagedAttentionPrefillDescriptor(desc));
+            desc = nullptr;
+        }
+    });
+
+void calculate(Tensor out, Tensor q, Tensor k_cache, Tensor v_cache, Tensor block_tables, Tensor cache_lens, Tensor seq_lens, Tensor seq_offsets, std::optional<Tensor> alibi_slopes, float scale) {
+    size_t seed = hash_combine(out, q, k_cache, v_cache, block_tables, cache_lens, seq_lens, seq_offsets, alibi_slopes, scale);
+
+    auto device = context::getDevice();
+    auto &cache = caches.getCache(device);
+
+    auto desc_opt = cache.get(seed);
+    infiniopPagedAttentionPrefillDescriptor_t desc = nullptr;
+
+    if (!desc_opt) {
+        INFINICORE_CHECK_ERROR(infiniopCreatePagedAttentionPrefillDescriptor(
+            context::getInfiniopHandle(device), &desc,
+            out->desc(), q->desc(), k_cache->desc(), v_cache->desc(), block_tables->desc(),
+            cache_lens->desc(), seq_lens->desc(), seq_offsets->desc(),
+            alibi_slopes.has_value() ? alibi_slopes.value()->desc() : nullptr,
+            scale));
+        cache.put(seed, desc);
+    } else {
+        desc = *desc_opt;
+    }
+
+    size_t workspace_size = 0;
+    INFINICORE_CHECK_ERROR(infiniopGetPagedAttentionPrefillWorkspaceSize(desc, &workspace_size));
+    std::shared_ptr<Memory> workspace = context::allocateMemory(workspace_size);
+
+    INFINICORE_CHECK_ERROR(infiniopPagedAttentionPrefill(
+        desc, workspace->data(), workspace_size,
+        out->data(), q->data(), k_cache->data(), v_cache->data(), block_tables->data(), cache_lens->data(), seq_lens->data(), seq_offsets->data(),
+        alibi_slopes.has_value() ? alibi_slopes.value()->data() : nullptr,
+        context::getStream()));
+}
+
+static bool registered = []() {
+    PagedAttentionPrefill::dispatcher().registerAll(&calculate, false);
+    return true;
+}();
+
+} // namespace infinicore::op::paged_attention_prefill_impl::infiniop

src/infiniop/ops/paged_attention_prefill/cuda/kernel.cuh

@@ -4,10 +4,10 @@
 namespace op::paged_attention_prefill::cuda {
 
 // 辅助函数：二分查找确定当前 global_token_idx 属于哪个 sequence
-__device__ __forceinline__ int find_seq_id(int token_idx, const int64_t *offset, int num_seqs) {
-    int low = 0, high = num_seqs - 1;
+__device__ __forceinline__ size_t find_seq_id(size_t token_idx, const int64_t *offset, size_t num_seqs) {
+    size_t low = 0, high = num_seqs - 1;
     while (low <= high) {
-        int mid = (low + high) >> 1;
+        size_t mid = (low + high) >> 1;
         if (token_idx >= offset[mid] && token_idx < offset[mid + 1]) {
             return mid;
         } else if (token_idx < offset[mid]) {
@@ -32,22 +32,22 @@ __global__ void pagedAttentionPrefillKernel(
     const size_t num_seqs) {
 
     // --- 使用 2D Grid 坐标 ---
-    const int global_token_idx = blockIdx.x; // 展平后的全局 token 索引
-    const int head_idx = blockIdx.y;         // Head 索引
-    const int dim_idx = threadIdx.x;         // Head 内部维度
+    const size_t global_token_idx = blockIdx.x; // 展平后的全局 token 索引
+    const size_t head_idx = blockIdx.y;         // Head 索引
+    const size_t dim_idx = threadIdx.x;         // Head 内部维度
 
     if (dim_idx >= head_size) {
         return;
     }
 
     // --- 通过二分查找 offset 找到所属的 seq_idx ---
-    int seq_idx = find_seq_id(global_token_idx, offset_, num_seqs);
+    size_t seq_idx = find_seq_id(global_token_idx, offset_, num_seqs);
 
     // --- 获取该 Sequence 本次 Prefill 的长度
     const int64_t cur_new_len = seq_lens_[seq_idx];
 
     // --- 该 token 在当前序列中的相对位置
-    int q_token_idx = global_token_idx - offset_[seq_idx];
+    size_t q_token_idx = global_token_idx - offset_[seq_idx];
 
     const Tdata *q_ptr_base = q_ + global_token_idx * num_heads * head_size + head_idx * head_size;
     Tdata *out_ptr = out_ + global_token_idx * num_heads * head_size + head_idx * head_size;
@@ -65,14 +65,14 @@ __global__ void pagedAttentionPrefillKernel(
 
     // Pass 1: 计算 Score 并找最大值
     Tcompute max_score = -FLT_MAX;
-    for (int t = 0; t <= causal_limit; ++t) {
+    for (size_t t = 0; t <= causal_limit; ++t) {
         const int64_t b_idx = t / block_size;
         const int64_t t_off = t % block_size;
         const int64_t physical_block_id = block_table[b_idx];
         const Tdata *k_vec = k_cache_ + physical_block_id * kv_block_stride + kv_head_idx * kv_head_stride + t_off * head_size;
 
         Tcompute score = 0.0f;
-        for (int d = 0; d < head_size; ++d) {
+        for (size_t d = 0; d < head_size; ++d) {
             score += static_cast<Tcompute>(q_ptr_base[d]) * static_cast<Tcompute>(k_vec[d]);
         }
         score *= static_cast<Tcompute>(scale);
@@ -86,14 +86,14 @@ __global__ void pagedAttentionPrefillKernel(
 
     // Pass 2: 计算 Sum of Exp
     Tcompute sum_exp = 0.0f;
-    for (int t = 0; t <= causal_limit; ++t) {
+    for (size_t t = 0; t <= causal_limit; ++t) {
         const int64_t b_idx = t / block_size;
         const int64_t t_off = t % block_size;
         const int64_t physical_block_id = block_table[b_idx];
         const Tdata *k_vec = k_cache_ + physical_block_id * kv_block_stride + kv_head_idx * kv_head_stride + t_off * head_size;
 
         Tcompute score = 0.0f;
-        for (int d = 0; d < head_size; ++d) {
+        for (size_t d = 0; d < head_size; ++d) {
             score += static_cast<Tcompute>(q_ptr_base[d]) * static_cast<Tcompute>(k_vec[d]);
         }
         score *= static_cast<Tcompute>(scale);
@@ -106,14 +106,14 @@ __global__ void pagedAttentionPrefillKernel(
     // Pass 3: 加权求和得到输出
     Tcompute acc = 0.0f;
     Tcompute inv_sum = 1.0f / (sum_exp + 1e-6f);
-    for (int t = 0; t <= causal_limit; ++t) {
+    for (size_t t = 0; t <= causal_limit; ++t) {
         const int64_t b_idx = t / block_size;
         const int64_t t_off = t % block_size;
         const int64_t physical_block_id = block_table[b_idx];
 
         const Tdata *k_vec = k_cache_ + physical_block_id * kv_block_stride + kv_head_idx * kv_head_stride + t_off * head_size;
         Tcompute score = 0.0f;
-        for (int d = 0; d < head_size; ++d) {
+        for (size_t d = 0; d < head_size; ++d) {
             score += static_cast<Tcompute>(q_ptr_base[d]) * static_cast<Tcompute>(k_vec[d]);
         }
         score *= static_cast<Tcompute>(scale);

xmake/nvidia.lua

@@ -55,7 +55,7 @@ target("infiniop-nvidia")
         end
     end
 
-    add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations")
+    add_cuflags("-Xcompiler=-Wno-error=deprecated-declarations", "-Xcompiler=-Wno-error=unused-function")
 
     local arch_opt = get_config("cuda_arch")
     if arch_opt and type(arch_opt) == "string" then