From 11cd06890b177cd7816eeb26d7cc7a98dff246c8 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Thu, 27 Aug 2026 08:00:31 +0000 Subject: [PATCH 1/2] fix(iluvatar): avoid unavailable CUDA driver target --- src/CMakeLists.txt | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 323e1303e..20196f657 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -236,7 +236,9 @@ if(WITH_ILUVATAR) target_sources(infiniops PRIVATE ${ILUVATAR_SOURCES}) find_package(CUDAToolkit REQUIRED) - target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas CUDA::cuda_driver) + # CoreX does not provide FindCUDAToolkit's CUDA::cuda_driver target. The + # Iluvatar backend uses the runtime and BLAS APIs directly. + target_link_libraries(infiniops PUBLIC CUDA::cudart CUDA::cublas) list(APPEND DEVICE_LIST "iluvatar") endif() From 572f6b8c5b7c456a092310b22125c45c41cfb377 Mon Sep 17 00:00:00 2001 From: gongchensu Date: Thu, 27 Aug 2026 08:36:45 +0000 Subject: [PATCH 2/2] feat(iluvatar): add native canonical attention providers --- src/CMakeLists.txt | 6 +- .../ops/flash_attn_varlen_func/kernel.h | 95 +++++++++++++++ .../ops/flash_attn_with_kvcache/kernel.h | 110 ++++++++++++++++++ .../ops/reshape_and_cache_flash/kernel.h | 10 +- .../ops/paged_attention_infinilm/kernel.cuh | 15 ++- .../kernel.cuh | 4 +- .../paged_attention_prefill_infinilm/kernel.h | 18 ++- .../cuda/ops/reshape_and_cache_flash/kernel.h | 12 +- tests/conftest.py | 18 +++ tests/test_flash_attn_varlen_func.py | 13 +++ tests/test_flash_attn_with_kvcache.py | 26 ++++- 11 files changed, 306 insertions(+), 21 deletions(-) create mode 100644 src/native/cuda/iluvatar/ops/flash_attn_varlen_func/kernel.h create mode 100644 src/native/cuda/iluvatar/ops/flash_attn_with_kvcache/kernel.h diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 20196f657..b151d9b0f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -585,10 +585,14 @@ if(WITH_NVIDIA) list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum) endif() -if((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH) +if(((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH) OR WITH_ILUVATAR) list(APPEND _infini_ops_smoke_ops flash_attn_varlen_func) endif() +if(WITH_ILUVATAR) + list(APPEND _infini_ops_smoke_ops flash_attn_with_kvcache) +endif() + if(INFINI_OPS_SMOKE_BUILD) if(NOT INFINI_OPS_OPS) set(INFINI_OPS_OPS "${_infini_ops_smoke_ops}" CACHE STRING diff --git a/src/native/cuda/iluvatar/ops/flash_attn_varlen_func/kernel.h b/src/native/cuda/iluvatar/ops/flash_attn_varlen_func/kernel.h new file mode 100644 index 000000000..6b53ad3c9 --- /dev/null +++ b/src/native/cuda/iluvatar/ops/flash_attn_varlen_func/kernel.h @@ -0,0 +1,95 @@ +#ifndef INFINI_OPS_ILUVATAR_FLASH_ATTN_VARLEN_FUNC_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_FLASH_ATTN_VARLEN_FUNC_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/flash_attn_varlen_func.h" +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/paged_attention_prefill_infinilm/kernel.h" + +namespace infini::ops { +namespace flash_attn_varlen_func_iluvatar_detail { + +inline Tensor LegacyCacheView(Tensor cache) { + return {cache.data(), + {cache.size(0), cache.size(2), cache.size(1), cache.size(3)}, + cache.dtype(), + cache.device(), + {cache.stride(0), cache.stride(2), cache.stride(1), cache.stride(3)}}; +} + +// The compatibility kernel consumes one total KV length per sequence. This +// metadata view keeps the canonical cumulative buffer and lets the native +// kernel derive adjacent differences without allocating another device tensor. +inline Tensor CumulativeLengthsView(Tensor cumulative_lengths) { + return {cumulative_lengths.data(), + {cumulative_lengths.size(0) - 1}, + cumulative_lengths.dtype(), + cumulative_lengths.device(), + {cumulative_lengths.stride(0)}}; +} + +} // namespace flash_attn_varlen_func_iluvatar_detail + +template <> +class Operator + : public FlashAttnVarlenFunc { + public: + using FlashAttnVarlenFunc::FlashAttnVarlenFunc; + using FlashAttnVarlenFunc::operator(); + + 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, + const std::optional block_table, + const int64_t /*max_seqlen_q*/, + const int64_t /*max_seqlen_k*/, const double dropout_p, + const std::optional softmax_scale, const bool causal, + const std::vector window_size, const double softcap, + const bool deterministic, const bool return_attn_probs, + Tensor out, std::optional softmax_lse, + std::optional s_dmask) const override { + assert((q.size(2) == 64 || q.size(2) == 128) && + "Iluvatar native `FlashAttnVarlenFunc` supports only head " + "dimensions 64 and 128."); + assert(block_table && dropout_p == 0.0 && causal && + window_size == std::vector({-1, -1}) && softcap == 0.0 && + !deterministic && !return_attn_probs && !softmax_lse && !s_dmask && + (!alibi_slopes || alibi_slopes->ndim() == 1) && + "Iluvatar native `FlashAttnVarlenFunc` supports only causal paged " + "inference without dropout, local windows, softcap, deterministic " + "mode, or auxiliary outputs."); + + auto legacy_k_cache = + flash_attn_varlen_func_iluvatar_detail::LegacyCacheView(k); + auto legacy_v_cache = + flash_attn_varlen_func_iluvatar_detail::LegacyCacheView(v); + auto cumulative_lengths = + flash_attn_varlen_func_iluvatar_detail::CumulativeLengthsView( + cu_seqlens_k); + const float scale = static_cast(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(q.size(2))))); + + CudaPagedAttentionPrefillInfinilm, true> + provider{q, + legacy_k_cache, + legacy_v_cache, + *block_table, + cumulative_lengths, + cu_seqlens_q, + alibi_slopes, + scale, + out}; + provider.set_stream(stream_); + provider(q, legacy_k_cache, legacy_v_cache, *block_table, + cumulative_lengths, cu_seqlens_q, alibi_slopes, scale, out); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ILUVATAR_FLASH_ATTN_VARLEN_FUNC_KERNEL_H_ diff --git a/src/native/cuda/iluvatar/ops/flash_attn_with_kvcache/kernel.h b/src/native/cuda/iluvatar/ops/flash_attn_with_kvcache/kernel.h new file mode 100644 index 000000000..574a4e4b1 --- /dev/null +++ b/src/native/cuda/iluvatar/ops/flash_attn_with_kvcache/kernel.h @@ -0,0 +1,110 @@ +#ifndef INFINI_OPS_ILUVATAR_FLASH_ATTN_WITH_KVCACHE_KERNEL_H_ +#define INFINI_OPS_ILUVATAR_FLASH_ATTN_WITH_KVCACHE_KERNEL_H_ + +#include +#include +#include +#include + +#include "base/flash_attn_with_kvcache.h" +#include "native/cuda/iluvatar/caster.cuh" +#include "native/cuda/iluvatar/runtime_.h" +#include "native/cuda/ops/paged_attention_infinilm/kernel.h" + +namespace infini::ops { +namespace flash_attn_with_kvcache_iluvatar_detail { + +inline Tensor DecodeView(Tensor tensor) { + return {tensor.data(), + {tensor.size(0), tensor.size(2), tensor.size(3)}, + tensor.dtype(), + tensor.device(), + {tensor.stride(0), tensor.stride(2), tensor.stride(3)}}; +} + +inline Tensor LegacyCacheView(Tensor cache) { + return {cache.data(), + {cache.size(0), cache.size(2), cache.size(1), cache.size(3)}, + cache.dtype(), + cache.device(), + {cache.stride(0), cache.stride(2), cache.stride(1), cache.stride(3)}}; +} + +} // namespace flash_attn_with_kvcache_iluvatar_detail + +template <> +class Operator + : public FlashAttnWithKvcache { + public: + using FlashAttnWithKvcache::FlashAttnWithKvcache; + using FlashAttnWithKvcache::operator(); + + std::size_t workspace_size_in_bytes() const override { + return static_cast(kMaxSplits) * q_shape_[0] * q_shape_[2] * + (head_size_ + 2) * sizeof(float); + } + + void operator()(const Tensor, Tensor, Tensor, const std::optional, + const std::optional, const std::optional, + const std::optional, const int64_t, + const std::optional, const std::optional, + const std::optional, const std::optional, + const std::optional, const bool, + const std::vector, const double, const bool, + const int64_t, const bool, Tensor, + std::optional) const override { + assert(false && + "Iluvatar native `FlashAttnWithKvcache` requires tensor " + "`cache_seqlens`."); + } + + 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 { + assert((head_size_ == 64 || head_size_ == 128) && + "Iluvatar native `FlashAttnWithKvcache` supports only head " + "dimensions 64 and 128."); + assert(q.size(1) == 1 && !k && !v && !rotary_cos && !rotary_sin && + cache_seqlens && !cache_batch_idx && !cache_leftpad && block_table && + causal && window_size == std::vector({-1, -1}) && + softcap == 0.0 && !return_softmax_lse && !softmax_lse && + (!alibi_slopes || alibi_slopes->ndim() == 1) && + "Iluvatar native `FlashAttnWithKvcache` supports only causal paged " + "decode without KV update, rotary inputs, local windows, softcap, " + "or auxiliary outputs."); + + auto q_view = flash_attn_with_kvcache_iluvatar_detail::DecodeView(q); + auto out_view = flash_attn_with_kvcache_iluvatar_detail::DecodeView(out); + auto legacy_k_cache = + flash_attn_with_kvcache_iluvatar_detail::LegacyCacheView(k_cache); + auto legacy_v_cache = + flash_attn_with_kvcache_iluvatar_detail::LegacyCacheView(v_cache); + const float scale = static_cast(softmax_scale.value_or( + 1.0 / std::sqrt(static_cast(head_size_)))); + + CudaPagedAttentionInfinilm> provider{ + q_view, legacy_k_cache, legacy_v_cache, *block_table, + *cache_seqlens, alibi_slopes, scale, out_view}; + provider.set_stream(stream_); + provider.set_workspace(workspace_); + provider.set_workspace_size_in_bytes(workspace_size_in_bytes_); + provider(q_view, legacy_k_cache, legacy_v_cache, *block_table, + *cache_seqlens, alibi_slopes, scale, out_view); + } +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_ILUVATAR_FLASH_ATTN_WITH_KVCACHE_KERNEL_H_ diff --git a/src/native/cuda/iluvatar/ops/reshape_and_cache_flash/kernel.h b/src/native/cuda/iluvatar/ops/reshape_and_cache_flash/kernel.h index e13a36f6d..039b88060 100644 --- a/src/native/cuda/iluvatar/ops/reshape_and_cache_flash/kernel.h +++ b/src/native/cuda/iluvatar/ops/reshape_and_cache_flash/kernel.h @@ -11,10 +11,14 @@ namespace infini::ops { template <> class Operator - : public CudaReshapeAndCacheFlash> { + : public CudaReshapeAndCacheFlash, 128> { public: - using CudaReshapeAndCacheFlash< - Runtime>::CudaReshapeAndCacheFlash; + // BI-V150 reports a 2048-thread device limit, but CoreX faults when this + // kernel uses the resulting 1024/2048-thread launch. The operator supports + // head dimensions up to 128 here, so 128 threads cover every element while + // keeping the workaround local to the Iluvatar provider. + using CudaReshapeAndCacheFlash, + 128>::CudaReshapeAndCacheFlash; }; } // namespace infini::ops diff --git a/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh b/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh index 6906e61e7..6cf076024 100644 --- a/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh +++ b/src/native/cuda/ops/paged_attention_infinilm/kernel.cuh @@ -10,6 +10,13 @@ namespace op::paged_attention::cuda { +// Standalone InfiniOps defines WITH_ILUVATAR; the legacy InfiniCore xmake +// integration defines ENABLE_ILUVATAR_API. Keep both build paths on the +// Iluvatar-safe warp and shared-memory implementation. +#if defined(WITH_ILUVATAR) || defined(ENABLE_ILUVATAR_API) +#define INFINI_OPS_PAGED_ATTENTION_ILUVATAR 1 +#endif + struct OnlineSoftmaxState { float m = -INFINITY; @@ -25,7 +32,7 @@ struct OnlineSoftmaxState { }; __device__ __forceinline__ float WarpReduceSum(float x) { -#if defined(ENABLE_ILUVATAR_API) +#if defined(INFINI_OPS_PAGED_ATTENTION_ILUVATAR) // Iluvatar may use warp size 64; __shfl_sync(0xffffffff) only covers 32 // threads. Use shared-memory tree reduce for portability across warp sizes. constexpr int kMaxWarps = 16; @@ -51,7 +58,7 @@ __device__ __forceinline__ float WarpReduceSum(float x) { } __device__ __forceinline__ float WarpBroadcast(float x, int src_lane) { -#if defined(ENABLE_ILUVATAR_API) +#if defined(INFINI_OPS_PAGED_ATTENTION_ILUVATAR) __shared__ float _bcast_buf[16]; const int warp_id = threadIdx.x / 32; if ((threadIdx.x & 31) == src_lane) { @@ -65,7 +72,7 @@ __device__ __forceinline__ float WarpBroadcast(float x, int src_lane) { } __device__ __forceinline__ float WarpReduceMax(float x) { -#if defined(ENABLE_ILUVATAR_API) +#if defined(INFINI_OPS_PAGED_ATTENTION_ILUVATAR) __shared__ float _reduce_buf[16 * 32]; const int lane = threadIdx.x & 31; const int warp_id = threadIdx.x / 32; @@ -89,7 +96,7 @@ __device__ __forceinline__ float WarpReduceMax(float x) { } __device__ __forceinline__ unsigned int CvtaToShared(const void* ptr) { -#if defined(ENABLE_ILUVATAR_API) +#if defined(INFINI_OPS_PAGED_ATTENTION_ILUVATAR) return static_cast(reinterpret_cast(ptr)); #else return static_cast(__cvta_generic_to_shared(ptr)); diff --git a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.cuh b/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.cuh index 04ce64f11..9ca33811c 100644 --- a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.cuh +++ b/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.cuh @@ -7,7 +7,7 @@ #include #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) || \ - defined(ENABLE_ILUVATAR_API) + defined(WITH_ILUVATAR) || defined(ENABLE_ILUVATAR_API) #include #include #include @@ -481,7 +481,7 @@ __global__ void PagedAttentionPrefillWarpGlobalKernel( if (lane == 0) { inv_l = 1.0f / (l + 1e-6f); } -#ifdef ENABLE_ILUVATAR_API +#ifdef INFINI_OPS_PAGED_ATTENTION_ILUVATAR inv_l = op::paged_attention::cuda::WarpBroadcast(inv_l, 0); #else inv_l = __shfl_sync(0xffffffff, inv_l, 0); diff --git a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h b/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h index 598ead224..a1b2a10e2 100644 --- a/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h +++ b/src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h @@ -17,7 +17,7 @@ namespace infini::ops { using PagedAttentionPrefillInfinilmIndexTypes = List; -template +template class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { public: using PagedAttentionPrefillInfinilm::PagedAttentionPrefillInfinilm; @@ -52,6 +52,16 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { using TIndex = TypeMapType(list_tag)>; constexpr int kHeadSize = ListGet<2>(list_tag); + const auto* sequence_lengths = + reinterpret_cast(seq_lens.data()); + const auto total_kv_lengths = [&]() { + if constexpr (kCumulativeSequenceLengths) { + return op::paged_attention_prefill::cuda:: + CumulativeSequenceLengths{sequence_lengths}; + } else { + return sequence_lengths; + } + }(); if constexpr (kHeadSize == 128) { if (block_size_ == 256) { @@ -68,7 +78,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { reinterpret_cast(k_cache.data()), reinterpret_cast(v_cache.data()), reinterpret_cast(block_tables.data()), - reinterpret_cast(seq_lens.data()), + total_kv_lengths, reinterpret_cast(cum_seq_lens_q.data()), alibi_slopes.has_value() ? reinterpret_cast(alibi_slopes->data()) @@ -91,7 +101,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { reinterpret_cast(k_cache.data()), reinterpret_cast(v_cache.data()), reinterpret_cast(block_tables.data()), - reinterpret_cast(seq_lens.data()), + total_kv_lengths, reinterpret_cast(cum_seq_lens_q.data()), alibi_slopes.has_value() ? reinterpret_cast(alibi_slopes->data()) @@ -116,7 +126,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm { reinterpret_cast(k_cache.data()), reinterpret_cast(v_cache.data()), reinterpret_cast(block_tables.data()), - reinterpret_cast(seq_lens.data()), + total_kv_lengths, reinterpret_cast(cum_seq_lens_q.data()), alibi_slopes.has_value() ? reinterpret_cast(alibi_slopes->data()) diff --git a/src/native/cuda/ops/reshape_and_cache_flash/kernel.h b/src/native/cuda/ops/reshape_and_cache_flash/kernel.h index 54294dcab..4dc591c64 100644 --- a/src/native/cuda/ops/reshape_and_cache_flash/kernel.h +++ b/src/native/cuda/ops/reshape_and_cache_flash/kernel.h @@ -14,7 +14,10 @@ namespace infini::ops { -template +// Backends may lower the launch cap for this kernel without changing the +// default launch policy used by other CUDA-compatible platforms. +template ::value> class CudaReshapeAndCacheFlash : public ReshapeAndCacheFlash { public: using ReshapeAndCacheFlash::ReshapeAndCacheFlash; @@ -32,13 +35,12 @@ class CudaReshapeAndCacheFlash : public ReshapeAndCacheFlash { static_cast(stream_ ? stream_ : 0); int block_size = std::min(RuntimeUtils::GetOptimalBlockSize(), - BackendMaxBlockSize::value); + kMaxBlockSize); dim3 grid(static_cast(num_heads_), static_cast(num_tokens_)); - DispatchFunc< - ConcatType, ReducedFloatTypes>, - SupportedCudaBlockSizesType::value>>( + DispatchFunc, ReducedFloatTypes>, + SupportedCudaBlockSizesType>( {static_cast(dtype_), block_size}, [&](auto list_tag) { using T = TypeMapType(list_tag)>; diff --git a/tests/conftest.py b/tests/conftest.py index 2d63e5ec8..e206fee3f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -497,6 +497,24 @@ def _is_smoke_cutlass_scaled_mm_case(params): def _is_smoke_flash_attn_varlen_func_case(params): + if ( + params.get("device") == "cuda" + and params.get("implementation_index") == 0 + ): + return ( + params.get("q_lens") == (2, 3) + and params.get("k_lens") == (130, 300) + and params.get("num_heads") == 4 + and params.get("num_kv_heads") == 2 + and params.get("causal") is True + and params.get("window_size") == (-1, -1) + and params.get("scale") is None + and params.get("paged") is True + and params.get("use_alibi") is True + and params.get("head_dim") == 64 + and params.get("dtype") == torch.float16 + ) + dense = ( params.get("q_lens") == (3, 5) and params.get("k_lens") == (4, 5) diff --git a/tests/test_flash_attn_varlen_func.py b/tests/test_flash_attn_varlen_func.py index 5bdf5d514..0682facb2 100644 --- a/tests/test_flash_attn_varlen_func.py +++ b/tests/test_flash_attn_varlen_func.py @@ -60,6 +60,11 @@ def test_flash_attn_varlen_func( if device == "musa" and not paged and causal and q_lens != k_lens: pytest.skip("TorchMusa causal FlashAttention requires matching Q/K lengths") + if device == "cuda" and implementation_index == 0 and not ( + paged and causal and window_size == (-1, -1) + ): + pytest.skip("Iluvatar native provider supports causal paged inference") + if device == "cuda" and (paged or use_alibi) and implementation_index == 8: pytest.skip("paged KV cache and ALiBi require the linked provider") @@ -189,6 +194,8 @@ def test_flash_attn_varlen_func( def test_flash_attn_varlen_func_non_default_stream(device, implementation_index): if device != "cuda": pytest.skip("non-default CUDA streams require the NVIDIA backend") + if implementation_index == 0: + pytest.skip("Iluvatar native provider requires a paged KV cache") dtype = torch.float16 q_lens = (3, 5) @@ -243,6 +250,8 @@ def test_flash_attn_varlen_func_non_default_stream(device, implementation_index) def test_flash_attn_varlen_func_default_stream(device, implementation_index): if device != "cuda": pytest.skip("CUDA stream coverage requires the NVIDIA backend") + if implementation_index == 0: + pytest.skip("Iluvatar native provider requires a paged KV cache") q = torch.randn((5, 4, 64), dtype=torch.float16, device=device) k = torch.randn_like(q) @@ -285,6 +294,8 @@ def test_flash_attn_varlen_func_default_stream(device, implementation_index): def test_flash_attn_varlen_func_defaults(device, implementation_index): if device not in ("cuda", "musa"): pytest.skip("FlashAttention requires the NVIDIA or Moore backend") + if device == "cuda" and implementation_index == 0: + pytest.skip("Iluvatar native provider requires a paged KV cache") q = torch.randn((5, 4, 64), dtype=torch.float16, device=device) k = torch.randn((5, 4, 64), dtype=torch.float16, device=device) @@ -319,6 +330,8 @@ def test_flash_attn_varlen_func_defaults(device, implementation_index): def test_flash_attn_varlen_func_device_guard(): + if 0 in infini.ops.FlashAttnVarlenFunc.active_implementation_indices("iluvatar"): + pytest.skip("Iluvatar native provider requires a paged KV cache") if torch.cuda.device_count() < 2: pytest.skip("device-guard coverage requires at least two NVIDIA GPUs") diff --git a/tests/test_flash_attn_with_kvcache.py b/tests/test_flash_attn_with_kvcache.py index 366f1c795..f761bbc00 100644 --- a/tests/test_flash_attn_with_kvcache.py +++ b/tests/test_flash_attn_with_kvcache.py @@ -7,7 +7,10 @@ from tests.utils import get_stream -flash_attn = pytest.importorskip("flash_attn") +try: + import flash_attn +except ImportError: + flash_attn = None if not hasattr(infini.ops, "FlashAttnWithKvcache"): @@ -37,6 +40,10 @@ def test_flash_attn_with_kvcache_dense( ): if device not in ("cuda", "mlu"): pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") + if device == "cuda" and implementation_index == 0: + pytest.skip("Iluvatar native provider supports paged decode only") + if device == "cuda" and flash_attn is None: + pytest.skip("flash_attn is required as the dense reference") batch_size, cache_size = 2, 16 num_heads, num_kv_heads, head_size = 4, 2, 64 @@ -136,6 +143,7 @@ def test_flash_attn_with_kvcache_dense( torch.testing.assert_close(actual_v_cache, expected_v_cache, rtol=0, atol=0) +@pytest.mark.smoke def test_flash_attn_with_kvcache_paged(device, implementation_index): if device not in ("cuda", "mlu"): pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") @@ -155,7 +163,7 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index): v_cache = torch.randn_like(k_cache) cache_seqlens = torch.tensor((130, 300), dtype=torch.int32, device=device) block_table = torch.tensor(((0, 1), (2, 3)), dtype=torch.int32, device=device) - if device == "mlu": + if device == "mlu" or (device == "cuda" and implementation_index == 0): expected, _ = _reference_flash_attn_with_kvcache( q, k_cache, @@ -165,6 +173,8 @@ def test_flash_attn_with_kvcache_paged(device, implementation_index): causal=True, ) else: + if flash_attn is None: + pytest.skip("flash_attn is required as the paged reference") expected = flash_attn.flash_attn_with_kvcache( q, k_cache, @@ -209,6 +219,10 @@ def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( ): if device not in ("cuda", "mlu"): pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") + if device == "cuda" and implementation_index == 0: + pytest.skip("Iluvatar native provider requires tensor cache lengths") + if device == "cuda" and flash_attn is None: + pytest.skip("flash_attn is required as the dense reference") q = torch.randn((2, 1, 4, 60), dtype=torch.float16, device=device) k_cache = torch.randn((3, 8, 2, 60), dtype=torch.float16, device=device) @@ -266,6 +280,10 @@ def test_flash_attn_with_kvcache_scalar_seqlens_with_cache_batch_idx( def test_flash_attn_with_kvcache_defaults(device, implementation_index): if device not in ("cuda", "mlu"): pytest.skip("FlashAttention FA2 requires the NVIDIA or Cambricon backend") + if device == "cuda" and implementation_index == 0: + pytest.skip("Iluvatar native provider supports paged decode only") + if device == "cuda" and flash_attn is None: + pytest.skip("flash_attn is required as the dense reference") q = torch.randn((2, 1, 4, 64), dtype=torch.float16, device=device) k_cache = torch.randn((2, 8, 2, 64), dtype=torch.float16, device=device) @@ -297,6 +315,10 @@ def test_flash_attn_with_kvcache_non_default_stream(device, implementation_index stream_attribute = "mlu_stream" else: pytest.skip("stream coverage requires an accelerator backend") + if device == "cuda" and implementation_index == 0: + pytest.skip("Iluvatar native provider supports paged decode only") + if device == "cuda" and flash_attn is None: + pytest.skip("flash_attn is required as the dense reference") q = torch.randn((2, 1, 4, 64), dtype=torch.float16, device=device) k_cache = torch.randn((2, 8, 2, 64), dtype=torch.float16, device=device)