feat(iluvatar): add native canonical attention providers

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()
@@ -583,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

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 <cassert>
+#include <cmath>
+#include <optional>
+#include <vector>
+
+#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<FlashAttnVarlenFunc, Device::Type::kIluvatar>
+    : 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<Tensor> alibi_slopes,
+                  const std::optional<Tensor> block_table,
+                  const int64_t /*max_seqlen_q*/,
+                  const int64_t /*max_seqlen_k*/, const double dropout_p,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool deterministic, const bool return_attn_probs,
+                  Tensor out, std::optional<Tensor> softmax_lse,
+                  std::optional<Tensor> 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<int64_t>({-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<float>(softmax_scale.value_or(
+        1.0 / std::sqrt(static_cast<double>(q.size(2)))));
+
+    CudaPagedAttentionPrefillInfinilm<Runtime<Device::Type::kIluvatar>, 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_

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 <cassert>
+#include <cmath>
+#include <optional>
+#include <vector>
+
+#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<FlashAttnWithKvcache, Device::Type::kIluvatar>
+    : public FlashAttnWithKvcache {
+ public:
+  using FlashAttnWithKvcache::FlashAttnWithKvcache;
+  using FlashAttnWithKvcache::operator();
+
+  std::size_t workspace_size_in_bytes() const override {
+    return static_cast<std::size_t>(kMaxSplits) * q_shape_[0] * q_shape_[2] *
+           (head_size_ + 2) * sizeof(float);
+  }
+
+  void operator()(const Tensor, Tensor, Tensor, const std::optional<Tensor>,
+                  const std::optional<Tensor>, const std::optional<Tensor>,
+                  const std::optional<Tensor>, const int64_t,
+                  const std::optional<Tensor>, const std::optional<Tensor>,
+                  const std::optional<Tensor>, const std::optional<Tensor>,
+                  const std::optional<double>, const bool,
+                  const std::vector<int64_t>, const double, const bool,
+                  const int64_t, const bool, Tensor,
+                  std::optional<Tensor>) 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<Tensor> k, const std::optional<Tensor> v,
+                  const std::optional<Tensor> rotary_cos,
+                  const std::optional<Tensor> rotary_sin,
+                  const std::optional<Tensor> cache_seqlens,
+                  const std::optional<Tensor> cache_batch_idx,
+                  const std::optional<Tensor> cache_leftpad,
+                  const std::optional<Tensor> block_table,
+                  const std::optional<Tensor> alibi_slopes,
+                  const std::optional<double> softmax_scale, const bool causal,
+                  const std::vector<int64_t> window_size, const double softcap,
+                  const bool /*rotary_interleaved*/,
+                  const int64_t /*num_splits*/, const bool return_softmax_lse,
+                  Tensor out,
+                  std::optional<Tensor> 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<int64_t>({-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<float>(softmax_scale.value_or(
+        1.0 / std::sqrt(static_cast<double>(head_size_))));
+
+    CudaPagedAttentionInfinilm<Runtime<Device::Type::kIluvatar>> 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_

src/native/cuda/iluvatar/ops/reshape_and_cache_flash/kernel.h

@@ -11,10 +11,14 @@ namespace infini::ops {
 
 template <>
 class Operator<ReshapeAndCacheFlash, Device::Type::kIluvatar>
-    : public CudaReshapeAndCacheFlash<Runtime<Device::Type::kIluvatar>> {
+    : public CudaReshapeAndCacheFlash<Runtime<Device::Type::kIluvatar>, 128> {
  public:
-  using CudaReshapeAndCacheFlash<
-      Runtime<Device::Type::kIluvatar>>::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<Runtime<Device::Type::kIluvatar>,
+                                 128>::CudaReshapeAndCacheFlash;
 };
 
 }  // namespace infini::ops

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<unsigned int>(reinterpret_cast<uintptr_t>(ptr));
 #else
   return static_cast<unsigned int>(__cvta_generic_to_shared(ptr));

src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.cuh

@@ -7,7 +7,7 @@
 #include <cstdint>
 
 #if defined(ENABLE_NVIDIA_API) || defined(ENABLE_ALI_API) || \
-    defined(ENABLE_ILUVATAR_API)
+    defined(WITH_ILUVATAR) || defined(ENABLE_ILUVATAR_API)
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
@@ -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);

src/native/cuda/ops/paged_attention_prefill_infinilm/kernel.h

@@ -17,7 +17,7 @@ namespace infini::ops {
 using PagedAttentionPrefillInfinilmIndexTypes =
     List<DataType::kInt32, DataType::kInt64, DataType::kUInt32>;
 
-template <typename Backend>
+template <typename Backend, bool kCumulativeSequenceLengths = false>
 class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
  public:
   using PagedAttentionPrefillInfinilm::PagedAttentionPrefillInfinilm;
@@ -52,6 +52,16 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
           using TIndex =
               TypeMapType<Backend::kDeviceType, ListGet<1>(list_tag)>;
           constexpr int kHeadSize = ListGet<2>(list_tag);
+          const auto* sequence_lengths =
+              reinterpret_cast<const TIndex*>(seq_lens.data());
+          const auto total_kv_lengths = [&]() {
+            if constexpr (kCumulativeSequenceLengths) {
+              return op::paged_attention_prefill::cuda::
+                  CumulativeSequenceLengths<TIndex>{sequence_lengths};
+            } else {
+              return sequence_lengths;
+            }
+          }();
 
           if constexpr (kHeadSize == 128) {
             if (block_size_ == 256) {
@@ -68,7 +78,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
                       reinterpret_cast<const TData*>(k_cache.data()),
                       reinterpret_cast<const TData*>(v_cache.data()),
                       reinterpret_cast<const TIndex*>(block_tables.data()),
-                      reinterpret_cast<const TIndex*>(seq_lens.data()),
+                      total_kv_lengths,
                       reinterpret_cast<const TIndex*>(cum_seq_lens_q.data()),
                       alibi_slopes.has_value()
                           ? reinterpret_cast<const float*>(alibi_slopes->data())
@@ -91,7 +101,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
                       reinterpret_cast<const TData*>(k_cache.data()),
                       reinterpret_cast<const TData*>(v_cache.data()),
                       reinterpret_cast<const TIndex*>(block_tables.data()),
-                      reinterpret_cast<const TIndex*>(seq_lens.data()),
+                      total_kv_lengths,
                       reinterpret_cast<const TIndex*>(cum_seq_lens_q.data()),
                       alibi_slopes.has_value()
                           ? reinterpret_cast<const float*>(alibi_slopes->data())
@@ -116,7 +126,7 @@ class CudaPagedAttentionPrefillInfinilm : public PagedAttentionPrefillInfinilm {
                     reinterpret_cast<const TData*>(k_cache.data()),
                     reinterpret_cast<const TData*>(v_cache.data()),
                     reinterpret_cast<const TIndex*>(block_tables.data()),
-                    reinterpret_cast<const TIndex*>(seq_lens.data()),
+                    total_kv_lengths,
                     reinterpret_cast<const TIndex*>(cum_seq_lens_q.data()),
                     alibi_slopes.has_value()
                         ? reinterpret_cast<const float*>(alibi_slopes->data())