fix: support ascend paged graph runtime

Synchronize RoPE cache copies before the cache is consumed by Ascend kernels.

Rename paged attention block_idx locals to page_block_idx to avoid AscendC kernel symbol conflicts while preserving the same block table address calculation.

Detect common Ascend driver library layouts before linking libascend_hal.so instead of relying on one hard-coded driver path.

Author: zhuyue

src/infinicore/nn/rope.cc

@@ -1,6 +1,7 @@
 #include "infinicore/nn/rope.hpp"
 #include "../../utils.h"
 #include "../utils.hpp"
+#include "infinicore/context/context.hpp"
 #include "infinicore/ops.hpp"
 #include <algorithm>
 #include <cassert>
@@ -11,6 +12,13 @@
 #include <vector>
 
 namespace infinicore::nn {
+namespace {
+
+void sync_cache_copy() {
+    infinicore::context::syncStream();
+}
+
+} // namespace
 
 RoPE::RoPE(size_t head_dim,
            size_t rotary_dim,
@@ -82,6 +90,7 @@ void RoPE::initialize_cache() {
         auto cos_f32_cpu = Tensor::from_blob(cos_data.data(), {max_seq_len_, cache_dim}, DataType::F32, cpu_device);
         sin_cache_->copy_from(sin_f32_cpu);
         cos_cache_->copy_from(cos_f32_cpu);
+        sync_cache_copy();
     } else if (dtype_ == DataType::BF16) {
         // Convert F32 to BF16 using the same conversion as Python's ml_dtypes.bfloat16
         // This uses round-to-nearest-even (matching _f32_to_bf16 implementation)
@@ -99,6 +108,7 @@ void RoPE::initialize_cache() {
         // copy_from handles cross-device copying to target device
         sin_cache_->copy_from(sin_bf16_cpu);
         cos_cache_->copy_from(cos_bf16_cpu);
+        sync_cache_copy();
     } else if (dtype_ == DataType::F16) {
         // Convert F32 to F16
         std::vector<fp16_t> sin_f16_data(max_seq_len_ * cache_dim);
@@ -114,6 +124,7 @@ void RoPE::initialize_cache() {
 
         sin_cache_->copy_from(sin_f16_cpu);
         cos_cache_->copy_from(cos_f16_cpu);
+        sync_cache_copy();
     } else {
         throw std::runtime_error(
             "RoPE cache dtype conversion not yet supported for dtype: "

src/infiniop/ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp

@@ -164,22 +164,22 @@ class PagedAttentionKernel {
 
 private:
     __aicore__ inline ptrdiff_t keyBase(size_t seq_idx, size_t kv_head_idx, size_t token_idx) {
-        const size_t block_idx = token_idx / _page_block_size;
+        const size_t page_block_idx = token_idx / _page_block_size;
         const size_t token_offset = token_idx % _page_block_size;
         const int64_t physical_block = loadIndex(
             _block_tables_gm,
-            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(block_idx));
+            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(page_block_idx));
         return static_cast<ptrdiff_t>(physical_block) * _k_batch_stride
              + static_cast<ptrdiff_t>(kv_head_idx) * _k_head_stride
              + static_cast<ptrdiff_t>(token_offset) * _k_row_stride;
     }
 
     __aicore__ inline ptrdiff_t valueBase(size_t seq_idx, size_t kv_head_idx, size_t token_idx) {
-        const size_t block_idx = token_idx / _page_block_size;
+        const size_t page_block_idx = token_idx / _page_block_size;
         const size_t token_offset = token_idx % _page_block_size;
         const int64_t physical_block = loadIndex(
             _block_tables_gm,
-            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(block_idx));
+            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(page_block_idx));
         return static_cast<ptrdiff_t>(physical_block) * _v_batch_stride
              + static_cast<ptrdiff_t>(kv_head_idx) * _v_head_stride
              + static_cast<ptrdiff_t>(token_offset) * _v_row_stride;

src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp

@@ -197,22 +197,22 @@ class PagedAttentionPrefillKernel {
     }
 
     __aicore__ inline ptrdiff_t keyBase(size_t seq_idx, size_t kv_head_idx, size_t token_idx) {
-        const size_t block_idx = token_idx / _page_block_size;
+        const size_t page_block_idx = token_idx / _page_block_size;
         const size_t token_offset = token_idx % _page_block_size;
         const int64_t physical_block = loadPrefillIndex(
             _block_tables_gm,
-            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(block_idx));
+            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(page_block_idx));
         return static_cast<ptrdiff_t>(physical_block) * _k_batch_stride
              + static_cast<ptrdiff_t>(kv_head_idx) * _k_head_stride
              + static_cast<ptrdiff_t>(token_offset) * _k_row_stride;
     }
 
     __aicore__ inline ptrdiff_t valueBase(size_t seq_idx, size_t kv_head_idx, size_t token_idx) {
-        const size_t block_idx = token_idx / _page_block_size;
+        const size_t page_block_idx = token_idx / _page_block_size;
         const size_t token_offset = token_idx % _page_block_size;
         const int64_t physical_block = loadPrefillIndex(
             _block_tables_gm,
-            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(block_idx));
+            static_cast<ptrdiff_t>(seq_idx) * _block_table_batch_stride + static_cast<ptrdiff_t>(page_block_idx));
         return static_cast<ptrdiff_t>(physical_block) * _v_batch_stride
              + static_cast<ptrdiff_t>(kv_head_idx) * _v_head_stride
              + static_cast<ptrdiff_t>(token_offset) * _v_row_stride;

xmake/ascend.lua

@@ -10,7 +10,14 @@ add_links("libascendcl.so")
 add_links("libnnopbase.so")
 add_links("libopapi.so")
 add_links("libruntime.so")
-add_linkdirs(ASCEND_HOME .. "/../../driver/lib64/driver")
+for _, driver_dir in ipairs({
+    path.join(ASCEND_HOME, "../driver/lib64/driver"),
+    path.join(ASCEND_HOME, "../../driver/lib64/driver"),
+}) do
+    if os.isdir(driver_dir) then
+        add_linkdirs(driver_dir)
+    end
+end
 add_links("libascend_hal.so")
 local builddir = string.format(
         "%s/build/%s/%s/%s",