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..c824c63b4 --- /dev/null +++ b/src/native/cuda/moore/ops/flash_attn_with_kvcache/paged.mu @@ -0,0 +1,123 @@ +#include +#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; + std::abort(); +} + +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.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,