diff --git a/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp b/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp index a915caf6b..a9f4d70e1 100644 --- a/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp +++ b/src/infiniop/ops/paged_attention/ascend/paged_attention_ascend_kernel.cpp @@ -121,15 +121,15 @@ class PagedAttentionKernel { float max_score = -3.4028234663852886e38f; float sum_exp = 0.0f; - for (size_t block_idx = 0; block_idx < _max_num_blocks_per_seq; ++block_idx) { - size_t token_start = block_idx * _page_block_size; + for (size_t page_block_idx = 0; page_block_idx < _max_num_blocks_per_seq; ++page_block_idx) { + size_t token_start = page_block_idx * _page_block_size; if (token_start >= seq_len) { break; } const int64_t physical_block = loadIndex( _block_tables_gm, - static_cast(seq_idx) * _block_table_batch_stride + static_cast(block_idx)); + static_cast(seq_idx) * _block_table_batch_stride + static_cast(page_block_idx)); size_t token_end = token_start + _page_block_size; if (token_end > seq_len) { diff --git a/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp b/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp index e1fce2785..db90b2056 100644 --- a/src/infiniop/ops/paged_attention_prefill/ascend/paged_attention_prefill_ascend_kernel.cpp +++ b/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(seq_idx) * _block_table_batch_stride + static_cast(block_idx)); + static_cast(seq_idx) * _block_table_batch_stride + static_cast(page_block_idx)); return static_cast(physical_block) * _k_batch_stride + static_cast(kv_head_idx) * _k_head_stride + static_cast(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(seq_idx) * _block_table_batch_stride + static_cast(block_idx)); + static_cast(seq_idx) * _block_table_batch_stride + static_cast(page_block_idx)); return static_cast(physical_block) * _v_batch_stride + static_cast(kv_head_idx) * _v_head_stride + static_cast(token_offset) * _v_row_stride;