fix: handle Ascend paged flash attention layout

src/native/ascend/ops/flash_attn_varlen_func/kernel.h

@@ -72,6 +72,30 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kAscend>
 
   using FlashAttnVarlenFunc::operator();
 
+  // CANN paged attention consumes a BnBsH cache. InfiniOps exposes the same
+  // storage as BnBsND, so flatten the head dimensions in the ACL descriptor.
+  static aclTensor* BuildPagedCacheAclTensor(const Tensor& tensor) {
+    const std::vector<int64_t> shape{
+        static_cast<int64_t>(tensor.size(0)),
+        static_cast<int64_t>(tensor.size(1)),
+        static_cast<int64_t>(tensor.size(2) * tensor.size(3)),
+    };
+    std::vector<int64_t> strides(shape.size());
+    int64_t stride = 1;
+    for (int64_t index = static_cast<int64_t>(shape.size()) - 1; index >= 0;
+         --index) {
+      strides[index] = stride;
+      stride *= shape[index];
+    }
+    const std::vector<int64_t> storage_shape{stride};
+    return aclCreateTensor(shape.data(), static_cast<int64_t>(shape.size()),
+                           ascend::ToAclDtype(tensor.dtype()), strides.data(),
+                           /*storageOffset=*/0, ACL_FORMAT_ND,
+                           storage_shape.data(),
+                           static_cast<int64_t>(storage_shape.size()),
+                           const_cast<void*>(tensor.data()));
+  }
+
   void operator()(const Tensor q, const Tensor k, const Tensor v,
                   const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
                   const std::optional<Tensor> alibi_slopes,
@@ -86,15 +110,18 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kAscend>
     ValidateSupportedOptions(alibi_slopes, return_attn_probs, softmax_lse,
                              s_dmask);
     ValidateTensors(q, k, v, block_table, out);
+    const bool paged = block_table.has_value();
 
     auto stream = static_cast<aclrtStream>(stream_);
     auto actual_seq_lengths =
-        MakeCumulativeLengths(cu_seqlens_q, q_lengths_i64_, stream);
+        MakeSequenceLengths(cu_seqlens_q, q_lengths_i64_, stream,
+                            /*cumulative=*/true);
     auto actual_seq_lengths_kv =
-        MakeCumulativeLengths(cu_seqlens_k, k_lengths_i64_, stream);
+        MakeSequenceLengths(cu_seqlens_k, k_lengths_i64_, stream,
+                            /*cumulative=*/!block_table.has_value());
     auto t_q = q_cache_.get(const_cast<void*>(q.data()));
-    auto t_k = ascend::BuildAclTensor(k);
-    auto t_v = ascend::BuildAclTensor(v);
+    auto t_k = paged ? BuildPagedCacheAclTensor(k) : ascend::BuildAclTensor(k);
+    auto t_v = paged ? BuildPagedCacheAclTensor(v) : ascend::BuildAclTensor(v);
     auto t_out = out_cache_.get(out.data());
     auto t_attention_mask =
         attention_mask_data_ ? attention_mask_cache_.get(attention_mask_data_)
@@ -121,7 +148,6 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kAscend>
       if (next_tokens < 0) next_tokens = max_token_count;
     }
 
-    const bool paged = block_table.has_value();
     auto scale = softmax_scale.value_or(1.0 / std::sqrt(q.size(2)));
     const auto num_key_value_heads = paged ? k.size(2) : k.size(1);
     const auto block_size = paged ? k.size(1) : 0;
@@ -242,9 +268,11 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kAscend>
         {mask_size, mask_size}, ACL_BOOL, attention_mask_data_);
   }
 
-  aclIntArray* MakeCumulativeLengths(const Tensor cu_seqlens,
-                                     std::vector<int64_t>& lengths,
-                                     aclrtStream stream) const {
+  // Dense TND K/V uses cumulative endpoints, while paged K/V uses per-batch
+  // lengths. Query is always dense TND and therefore cumulative.
+  aclIntArray* MakeSequenceLengths(const Tensor cu_seqlens,
+                                   std::vector<int64_t>& lengths,
+                                   aclrtStream stream, bool cumulative) const {
     const auto bytes = cu_i32_host_.size() * sizeof(cu_i32_host_[0]);
     auto ret = aclrtMemcpyAsync(cu_i32_host_.data(), bytes, cu_seqlens.data(),
                                 bytes, ACL_MEMCPY_DEVICE_TO_HOST, stream);
@@ -254,9 +282,12 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kAscend>
     assert(ret == ACL_SUCCESS &&
            "Ascend `FlashAttnVarlenFunc` failed to synchronize lengths");
 
-    std::transform(cu_i32_host_.begin() + 1, cu_i32_host_.end(),
-                   lengths.begin(),
-                   [](int32_t length) { return static_cast<int64_t>(length); });
+    std::transform(
+        cu_i32_host_.begin() + 1, cu_i32_host_.end(), cu_i32_host_.begin(),
+        lengths.begin(), [cumulative](int32_t current, int32_t previous) {
+          return static_cast<int64_t>(cumulative ? current
+                                                 : current - previous);
+        });
     return aclCreateIntArray(lengths.data(), lengths.size());
   }
 

tests/test_flash_attn_varlen_func.py

@@ -210,6 +210,84 @@ def test_flash_attn_varlen_func(
             torch.testing.assert_close(s_dmask, expected_auxiliary[4])
 
 
+@pytest.mark.parametrize("causal", (False, True))
+@pytest.mark.parametrize(
+    "dtype, rtol, atol",
+    (
+        (torch.float16, 2e-3, 2e-3),
+        (torch.bfloat16, 2e-2, 2e-2),
+    ),
+)
+def test_ascend_paged_prefill_follows_nontrivial_block_table(
+    device,
+    implementation_index,
+    causal,
+    dtype,
+    rtol,
+    atol,
+):
+    if device != "npu" or implementation_index != 0:
+        pytest.skip("coverage requires the native Ascend provider")
+
+    q_lens = (128, 256)
+    k_lens = q_lens
+    num_heads = 8
+    num_kv_heads = 2
+    head_dim = 64
+    page_size = 128
+    block_rows = ((2, 0), (1, 3))
+    num_blocks = 4
+
+    q = torch.randn((sum(q_lens), num_heads, head_dim), dtype=dtype, device=device)
+    k = torch.randn(
+        (num_blocks, page_size, num_kv_heads, head_dim),
+        dtype=dtype,
+        device=device,
+    )
+    v = torch.randn_like(k)
+    block_table = torch.tensor(block_rows, dtype=torch.int32, device=device)
+    cu_seqlens_q = _cumulative_lengths(q_lens, device)
+    cu_seqlens_k = _cumulative_lengths(k_lens, device)
+    out = torch.empty_like(q)
+
+    infini.ops.flash_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        None,
+        block_table,
+        max(q_lens),
+        max(k_lens),
+        0.0,
+        None,
+        causal,
+        (-1, -1),
+        0.0,
+        False,
+        False,
+        out,
+        None,
+        None,
+        stream=get_stream(q.device),
+        implementation_index=implementation_index,
+    )
+
+    expected = _reference_varlen_attention(
+        q,
+        k,
+        v,
+        q_lens,
+        k_lens,
+        None,
+        causal,
+        (-1, -1),
+        block_table,
+    )
+    torch.testing.assert_close(out, expected, rtol=rtol, atol=atol)
+
+
 def test_flash_attn_varlen_func_non_default_stream(device, implementation_index):
     if device == "cuda":
         accelerator = torch.cuda