feat(moore): support flash_attn_varlen_func

CMakeLists.txt

@@ -263,6 +263,30 @@ if(WITH_TORCH)
     find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED)
     set(TORCH_LIBRARIES ${TORCH_LIB} ${TORCH_CPU_LIB} ${C10_LIB})
 
+    if(WITH_MOORE)
+        execute_process(
+            COMMAND ${_TORCH_PYTHON} -c
+                "import pathlib, torch_musa; print(pathlib.Path(torch_musa.__file__).resolve().parent / 'share' / 'cmake' / 'TorchMusa')"
+            OUTPUT_VARIABLE TorchMusa_DIR
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            RESULT_VARIABLE _torch_musa_result
+        )
+        if(NOT _torch_musa_result EQUAL 0 OR NOT TorchMusa_DIR)
+            message(FATAL_ERROR "`WITH_MOORE` and `WITH_TORCH` require `torch_musa`.")
+        endif()
+        set(_infini_ops_torch_libraries "${TORCH_LIBRARIES}")
+        find_package(TorchMusa CONFIG REQUIRED)
+        set(TORCH_LIBRARIES "${_infini_ops_torch_libraries}")
+        unset(_infini_ops_torch_libraries)
+        get_filename_component(_torch_musa_lib_dir
+            "${TorchMusa_LIBRARY}" DIRECTORY)
+        list(APPEND TORCH_RUNTIME_DIRS "${_torch_musa_lib_dir}")
+        # Retain TorchMusa's dispatcher registrations. The active Moore block
+        # links the MUSA toolkit/runtime selected by InfiniOps.
+        list(APPEND TORCH_LIBRARIES
+            -Wl,--no-as-needed ${TorchMusa_LIBRARY} -Wl,--as-needed)
+    endif()
+
     if(WITH_NVIDIA)
         find_library(TORCH_CUDA_LIB torch_cuda HINTS ${_torch_lib_dirs} REQUIRED)
         find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED)

docs/build.md

@@ -58,6 +58,14 @@ toolchain. Operator-pruned builds require it only when
 `flash_attn_with_kvcache` is selected; the default smoke subset does not select
 this operator.
 
+Builds with both `WITH_MOORE=ON` and `WITH_TORCH=ON` require an installed
+TorchMusa package that provides `TorchMusaConfig.cmake`. The Moore
+`flash_attn_varlen_func` backend uses TorchMusa's ATen FlashAttention kernel.
+Runtime use additionally requires the provider to register
+`aten::_flash_attention_forward` for PrivateUse1. That provider currently
+supports global attention only; causal attention also requires matching query
+and key sequence lengths.
+
 The built wheel installs the InfiniOps Python extension and the InfiniRT shared
 library needed by the extension.
 

src/CMakeLists.txt

@@ -480,7 +480,7 @@ if(WITH_NVIDIA)
     list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum)
 endif()
 
-if(WITH_NVIDIA AND WITH_TORCH)
+if((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH)
     list(APPEND _infini_ops_smoke_ops flash_attn_varlen_func)
 endif()
 
@@ -618,14 +618,17 @@ if(WITH_TORCH)
         "${PROJECT_SOURCE_DIR}/generated/torch/*.cpp"
     )
 
-    if(NOT WITH_NVIDIA)
+    if(NOT WITH_NVIDIA AND NOT WITH_MOORE)
         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()
+    if(NOT WITH_MOORE)
+        list(FILTER TORCH_SOURCES EXCLUDE REGEX ".*/torch/moore/.*")
+    endif()
 
     set(INFINI_OPS_TORCH_UNITY_BATCH_SIZE "8" CACHE STRING
         "Number of torch sources to include in each generated unity translation unit; set to 1 to disable")
@@ -717,7 +720,8 @@ if(WITH_TORCH)
             list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include")
         endif()
 
-        # Vendor-specific defines required by forked `torch` headers.
+        # Select vendor-specific generated implementations and any definitions
+        # required by forked `torch` headers.
         set(_torch_extra_flags "")
         if(WITH_METAX)
             list(APPEND _torch_extra_flags "-DUSE_MACA=1" "-DWITH_METAX=1")

src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc

@@ -0,0 +1,40 @@
+#include "torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h"
+
+#include <ATen/ops/equal.h>
+#include <c10/util/Exception.h>
+
+#include "torch/tensor_.h"
+
+namespace infini::ops {
+
+void Operator<FlashAttnVarlenFunc, Device::Type::kMoore, 8>::operator()(
+    const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
+    const Tensor cu_seqlens_k, 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 std::optional<Tensor> alibi_slopes, const bool deterministic,
+    const bool return_attn_probs, const std::optional<Tensor> block_table,
+    Tensor out) const {
+  TORCH_CHECK(window_size[0] < 0 && window_size[1] < 0,
+              "TorchMusa FlashAttention does not support local windows");
+
+  if (causal) {
+    auto at_cu_seqlens_q = ToAtenTensor<Device::Type::kMoore>(
+        const_cast<void*>(cu_seqlens_q.data()), cu_seqlens_q_shape_,
+        cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_);
+    auto at_cu_seqlens_k = ToAtenTensor<Device::Type::kMoore>(
+        const_cast<void*>(cu_seqlens_k.data()), cu_seqlens_k_shape_,
+        cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
+    TORCH_CHECK(at::equal(at_cu_seqlens_q, at_cu_seqlens_k),
+                "TorchMusa FlashAttention requires matching query and key "
+                "sequence lengths for causal attention");
+  }
+
+  AtenFlashAttnVarlenFunc<Device::Type::kMoore>::operator()(
+      q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k,
+      dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes,
+      deterministic, return_attn_probs, block_table, out);
+}
+
+}  // namespace infini::ops

src/torch/moore/ops/flash_attn_varlen_func/flash_attn_varlen_func.h

@@ -0,0 +1,29 @@
+#ifndef INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_
+#define INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_
+
+#include "torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kMoore, 8>
+    : public AtenFlashAttnVarlenFunc<Device::Type::kMoore> {
+ public:
+  using AtenFlashAttnVarlenFunc::AtenFlashAttnVarlenFunc;
+  using AtenFlashAttnVarlenFunc::operator();
+
+  void operator()(const Tensor q, const Tensor k, const Tensor v,
+                  const Tensor cu_seqlens_q, const Tensor cu_seqlens_k,
+                  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 std::optional<Tensor> alibi_slopes,
+                  const bool deterministic, const bool return_attn_probs,
+                  const std::optional<Tensor> block_table,
+                  Tensor out) const override;
+};
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_TORCH_MOORE_FLASH_ATTN_VARLEN_FUNC_H_

src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.cc

@@ -1,17 +1,15 @@
 #include "torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h"
 
 #include <ATen/ops/_flash_attention_forward.h>
-#include <c10/cuda/CUDAGuard.h>
-#include <c10/cuda/CUDAStream.h>
-#include <cuda_runtime_api.h>
 
 #include <tuple>
 
 #include "torch/tensor_.h"
 
 namespace infini::ops {
 
-void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::operator()(
+template <Device::Type kDev>
+void AtenFlashAttnVarlenFunc<kDev>::operator()(
     const Tensor q, const Tensor k, const Tensor v, const Tensor cu_seqlens_q,
     const Tensor cu_seqlens_k, const int64_t max_seqlen_q,
     const int64_t max_seqlen_k, const double dropout_p,
@@ -26,50 +24,45 @@ void Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>::operator()(
   (void)return_attn_probs;
   (void)block_table;
 
-  const auto device_index = static_cast<c10::DeviceIndex>(device_index_);
-  const c10::cuda::CUDAGuard device_guard{device_index};
+  auto at_q = ToAtenTensor<kDev>(const_cast<void*>(q.data()), q_shape_,
+                                 q_strides_, q_dtype_, device_index_);
+  auto at_k = ToAtenTensor<kDev>(const_cast<void*>(k.data()), k_shape_,
+                                 k_strides_, k_dtype_, device_index_);
+  auto at_v = ToAtenTensor<kDev>(const_cast<void*>(v.data()), v_shape_,
+                                 v_strides_, v_dtype_, device_index_);
+  auto at_cu_seqlens_q = ToAtenTensor<kDev>(
+      const_cast<void*>(cu_seqlens_q.data()), cu_seqlens_q_shape_,
+      cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_);
+  auto at_cu_seqlens_k = ToAtenTensor<kDev>(
+      const_cast<void*>(cu_seqlens_k.data()), cu_seqlens_k_shape_,
+      cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
+  auto at_out = ToAtenTensor<kDev>(out.data(), out_shape_, out_strides_,
+                                   out_dtype_, device_index_);
 
-  const auto run = [&]() {
-    auto at_q = ToAtenTensor<Device::Type::kNvidia>(const_cast<void*>(q.data()),
-                                                    q_shape_, q_strides_,
-                                                    q_dtype_, device_index_);
-    auto at_k = ToAtenTensor<Device::Type::kNvidia>(const_cast<void*>(k.data()),
-                                                    k_shape_, k_strides_,
-                                                    k_dtype_, device_index_);
-    auto at_v = ToAtenTensor<Device::Type::kNvidia>(const_cast<void*>(v.data()),
-                                                    v_shape_, v_strides_,
-                                                    v_dtype_, device_index_);
-    auto at_cu_seqlens_q = ToAtenTensor<Device::Type::kNvidia>(
-        const_cast<void*>(cu_seqlens_q.data()), cu_seqlens_q_shape_,
-        cu_seqlens_q_strides_, cu_seqlens_q_dtype_, device_index_);
-    auto at_cu_seqlens_k = ToAtenTensor<Device::Type::kNvidia>(
-        const_cast<void*>(cu_seqlens_k.data()), cu_seqlens_k_shape_,
-        cu_seqlens_k_strides_, cu_seqlens_k_dtype_, device_index_);
-    auto at_out = ToAtenTensor<Device::Type::kNvidia>(
-        out.data(), out_shape_, out_strides_, out_dtype_, device_index_);
+  std::optional<int64_t> window_size_left;
+  std::optional<int64_t> window_size_right;
+  if (window_size[0] >= 0) {
+    window_size_left = window_size[0];
+  }
+  if (window_size[1] >= 0) {
+    window_size_right = window_size[1];
+  }
+  if (causal &&
+      (window_size_left.has_value() || window_size_right.has_value())) {
+    window_size_right = 0;
+  }
 
-    const std::optional<int64_t> window_size_left =
-        window_size[0] < 0 ? std::nullopt
-                           : std::optional<int64_t>{window_size[0]};
-    const std::optional<int64_t> window_size_right =
-        causal               ? std::optional<int64_t>{0}
-        : window_size[1] < 0 ? std::nullopt
-                             : std::optional<int64_t>{window_size[1]};
+  auto result = at::_flash_attention_forward(
+      at_q, at_k, at_v, at_cu_seqlens_q, at_cu_seqlens_k, max_seqlen_q,
+      max_seqlen_k, dropout_p, causal, false, softmax_scale, window_size_left,
+      window_size_right, std::nullopt, std::nullopt);
 
-    auto result = at::_flash_attention_forward(
-        at_q, at_k, at_v, at_cu_seqlens_q, at_cu_seqlens_k, max_seqlen_q,
-        max_seqlen_k, dropout_p, causal, false, softmax_scale, window_size_left,
-        window_size_right, std::nullopt, std::nullopt);
-
-    // ATen owns the returned tensor. Keep the InfiniOps trailing-output ABI by
-    // copying it into the caller-provided buffer on the selected CUDA stream.
-    at_out.copy_(std::get<0>(result));
-  };
-
-  const c10::cuda::CUDAStreamGuard stream_guard{
-      c10::cuda::getStreamFromExternal(reinterpret_cast<cudaStream_t>(stream_),
-                                       device_index)};
-  run();
+  // ATen owns the returned tensor. Keep the InfiniOps trailing-output ABI by
+  // copying it into the caller-provided buffer on the current ATen stream.
+  at_out.copy_(std::get<0>(result));
 }
 
+template class AtenFlashAttnVarlenFunc<Device::Type::kNvidia>;
+template class AtenFlashAttnVarlenFunc<Device::Type::kMoore>;
+
 }  // namespace infini::ops

src/torch/ops/flash_attn_varlen_func/flash_attn_varlen_func.h

@@ -5,9 +5,8 @@
 
 namespace infini::ops {
 
-template <>
-class Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>
-    : public FlashAttnVarlenFunc {
+template <Device::Type kDev>
+class AtenFlashAttnVarlenFunc : public FlashAttnVarlenFunc {
  public:
   using FlashAttnVarlenFunc::FlashAttnVarlenFunc;
   using FlashAttnVarlenFunc::operator();
@@ -24,6 +23,14 @@ class Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>
                   Tensor out) const override;
 };
 
+template <>
+class Operator<FlashAttnVarlenFunc, Device::Type::kNvidia, 8>
+    : public AtenFlashAttnVarlenFunc<Device::Type::kNvidia> {
+ public:
+  using AtenFlashAttnVarlenFunc::AtenFlashAttnVarlenFunc;
+  using AtenFlashAttnVarlenFunc::operator();
+};
+
 }  // namespace infini::ops
 
 #endif  // INFINI_OPS_TORCH_FLASH_ATTN_VARLEN_FUNC_H_