diff --git a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h index b83a2a7df..52a054b7a 100644 --- a/src/native/ascend/ops/flash_attn_varlen_func/kernel.h +++ b/src/native/ascend/ops/flash_attn_varlen_func/kernel.h @@ -72,6 +72,30 @@ class Operator 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 shape{ + static_cast(tensor.size(0)), + static_cast(tensor.size(1)), + static_cast(tensor.size(2) * tensor.size(3)), + }; + std::vector strides(shape.size()); + int64_t stride = 1; + for (int64_t index = static_cast(shape.size()) - 1; index >= 0; + --index) { + strides[index] = stride; + stride *= shape[index]; + } + const std::vector storage_shape{stride}; + return aclCreateTensor(shape.data(), static_cast(shape.size()), + ascend::ToAclDtype(tensor.dtype()), strides.data(), + /*storageOffset=*/0, ACL_FORMAT_ND, + storage_shape.data(), + static_cast(storage_shape.size()), + const_cast(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 alibi_slopes, @@ -86,15 +110,18 @@ class Operator 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(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(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 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 {mask_size, mask_size}, ACL_BOOL, attention_mask_data_); } - aclIntArray* MakeCumulativeLengths(const Tensor cu_seqlens, - std::vector& 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& 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 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(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(cumulative ? current + : current - previous); + }); return aclCreateIntArray(lengths.data(), lengths.size()); } diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index d8d15d7d0..c8a8badf4 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/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