refactor(linked)!: migrate flash_attn_with_kvcache provider

docs/build.md

@@ -51,11 +51,11 @@ python -m pip install .[dev] \
   --config-settings=cmake.define.WITH_NVIDIA=ON
 ```
 
-Full builds with both `WITH_NVIDIA=ON` and `WITH_TORCH=ON` include
-`flash_attn_with_kvcache` and require a compatible FlashAttention 2.7 Python
-distribution to be installed in the build environment. The distribution must
-provide the `flash_attn_2_cuda` shared library for the active PyTorch and CUDA
-toolchain. Operator-pruned builds require it only when
+Full builds with both `WITH_NVIDIA=ON` and `WITH_LINKED=ON` include
+`flash_attn_with_kvcache` and require a compatible FlashAttention Python
+distribution in the build environment. The distribution must provide a
+`flash_attn_2_cuda` shared library that exports the required C++ ABI for the
+active PyTorch and CUDA toolchain. Operator-pruned builds require it only when
 `flash_attn_with_kvcache` is selected; the default smoke subset does not select
 this operator.
 

src/CMakeLists.txt

@@ -620,73 +620,6 @@ if(WITH_LINKED)
         "Resolved ${_linked_source_count} linked operator source(s).")
 endif()
 
-set(_build_flash_attn_with_kvcache FALSE)
-if(WITH_NVIDIA AND WITH_TORCH AND
-   (NOT INFINI_OPS_OPS OR
-    "flash_attn_with_kvcache" IN_LIST _infini_ops_op_allowlist))
-    execute_process(
-        COMMAND ${_TORCH_PYTHON} -c
-            "import importlib.metadata; print(importlib.metadata.version('flash-attn'))"
-        OUTPUT_VARIABLE FLASH_ATTN_VERSION
-        OUTPUT_STRIP_TRAILING_WHITESPACE
-        RESULT_VARIABLE _flash_attn_version_result
-        ERROR_QUIET
-    )
-    if(NOT _flash_attn_version_result EQUAL 0 OR
-       NOT FLASH_ATTN_VERSION MATCHES "^2\\.7([.+-]|$)")
-        message(FATAL_ERROR
-            "`flash_attn_with_kvcache` requires installed FlashAttention 2.7; "
-            "found `${FLASH_ATTN_VERSION}`.")
-    endif()
-
-    execute_process(
-        COMMAND ${_TORCH_PYTHON} -c
-            "import importlib.util; s = importlib.util.find_spec('flash_attn_2_cuda'); print(s.origin if s else '')"
-        OUTPUT_VARIABLE FLASH_ATTN_LIBRARY
-        OUTPUT_STRIP_TRAILING_WHITESPACE
-        RESULT_VARIABLE _flash_attn_library_result
-        ERROR_QUIET
-    )
-    if(NOT _flash_attn_library_result EQUAL 0 OR
-       NOT EXISTS "${FLASH_ATTN_LIBRARY}")
-        message(FATAL_ERROR
-            "FlashAttention ${FLASH_ATTN_VERSION} does not provide the "
-            "`flash_attn_2_cuda` shared library.")
-    endif()
-
-    execute_process(
-        COMMAND ${CMAKE_NM} -D -C "${FLASH_ATTN_LIBRARY}"
-        OUTPUT_VARIABLE _flash_attn_symbols
-        RESULT_VARIABLE _flash_attn_nm_result
-        ERROR_QUIET
-    )
-    string(FIND "${_flash_attn_symbols}" "flash::mha_fwd_kvcache("
-        _flash_attn_kvcache_symbol)
-    if(NOT _flash_attn_nm_result EQUAL 0 OR
-       _flash_attn_kvcache_symbol EQUAL -1)
-        message(FATAL_ERROR
-            "FlashAttention ${FLASH_ATTN_VERSION} does not export the "
-            "required `flash::mha_fwd_kvcache` ABI.")
-    endif()
-
-    set(_build_flash_attn_with_kvcache TRUE)
-    get_filename_component(FLASH_ATTN_RUNTIME_DIR
-        "${FLASH_ATTN_LIBRARY}" DIRECTORY)
-    get_filename_component(FLASH_ATTN_LIBRARY_NAME
-        "${FLASH_ATTN_LIBRARY}" NAME)
-    target_link_directories(infiniops PRIVATE "${FLASH_ATTN_RUNTIME_DIR}")
-    # The Python extension has no SONAME, so link by its exact basename.
-    target_link_libraries(infiniops PRIVATE
-        "-l:${FLASH_ATTN_LIBRARY_NAME}")
-
-    list(APPEND TORCH_RUNTIME_DIRS
-        "${FLASH_ATTN_RUNTIME_DIR}" "$ORIGIN/..")
-    set_property(TARGET infiniops APPEND PROPERTY BUILD_RPATH
-        "${FLASH_ATTN_RUNTIME_DIR}")
-    message(STATUS
-        "Found FlashAttention ${FLASH_ATTN_VERSION}: ${FLASH_ATTN_LIBRARY}")
-endif()
-
 set(TORCH_SOURCES "")
 if(WITH_TORCH)
     # Auto-generate ATen-backed operator wrappers from `scripts/torch_ops.yaml`.
@@ -743,10 +676,6 @@ if(WITH_TORCH)
         list(FILTER TORCH_SOURCES EXCLUDE REGEX
             ".*/flash_attn_varlen_func/flash_attn_varlen_func\\.cc$")
     endif()
-    if(NOT _build_flash_attn_with_kvcache)
-        list(FILTER TORCH_SOURCES EXCLUDE REGEX
-            ".*/flash_attn_with_kvcache/flash_attn\\.cc$")
-    endif()
 endif()
 
 if(_infini_ops_linked_uses_torch)

src/linked/torch/nvidia/c10.h

@@ -0,0 +1,28 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <cuda_runtime_api.h>
+
+#include "linked/torch/c10.h"
+
+namespace infini::ops::linked::torch {
+
+template <>
+struct C10<Device::Type::kNvidia> {
+  static constexpr Device::Type kDeviceType = Device::Type::kNvidia;
+
+  using StreamGuard = c10::cuda::CUDAStreamGuard;
+
+  static c10::cuda::CUDAStream GetStreamFromExternal(void* stream,
+                                                     int device_index) {
+    return c10::cuda::getStreamFromExternal(
+        reinterpret_cast<cudaStream_t>(stream),
+        static_cast<c10::DeviceIndex>(device_index));
+  }
+};
+
+}  // namespace infini::ops::linked::torch
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_C10_H_

src/linked/torch/nvidia/flash_attn.yaml

@@ -0,0 +1,2 @@
+python_distribution_package: flash-attn
+library_glob: flash_attn_2_cuda*.so

src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.cc

@@ -0,0 +1,50 @@
+#include "linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h"
+
+namespace flash {
+
+std::vector<at::Tensor> mha_fwd_kvcache(
+    at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
+    std::optional<const at::Tensor>& cache_seqlens,
+    std::optional<const at::Tensor>& rotary_cos,
+    std::optional<const at::Tensor>& rotary_sin,
+    std::optional<const at::Tensor>& cache_batch_idx,
+    std::optional<const at::Tensor>& cache_leftpad,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
+    float softmax_scale, bool causal, int window_size_left,
+    int window_size_right, float softcap, bool rotary_interleaved,
+    int num_splits);
+
+}  // namespace flash
+
+namespace infini::ops::linked::torch::nvidia {
+
+std::vector<at::Tensor> FlashAttnKvcache::Call(
+    at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+    std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
+    std::optional<const at::Tensor>& cache_seqlens,
+    std::optional<const at::Tensor>& rotary_cos,
+    std::optional<const at::Tensor>& rotary_sin,
+    std::optional<const at::Tensor>& cache_batch_idx,
+    std::optional<const at::Tensor>& cache_leftpad,
+    std::optional<at::Tensor>& block_table,
+    std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
+    float softmax_scale, bool causal, int window_size_left,
+    int window_size_right, float softcap, bool rotary_interleaved,
+    int num_splits) {
+  return flash::mha_fwd_kvcache(
+      q, k_cache, v_cache, k, v, cache_seqlens, rotary_cos, rotary_sin,
+      cache_batch_idx, cache_leftpad, block_table, alibi_slopes, out,
+      softmax_scale, causal, window_size_left, window_size_right, softcap,
+      rotary_interleaved, num_splits);
+}
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+template class TorchFlashAttnWithKvcache<
+    ::infini::ops::linked::torch::nvidia::FlashAttnKvcache>;
+
+}  // namespace infini::ops::linked::torch

src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h

@@ -0,0 +1,50 @@
+#ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+#define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
+
+#include "linked/torch/nvidia/c10.h"
+#include "linked/torch/ops/flash_attn_with_kvcache.h"
+
+namespace infini::ops::linked::torch::nvidia {
+
+struct FlashAttnKvcache : C10<Device::Type::kNvidia> {
+  static std::vector<at::Tensor> Call(
+      at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache,
+      std::optional<const at::Tensor>& k, std::optional<const at::Tensor>& v,
+      std::optional<const at::Tensor>& cache_seqlens,
+      std::optional<const at::Tensor>& rotary_cos,
+      std::optional<const at::Tensor>& rotary_sin,
+      std::optional<const at::Tensor>& cache_batch_idx,
+      std::optional<const at::Tensor>& cache_leftpad,
+      std::optional<at::Tensor>& block_table,
+      std::optional<at::Tensor>& alibi_slopes, std::optional<at::Tensor>& out,
+      float softmax_scale, bool causal, int window_size_left,
+      int window_size_right, float softcap, bool rotary_interleaved,
+      int num_splits);
+};
+
+}  // namespace infini::ops::linked::torch::nvidia
+
+namespace infini::ops::linked::torch {
+
+extern template class TorchFlashAttnWithKvcache<
+    ::infini::ops::linked::torch::nvidia::FlashAttnKvcache>;
+
+}  // namespace infini::ops::linked::torch
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnWithKvcache, Device::Type::kNvidia, 16>
+    : public linked::torch::TorchFlashAttnWithKvcache<
+          linked::torch::nvidia::FlashAttnKvcache> {
+ public:
+  using linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::nvidia::FlashAttnKvcache>::TorchFlashAttnWithKvcache;
+
+  using linked::torch::TorchFlashAttnWithKvcache<
+      linked::torch::nvidia::FlashAttnKvcache>::operator();
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_

src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.yaml

@@ -0,0 +1,4 @@
+library: flash_attn
+required_symbols:
+  - >-
+    flash::mha_fwd_kvcache(at::Tensor&, at::Tensor const&, at::Tensor const&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor const>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, std::optional<at::Tensor>&, float, bool, int, int, float, bool, int)