From f249c9fc62387853a32bd01f00b4afe966a6a1b8 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 2 Sep 2026 12:48:00 +0800 Subject: [PATCH 1/3] feat(moore): support paged flash attention decode --- .../moore/ops/flash_attn_with_kvcache/paged.h | 48 ++++ .../ops/flash_attn_with_kvcache/paged.mu | 121 ++++++++++ tests/test_flash_attn_with_kvcache_moore.py | 225 ++++++++++++++++++ 3 files changed, 394 insertions(+) create mode 100644 src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.h create mode 100644 src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu create mode 100644 tests/test_flash_attn_with_kvcache_moore.py diff --git a/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.h b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.h new file mode 100644 index 000000000..d531556e3 --- /dev/null +++ b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.h @@ -0,0 +1,48 @@ +#ifndef INFINI_OPS_MOORE_FLASH_ATTN_WITH_KVCACHE_PAGED_H_ +#define INFINI_OPS_MOORE_FLASH_ATTN_WITH_KVCACHE_PAGED_H_ + +#include "base/flash_attn_with_kvcache.h" + +namespace infini::ops { + +template <> +class Operator + : public FlashAttnWithKvcache { + public: + using FlashAttnWithKvcache::FlashAttnWithKvcache; + using FlashAttnWithKvcache::operator(); + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override; + + void operator()(const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const override; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_MOORE_FLASH_ATTN_WITH_KVCACHE_PAGED_H_ diff --git a/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu new file mode 100644 index 000000000..6c60cf87d --- /dev/null +++ b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu @@ -0,0 +1,121 @@ +#include +#include + +#include "dispatcher.h" +#include "native/cuda/moore/caster.cuh" +#include "native/cuda/moore/ops/flash_attn_with_kvcache/paged.h" +#include "native/cuda/moore/runtime_.h" +#include "native/cuda/ops/paged_attention_infinilm/kernel.cuh" + +namespace infini::ops { + +void Operator::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, const int64_t cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const { + assert(false && + "Moore paged FlashAttnWithKvcache requires tensor cache_seqlens"); + (void)q; + (void)k_cache; + (void)v_cache; + (void)k; + (void)v; + (void)rotary_cos; + (void)rotary_sin; + (void)cache_seqlens; + (void)cache_batch_idx; + (void)cache_leftpad; + (void)block_table; + (void)alibi_slopes; + (void)softmax_scale; + (void)causal; + (void)window_size; + (void)softcap; + (void)rotary_interleaved; + (void)num_splits; + (void)return_softmax_lse; + (void)out; + (void)softmax_lse; +} + +void Operator::operator()( + const Tensor q, Tensor k_cache, Tensor v_cache, + const std::optional k, const std::optional v, + const std::optional rotary_cos, + const std::optional rotary_sin, + const std::optional cache_seqlens, + const std::optional cache_batch_idx, + const std::optional cache_leftpad, + const std::optional block_table, + const std::optional alibi_slopes, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool rotary_interleaved, const int64_t num_splits, + const bool return_softmax_lse, Tensor out, + std::optional softmax_lse) const { + assert(cache_seqlens.has_value() && block_table.has_value() && + "Moore FlashAttnWithKvcache requires paged tensor metadata"); + assert(!k.has_value() && !v.has_value() && !rotary_cos.has_value() && + !rotary_sin.has_value() && !cache_batch_idx.has_value() && + !cache_leftpad.has_value() && + "Moore FlashAttnWithKvcache supports read-only paged decode"); + assert(q_shape_[1] == 1 && + "Moore FlashAttnWithKvcache supports one decode token per batch"); + assert(causal && window_size[0] == -1 && window_size[1] == -1 && + softcap == 0.0 && num_splits == 0 && + "Moore FlashAttnWithKvcache supports global causal attention"); + assert(!return_softmax_lse && !softmax_lse.has_value() && + "Moore FlashAttnWithKvcache does not return softmax LSE"); + assert((head_size_ == 64 || head_size_ == 128) && + "Moore FlashAttnWithKvcache supports head sizes 64 and 128"); + assert((!alibi_slopes.has_value() || alibi_slopes_shape_.size() == 1) && + "Moore FlashAttnWithKvcache supports one-dimensional ALiBi slopes"); + + (void)rotary_interleaved; + + using Backend = Runtime; + using Index = int32_t; + const auto stream = static_cast(stream_ ? stream_ : 0); + const dim3 grid(static_cast(q_shape_[2]), + static_cast(q_shape_[0])); + const float scale = static_cast( + softmax_scale.value_or(1.0 / std::sqrt(static_cast(head_size_)))); + + DispatchFunc>( + {static_cast(q_dtype_), static_cast(head_size_)}, + [&](auto list_tag) { + using TData = TypeMapType(list_tag)>; + constexpr int kHeadSize = ListGet<1>(list_tag); + + PagedAttentionInfinilmDecodeWarpKernel + <<>>( + reinterpret_cast(out.data()), + reinterpret_cast(q.data()), + reinterpret_cast(k_cache.data()), + reinterpret_cast(v_cache.data()), + reinterpret_cast(block_table->data()), + reinterpret_cast(cache_seqlens->data()), + alibi_slopes.has_value() + ? reinterpret_cast(alibi_slopes->data()) + : nullptr, + q_shape_[2], k_cache_shape_[2], scale, block_table_shape_[1], + k_cache_shape_[1], k_cache_strides_[0], k_cache_strides_[2], + k_cache_strides_[1], v_cache_strides_[0], v_cache_strides_[2], + v_cache_strides_[1], q_strides_[0], q_strides_[2], + out_strides_[0], out_strides_[2], block_table_strides_[0], + cache_seqlens_strides_[0]); + }, + "MooreFlashAttnWithKvcache"); +} + +} // namespace infini::ops diff --git a/tests/test_flash_attn_with_kvcache_moore.py b/tests/test_flash_attn_with_kvcache_moore.py new file mode 100644 index 000000000..963f530b4 --- /dev/null +++ b/tests/test_flash_attn_with_kvcache_moore.py @@ -0,0 +1,225 @@ +import math + +import infini.ops +import pytest +import torch + +from tests.test_flash_attn_varlen_func import _reference_varlen_attention + + +if not hasattr(infini.ops, "FlashAttnWithKvcache"): + pytest.skip( + "`FlashAttnWithKvcache` is not available on this platform", + allow_module_level=True, + ) + + +@pytest.mark.parametrize("head_dim", (64, 128)) +@pytest.mark.parametrize( + "dtype, rtol, atol", + ( + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 2e-2, 2e-2), + ), +) +@pytest.mark.parametrize("num_kv_heads", (2, 4)) +@pytest.mark.parametrize("implementation_index", (8,)) +def test_moore_paged_flash_attn_with_kvcache( + device, implementation_index, num_kv_heads, head_dim, dtype, rtol, atol +): + if device != "musa": + pytest.skip("paged Moore decode requires the Moore backend") + + cache_seqlens = (1, 255, 256, 257) + batch_size = len(cache_seqlens) + num_heads = 4 + page_size = 256 + q = torch.randn( + (batch_size, 1, num_heads, head_dim + 5), dtype=dtype, device=device + )[..., :head_dim] + k_cache = torch.randn( + (5, page_size, num_kv_heads, head_dim + 7), dtype=dtype, device=device + )[..., :head_dim] + v_cache = torch.randn( + (5, page_size, num_kv_heads, head_dim + 11), dtype=dtype, device=device + )[..., :head_dim] + block_table = torch.tensor( + ((3, -1), (1, -1), (4, -1), (2, 0)), + dtype=torch.int32, + device=device, + ) + cache_seqlens_tensor = torch.tensor(cache_seqlens, dtype=torch.int32, device=device) + alibi_slopes = torch.linspace( + 0.01, 0.04, num_heads, dtype=torch.float32, device=device + ) + out = torch.full( + (batch_size, 1, num_heads, head_dim + 13), + math.nan, + dtype=dtype, + device=device, + )[..., :head_dim] + k_before = k_cache.clone() + v_before = v_cache.clone() + + expected = _reference_varlen_attention( + q[:, 0], + k_cache, + v_cache, + (1,) * batch_size, + cache_seqlens, + 0.125, + True, + (-1, -1), + block_table, + alibi_slopes, + ).unsqueeze(1) + + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens_tensor, + None, + None, + block_table, + alibi_slopes, + 0.125, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=torch.musa.current_stream().musa_stream, + implementation_index=implementation_index, + ) + + torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) + torch.testing.assert_close(k_cache, k_before, rtol=0, atol=0) + torch.testing.assert_close(v_cache, v_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("implementation_index", (8,)) +def test_moore_paged_flash_attn_with_kvcache_uses_stream_and_current_lengths( + device, implementation_index +): + if device != "musa": + pytest.skip("paged Moore stream coverage requires the Moore backend") + + head_dim = 64 + q = torch.randn((2, 1, 4, head_dim), dtype=torch.float16, device=device) + k_cache = torch.randn((3, 256, 2, head_dim), dtype=torch.float16, device=device) + v_cache = torch.randn_like(k_cache) + block_table = torch.tensor(((2, -1), (1, 0)), dtype=torch.int32, device=device) + cache_seqlens = torch.tensor((0, 257), dtype=torch.int32, device=device) + out = torch.full_like(q, math.nan) + expected = _reference_varlen_attention( + q[1:, 0], + k_cache, + v_cache, + (1,), + (257,), + None, + True, + (-1, -1), + block_table[1:], + ).cpu() + updated_block_table = torch.tensor( + ((0, -1), (2, -1)), dtype=torch.int32, device=device + ) + updated_cache_seqlens = torch.tensor((1, 256), dtype=torch.int32, device=device) + updated_expected = _reference_varlen_attention( + q[:, 0], + k_cache, + v_cache, + (1, 1), + (1, 256), + None, + True, + (-1, -1), + updated_block_table, + ).cpu() + torch.musa.synchronize() + + stream = torch.musa.Stream() + stream.wait_stream(torch.musa.current_stream()) + torch.musa._sleep(50_000_000) + try: + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens, + None, + None, + block_table, + None, + None, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=stream.musa_stream, + implementation_index=implementation_index, + ) + + stream.synchronize() + with torch.musa.stream(stream): + actual = out.cpu() + torch.testing.assert_close(actual[1, 0], expected[0], rtol=1e-2, atol=1e-2) + torch.testing.assert_close( + actual[0], torch.zeros_like(actual[0]), rtol=0, atol=0 + ) + + with torch.musa.stream(stream): + cache_seqlens.copy_(updated_cache_seqlens) + block_table.copy_(updated_block_table) + out.fill_(math.nan) + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens, + None, + None, + block_table, + None, + None, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=stream.musa_stream, + implementation_index=implementation_index, + ) + + stream.synchronize() + with torch.musa.stream(stream): + updated_actual = out.cpu() + torch.testing.assert_close( + updated_actual[:, 0], updated_expected, rtol=1e-2, atol=1e-2 + ) + finally: + torch.musa.synchronize() From d6048c21e09fc16b1b0cb4e7d611746cd4a9a929 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Wed, 2 Sep 2026 18:00:23 +0800 Subject: [PATCH 2/3] fix(moore): fail closed for scalar cache lengths --- src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu index 6c60cf87d..c824c63b4 100644 --- a/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu +++ b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu @@ -1,5 +1,6 @@ #include #include +#include #include "dispatcher.h" #include "native/cuda/moore/caster.cuh" @@ -46,6 +47,7 @@ void Operator::operator()( (void)return_softmax_lse; (void)out; (void)softmax_lse; + std::abort(); } void Operator::operator()( From 65e2293044121fb2b9f599fae36fb82ca6cea52d Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Fri, 4 Sep 2026 21:09:22 +0800 Subject: [PATCH 3/3] test: share flash attention KV-cache coverage --- tests/test_flash_attn_with_kvcache.py | 233 +++++++++++++++++++- tests/test_flash_attn_with_kvcache_moore.py | 225 ------------------- 2 files changed, 225 insertions(+), 233 deletions(-) delete mode 100644 tests/test_flash_attn_with_kvcache_moore.py diff --git a/tests/test_flash_attn_with_kvcache.py b/tests/test_flash_attn_with_kvcache.py index 366f1c795..3499c2378 100644 --- a/tests/test_flash_attn_with_kvcache.py +++ b/tests/test_flash_attn_with_kvcache.py @@ -7,9 +7,6 @@ from tests.utils import get_stream -flash_attn = pytest.importorskip("flash_attn") - - if not hasattr(infini.ops, "FlashAttnWithKvcache"): pytest.skip( "`FlashAttnWithKvcache` is not available on this platform", @@ -17,6 +14,10 @@ ) +def _get_flash_attn(): + return pytest.importorskip("flash_attn") + + @pytest.mark.parametrize("cache_seqlens_kind", ("tensor", "scalar")) @pytest.mark.parametrize("append_kv", (False, True)) @pytest.mark.parametrize( @@ -79,7 +80,7 @@ def test_flash_attn_with_kvcache_dense( window_size=(4, 0), ) else: - expected, expected_softmax_lse = flash_attn.flash_attn_with_kvcache( + expected, expected_softmax_lse = _get_flash_attn().flash_attn_with_kvcache( q, expected_k_cache, expected_v_cache, @@ -165,7 +166,7 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index): causal=True, ) else: - expected = flash_attn.flash_attn_with_kvcache( + expected = _get_flash_attn().flash_attn_with_kvcache( q, k_cache, v_cache, @@ -225,7 +226,7 @@ def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( cache_batch_idx=cache_batch_idx, ) else: - expected = flash_attn.flash_attn_with_kvcache( + expected = _get_flash_attn().flash_attn_with_kvcache( q, k_cache, v_cache, @@ -273,7 +274,7 @@ def test_flash_attn_with_kvcache_defaults(device, implementation_index): if device == "mlu": expected, _ = _reference_flash_attn_with_kvcache(q, k_cache, v_cache) else: - expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) + expected = _get_flash_attn().flash_attn_with_kvcache(q, k_cache, v_cache) actual = torch.empty_like(q) infini.ops.flash_attn_with_kvcache( @@ -304,7 +305,7 @@ def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index if device == "mlu": expected, _ = _reference_flash_attn_with_kvcache(q, k_cache, v_cache) else: - expected = flash_attn.flash_attn_with_kvcache(q, k_cache, v_cache) + expected = _get_flash_attn().flash_attn_with_kvcache(q, k_cache, v_cache) actual = torch.empty_like(q) stream = accelerator.Stream() stream.wait_stream(accelerator.current_stream()) @@ -322,6 +323,215 @@ def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index torch.testing.assert_close(actual, expected, rtol=2e-3, atol=2e-3) +@pytest.mark.parametrize("head_dim", (64, 128)) +@pytest.mark.parametrize( + "dtype, rtol, atol", + ( + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 2e-2, 2e-2), + ), +) +@pytest.mark.parametrize("num_kv_heads", (2, 4)) +@pytest.mark.parametrize("implementation_index", (8,)) +def test_flash_attn_with_kvcache_paged_moore_decode_matrix( + device, implementation_index, num_kv_heads, head_dim, dtype, rtol, atol +): + if device != "musa": + pytest.skip("paged decode matrix requires the Moore backend") + + cache_seqlens = (1, 255, 256, 257) + batch_size = len(cache_seqlens) + num_heads = 4 + page_size = 256 + q = torch.randn( + (batch_size, 1, num_heads, head_dim + 5), dtype=dtype, device=device + )[..., :head_dim] + k_cache = torch.randn( + (5, page_size, num_kv_heads, head_dim + 7), dtype=dtype, device=device + )[..., :head_dim] + v_cache = torch.randn( + (5, page_size, num_kv_heads, head_dim + 11), dtype=dtype, device=device + )[..., :head_dim] + block_table = torch.tensor( + ((3, -1), (1, -1), (4, -1), (2, 0)), + dtype=torch.int32, + device=device, + ) + cache_seqlens_tensor = torch.tensor(cache_seqlens, dtype=torch.int32, device=device) + alibi_slopes = torch.linspace( + 0.01, 0.04, num_heads, dtype=torch.float32, device=device + ) + out = torch.full( + (batch_size, 1, num_heads, head_dim + 13), + math.nan, + dtype=dtype, + device=device, + )[..., :head_dim] + k_before = k_cache.clone() + v_before = v_cache.clone() + + expected, _ = _reference_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=cache_seqlens_tensor, + block_table=block_table, + softmax_scale=0.125, + causal=True, + alibi_slopes=alibi_slopes, + ) + + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens_tensor, + None, + None, + block_table, + alibi_slopes, + 0.125, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=get_stream(q.device), + implementation_index=implementation_index, + ) + + torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) + torch.testing.assert_close(k_cache, k_before, rtol=0, atol=0) + torch.testing.assert_close(v_cache, v_before, rtol=0, atol=0) + + +@pytest.mark.parametrize("implementation_index", (8,)) +def test_flash_attn_with_kvcache_paged_moore_uses_stream_and_current_metadata( + device, implementation_index +): + if device != "musa": + pytest.skip("paged stream metadata coverage requires the Moore backend") + + accelerator = torch.musa + head_dim = 64 + q = torch.randn((2, 1, 4, head_dim), dtype=torch.float16, device=device) + k_cache = torch.randn((3, 256, 2, head_dim), dtype=torch.float16, device=device) + v_cache = torch.randn_like(k_cache) + block_table = torch.tensor(((2, -1), (1, 0)), dtype=torch.int32, device=device) + cache_seqlens = torch.tensor((0, 257), dtype=torch.int32, device=device) + out = torch.full_like(q, math.nan) + expected, _ = _reference_flash_attn_with_kvcache( + q[1:], + k_cache, + v_cache, + cache_seqlens=cache_seqlens[1:], + block_table=block_table[1:], + causal=True, + ) + expected = expected.cpu() + updated_block_table = torch.tensor( + ((0, -1), (2, -1)), dtype=torch.int32, device=device + ) + updated_cache_seqlens = torch.tensor((1, 256), dtype=torch.int32, device=device) + updated_expected, _ = _reference_flash_attn_with_kvcache( + q, + k_cache, + v_cache, + cache_seqlens=updated_cache_seqlens, + block_table=updated_block_table, + causal=True, + ) + updated_expected = updated_expected.cpu() + accelerator.synchronize() + + stream = accelerator.Stream() + stream.wait_stream(accelerator.current_stream()) + + # Keep the current stream busy after the target dependency is recorded. A + # provider that ignores the raw stream then leaves output untouched at target sync. + accelerator._sleep(50_000_000) + try: + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens, + None, + None, + block_table, + None, + None, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=stream.musa_stream, + implementation_index=implementation_index, + ) + + stream.synchronize() + with accelerator.stream(stream): + actual = out.cpu() + torch.testing.assert_close(actual[1], expected[0], rtol=1e-2, atol=1e-2) + torch.testing.assert_close( + actual[0], torch.zeros_like(actual[0]), rtol=0, atol=0 + ) + + with accelerator.stream(stream): + cache_seqlens.copy_(updated_cache_seqlens) + block_table.copy_(updated_block_table) + out.fill_(math.nan) + infini.ops.flash_attn_with_kvcache( + q, + k_cache, + v_cache, + None, + None, + None, + None, + cache_seqlens, + None, + None, + block_table, + None, + None, + True, + (-1, -1), + 0.0, + True, + 0, + False, + out, + None, + stream=stream.musa_stream, + implementation_index=implementation_index, + ) + + stream.synchronize() + with accelerator.stream(stream): + updated_actual = out.cpu() + torch.testing.assert_close( + updated_actual, updated_expected, rtol=1e-2, atol=1e-2 + ) + finally: + accelerator.synchronize() + + def _reference_flash_attn_with_kvcache( q, k_cache, @@ -334,6 +544,7 @@ def _reference_flash_attn_with_kvcache( softmax_scale=None, causal=False, window_size=(-1, -1), + alibi_slopes=None, ): batch_size, query_length, num_heads, _ = q.shape if cache_seqlens is None: @@ -379,6 +590,12 @@ def _reference_flash_attn_with_kvcache( scale = softmax_scale if softmax_scale is not None else q.size(-1) ** -0.5 scores = torch.matmul(q_seq.float(), k_seq.float().transpose(-2, -1)) scores *= scale + if alibi_slopes is not None: + slopes = alibi_slopes if alibi_slopes.ndim == 1 else alibi_slopes[batch] + query_positions = torch.arange(query_length, device=q.device).unsqueeze(1) + key_positions = torch.arange(length, device=q.device).unsqueeze(0) + distance = (query_positions + length - query_length - key_positions).abs() + scores += -slopes[:, None, None] * distance mask = _attention_mask( query_length, length, diff --git a/tests/test_flash_attn_with_kvcache_moore.py b/tests/test_flash_attn_with_kvcache_moore.py deleted file mode 100644 index 963f530b4..000000000 --- a/tests/test_flash_attn_with_kvcache_moore.py +++ /dev/null @@ -1,225 +0,0 @@ -import math - -import infini.ops -import pytest -import torch - -from tests.test_flash_attn_varlen_func import _reference_varlen_attention - - -if not hasattr(infini.ops, "FlashAttnWithKvcache"): - pytest.skip( - "`FlashAttnWithKvcache` is not available on this platform", - allow_module_level=True, - ) - - -@pytest.mark.parametrize("head_dim", (64, 128)) -@pytest.mark.parametrize( - "dtype, rtol, atol", - ( - (torch.float16, 1e-2, 1e-2), - (torch.bfloat16, 2e-2, 2e-2), - ), -) -@pytest.mark.parametrize("num_kv_heads", (2, 4)) -@pytest.mark.parametrize("implementation_index", (8,)) -def test_moore_paged_flash_attn_with_kvcache( - device, implementation_index, num_kv_heads, head_dim, dtype, rtol, atol -): - if device != "musa": - pytest.skip("paged Moore decode requires the Moore backend") - - cache_seqlens = (1, 255, 256, 257) - batch_size = len(cache_seqlens) - num_heads = 4 - page_size = 256 - q = torch.randn( - (batch_size, 1, num_heads, head_dim + 5), dtype=dtype, device=device - )[..., :head_dim] - k_cache = torch.randn( - (5, page_size, num_kv_heads, head_dim + 7), dtype=dtype, device=device - )[..., :head_dim] - v_cache = torch.randn( - (5, page_size, num_kv_heads, head_dim + 11), dtype=dtype, device=device - )[..., :head_dim] - block_table = torch.tensor( - ((3, -1), (1, -1), (4, -1), (2, 0)), - dtype=torch.int32, - device=device, - ) - cache_seqlens_tensor = torch.tensor(cache_seqlens, dtype=torch.int32, device=device) - alibi_slopes = torch.linspace( - 0.01, 0.04, num_heads, dtype=torch.float32, device=device - ) - out = torch.full( - (batch_size, 1, num_heads, head_dim + 13), - math.nan, - dtype=dtype, - device=device, - )[..., :head_dim] - k_before = k_cache.clone() - v_before = v_cache.clone() - - expected = _reference_varlen_attention( - q[:, 0], - k_cache, - v_cache, - (1,) * batch_size, - cache_seqlens, - 0.125, - True, - (-1, -1), - block_table, - alibi_slopes, - ).unsqueeze(1) - - infini.ops.flash_attn_with_kvcache( - q, - k_cache, - v_cache, - None, - None, - None, - None, - cache_seqlens_tensor, - None, - None, - block_table, - alibi_slopes, - 0.125, - True, - (-1, -1), - 0.0, - True, - 0, - False, - out, - None, - stream=torch.musa.current_stream().musa_stream, - implementation_index=implementation_index, - ) - - torch.testing.assert_close(out, expected, rtol=rtol, atol=atol) - torch.testing.assert_close(k_cache, k_before, rtol=0, atol=0) - torch.testing.assert_close(v_cache, v_before, rtol=0, atol=0) - - -@pytest.mark.parametrize("implementation_index", (8,)) -def test_moore_paged_flash_attn_with_kvcache_uses_stream_and_current_lengths( - device, implementation_index -): - if device != "musa": - pytest.skip("paged Moore stream coverage requires the Moore backend") - - head_dim = 64 - q = torch.randn((2, 1, 4, head_dim), dtype=torch.float16, device=device) - k_cache = torch.randn((3, 256, 2, head_dim), dtype=torch.float16, device=device) - v_cache = torch.randn_like(k_cache) - block_table = torch.tensor(((2, -1), (1, 0)), dtype=torch.int32, device=device) - cache_seqlens = torch.tensor((0, 257), dtype=torch.int32, device=device) - out = torch.full_like(q, math.nan) - expected = _reference_varlen_attention( - q[1:, 0], - k_cache, - v_cache, - (1,), - (257,), - None, - True, - (-1, -1), - block_table[1:], - ).cpu() - updated_block_table = torch.tensor( - ((0, -1), (2, -1)), dtype=torch.int32, device=device - ) - updated_cache_seqlens = torch.tensor((1, 256), dtype=torch.int32, device=device) - updated_expected = _reference_varlen_attention( - q[:, 0], - k_cache, - v_cache, - (1, 1), - (1, 256), - None, - True, - (-1, -1), - updated_block_table, - ).cpu() - torch.musa.synchronize() - - stream = torch.musa.Stream() - stream.wait_stream(torch.musa.current_stream()) - torch.musa._sleep(50_000_000) - try: - infini.ops.flash_attn_with_kvcache( - q, - k_cache, - v_cache, - None, - None, - None, - None, - cache_seqlens, - None, - None, - block_table, - None, - None, - True, - (-1, -1), - 0.0, - True, - 0, - False, - out, - None, - stream=stream.musa_stream, - implementation_index=implementation_index, - ) - - stream.synchronize() - with torch.musa.stream(stream): - actual = out.cpu() - torch.testing.assert_close(actual[1, 0], expected[0], rtol=1e-2, atol=1e-2) - torch.testing.assert_close( - actual[0], torch.zeros_like(actual[0]), rtol=0, atol=0 - ) - - with torch.musa.stream(stream): - cache_seqlens.copy_(updated_cache_seqlens) - block_table.copy_(updated_block_table) - out.fill_(math.nan) - infini.ops.flash_attn_with_kvcache( - q, - k_cache, - v_cache, - None, - None, - None, - None, - cache_seqlens, - None, - None, - block_table, - None, - None, - True, - (-1, -1), - 0.0, - True, - 0, - False, - out, - None, - stream=stream.musa_stream, - implementation_index=implementation_index, - ) - - stream.synchronize() - with torch.musa.stream(stream): - updated_actual = out.cpu() - torch.testing.assert_close( - updated_actual[:, 0], updated_expected, rtol=1e-2, atol=1e-2 - ) - finally: - torch.musa.synchronize()