feat: support graph on moore

Author: wooway777

python/infinicore/ops/moore_mate_flash_attn.py

@@ -141,7 +141,8 @@ def moore_mate_flash_attn_decode(
     scale: float,
     block_size: int,
     max_seq_len: int,
-) -> torch.Tensor:
+    return_state: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, ...]:
     """
     Decode entry point with native flash_attn KV cache layout (B, P, H, D).
     No layout conversion is performed.
@@ -174,7 +175,7 @@ def moore_mate_flash_attn_decode(
         pack_gqa=pack_gqa,
     )
 
-    out, *_ = flash_attn_with_kvcache(
+    result = flash_attn_with_kvcache(
         q=q,
         k_cache=k_cache,
         v_cache=v_cache,
@@ -188,7 +189,18 @@ def moore_mate_flash_attn_decode(
         pack_gqa=pack_gqa,
         return_softmax_lse=True,
     )
-    return out
+    if not return_state:
+        return result[0]
+
+    # Mate allocates output/workspace and scheduler tensors while the graph is
+    # captured. Keep every referenced allocation alive for graph replay.
+    return (
+        *result,
+        cache_seqlens,
+        page_table,
+        cu_seqlens_q,
+        metadata,
+    )
 
 
 # =============================================================================

src/infinicore/ops/mha_kvcache/moore/mha_kvcache_flashattn_moore.cc

@@ -8,6 +8,7 @@
 #include <pybind11/pybind11.h>
 #include <stdexcept>
 #include <torch/csrc/utils/pybind.h>
+#include <vector>
 
 namespace infinicore::op::mha_kvcache_impl::flashattn_moore {
 
@@ -37,6 +38,7 @@ struct PlannedMeta {
     graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table;
     std::optional<graph::GraphTensor> alibi_slopes;
     float scale;
+    std::vector<at::Tensor> mate_state;
 };
 
 void *plan(Tensor out,
@@ -89,20 +91,26 @@ void run(void *planned_meta) {
         py::object py_seqlens_k = py::cast(seqlens_k);
         py::object py_blk_tbl = py::cast(block_table);
 
-        py::object result = wrapper.attr("moore_mate_flash_attn_decode")(
+        py::tuple result = wrapper.attr("moore_mate_flash_attn_decode")(
             py_q,
             py_k_cache,
             py_v_cache,
             py_blk_tbl,
             py_seqlens_k,
             p->scale,
             block_size,
-            max_seq_len);
+            max_seq_len,
+            true);
 
-        at::Tensor result_t = result.cast<at::Tensor>();
+        at::Tensor result_t = result[0].cast<at::Tensor>();
         out_tensor.copy_(result_t.unsqueeze(1));
 
-        result = py::none();
+        p->mate_state.clear();
+        p->mate_state.reserve(result.size());
+        for (py::handle item : result) {
+            p->mate_state.push_back(item.cast<at::Tensor>());
+        }
+
         py_q = py_k_cache = py_v_cache = py_seqlens_k = py_blk_tbl = py::none();
     } catch (const py::error_already_set &e) {
         throw std::runtime_error(

src/infiniop/elementwise/moore/elementwise_moore.h

@@ -6,12 +6,19 @@
 #include "../../devices/moore/moore_kernel_common.h"
 #include "elementwise_moore_api.h"
 
+#include <algorithm>
+
 namespace op::elementwise::moore {
 template <typename T>
 __device__ __forceinline__ const T *typedInputPtr(const void *ptr) {
     return reinterpret_cast<const T *>(ptr);
 }
 
+template <size_t N>
+struct InputPointerArray {
+    const void *values[N];
+};
+
 __device__ __forceinline__ size_t getOutputIndex(size_t idx, bool is_contiguous, size_t ndim,
                                                  const size_t *shape, const ptrdiff_t *strides) {
     return is_contiguous ? idx : device::moore::indexToOffset(idx, ndim, shape, strides);
@@ -50,20 +57,20 @@ INFINIOP_MOORE_KERNEL elementwiseKernel(
     const ptrdiff_t *__restrict__ output_strides,
     const ptrdiff_t *__restrict__ input_strides,
     Tdata *output,
-    const void *const *inputs,
+    InputPointerArray<N> inputs,
     size_t offset,
     Args... args) {
 
     size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
 
     if (idx < output_size) {
-        const Tdata *const *typed_inputs = reinterpret_cast<const Tdata *const *>(inputs);
         size_t out_idx = getOutputIndex(idx, output_contiguous, ndim, output_shape, output_strides);
         InputIndexer indexer{idx, ndim, input_contiguous, input_broadcasted, input_shapes, input_strides, output_strides};
 
         unpackInputsAndApply(
             [&](auto... Is) {
-                output[out_idx] = Op{}(typed_inputs[Is.value][indexer(Is.value)]..., std::forward<Args>(args)...);
+                output[out_idx] = Op{}(
+                    typedInputPtr<Tdata>(inputs.values[Is.value])[indexer(Is.value)]..., std::forward<Args>(args)...);
             },
             std::make_index_sequence<N>{});
     }
@@ -81,7 +88,7 @@ INFINIOP_MOORE_KERNEL elementwiseKernel(
     const ptrdiff_t *__restrict__ output_strides,
     const ptrdiff_t *__restrict__ input_strides,
     Tout *output,
-    const void *const *__restrict__ inputs,
+    InputPointerArray<sizeof...(Tin)> inputs,
     size_t offset) {
 
     size_t idx = blockIdx.x * blockDim.x + threadIdx.x + offset;
@@ -93,17 +100,56 @@ INFINIOP_MOORE_KERNEL elementwiseKernel(
         unpackInputsAndApply(
             [&](auto... Is) {
                 output[out_idx] = Op{}.template operator()<Tout, Tin...>(
-                    (typedInputPtr<Tin>(inputs[Is.value])[indexer(Is.value)])...);
+                    (typedInputPtr<Tin>(inputs.values[Is.value])[indexer(Is.value)])...);
             },
             std::index_sequence_for<Tin...>{});
     }
 }
 
 struct DeviceImpl::Opaque {
     std::shared_ptr<device::moore::Handle::Internal> internal;
+    void *device_meta = nullptr;
+    const bool *input_contiguous = nullptr;
+    const bool *input_broadcasted = nullptr;
+    const size_t *output_shape = nullptr;
+    const ptrdiff_t *output_strides = nullptr;
+    const size_t *input_shapes = nullptr;
+    const ptrdiff_t *input_strides = nullptr;
+    infiniStatus_t init_status = INFINI_STATUS_SUCCESS;
+
+    Opaque(const std::shared_ptr<device::moore::Handle::Internal> &internal_,
+           const op::elementwise::ElementwiseInfo &info)
+        : internal(internal_), init_status(initialize(info)) {}
+
+    ~Opaque() {
+        if (device_meta != nullptr) {
+            musaFree(device_meta);
+        }
+    }
+
+    infiniStatus_t initialize(const op::elementwise::ElementwiseInfo &info) {
+        const auto meta_size = info.getMetaMemSize();
+        if (meta_size == 0) {
+            return INFINI_STATUS_SUCCESS;
+        }
+
+        CHECK_MOORE(musaMalloc(&device_meta, meta_size));
+        CHECK_MOORE(musaMemcpy(device_meta,
+                               info.getMetaStart(),
+                               meta_size,
+                               musaMemcpyHostToDevice));
+
+        const auto ndim = info.getNdim();
+        const auto input_size = info.getInputSize();
+        output_shape = reinterpret_cast<const size_t *>(device_meta);
+        output_strides = reinterpret_cast<const ptrdiff_t *>(output_shape + ndim);
+        input_shapes = reinterpret_cast<const size_t *>(output_strides + ndim);
+        input_strides = reinterpret_cast<const ptrdiff_t *>(input_shapes + input_size * ndim);
+        input_contiguous = reinterpret_cast<const bool *>(input_strides + input_size * ndim);
+        input_broadcasted = input_contiguous + input_size;
 
-    Opaque(const std::shared_ptr<device::moore::Handle::Internal> &internal)
-        : internal(internal) {}
+        return INFINI_STATUS_SUCCESS;
+    }
 
     template <uint32_t BLOCK_SIZE, size_t N, typename Op, typename Tdata, typename... Args>
     infiniStatus_t calculateImpl(const op::elementwise::ElementwiseInfo &info,
@@ -136,42 +182,6 @@ struct DeviceImpl::Opaque {
     }
 
 private:
-    template <size_t N>
-    infiniStatus_t infoToDevice(
-        const op::elementwise::ElementwiseInfo &info,
-        void *workspace,
-        const void *const *h_inputs_arr,
-        const void **&d_inputs_arr,
-        const bool *&d_input_contiguous,
-        const bool *&d_input_broadcasted,
-        const size_t *&d_output_shape,
-        const ptrdiff_t *&d_output_strides,
-        const size_t *&d_input_shapes,
-        const ptrdiff_t *&d_input_strides,
-        musaStream_t stream) const {
-
-        constexpr auto input_size = N;
-        const auto ndim = info.getNdim();
-        constexpr auto input_arr_size = N * sizeof(*h_inputs_arr);
-        const int8_t *info_meta_start = info.getMetaStart();
-        const int8_t *d_meta_start = reinterpret_cast<int8_t *>(workspace) + input_arr_size;
-
-        // copy the input pointer array and meta to device
-        CHECK_MOORE(musaMemcpyAsync(workspace, h_inputs_arr, input_arr_size, musaMemcpyHostToDevice, stream));
-        CHECK_MOORE(musaMemcpyAsync((void *)d_meta_start, info_meta_start, info.getMetaMemSize(), musaMemcpyHostToDevice, stream));
-
-        // offset/assign the pointers
-        d_inputs_arr = reinterpret_cast<const void **>(workspace);
-        d_output_shape = reinterpret_cast<const size_t *>(d_meta_start);
-        d_output_strides = reinterpret_cast<const ptrdiff_t *>(d_output_shape + ndim);
-        d_input_shapes = reinterpret_cast<const size_t *>(d_output_strides + ndim);
-        d_input_strides = reinterpret_cast<const ptrdiff_t *>(d_input_shapes + input_size * ndim);
-        d_input_contiguous = reinterpret_cast<const bool *>(d_input_strides + input_size * ndim);
-        d_input_broadcasted = reinterpret_cast<const bool *>(d_input_contiguous + input_size);
-
-        return INFINI_STATUS_SUCCESS;
-    }
-
     template <uint32_t BLOCK_SIZE, size_t N, typename KernelFunc, typename Tout, typename... Args>
     infiniStatus_t launchElementwiseKernel(
         const op::elementwise::ElementwiseInfo &info,
@@ -187,19 +197,10 @@ struct DeviceImpl::Opaque {
             return INFINI_STATUS_SUCCESS;
         }
 
-        // Device pointers
-        const void **d_inputs_arr = nullptr;
-        const bool *d_input_contiguous = nullptr;
-        const bool *d_input_broadcasted = nullptr;
-        const size_t *d_output_shape = nullptr;
-        const ptrdiff_t *d_output_strides = nullptr;
-        const size_t *d_input_shapes = nullptr;
-        const ptrdiff_t *d_input_strides = nullptr;
-
-        CHECK_STATUS(infoToDevice<N>(info, workspace, inputs.data(), d_inputs_arr,
-                                     d_input_contiguous, d_input_broadcasted,
-                                     d_output_shape, d_output_strides,
-                                     d_input_shapes, d_input_strides, stream));
+        (void)workspace;
+        CHECK_OR_RETURN(inputs.size() == N, INFINI_STATUS_BAD_PARAM);
+        InputPointerArray<N> input_ptrs{};
+        std::copy_n(inputs.begin(), N, input_ptrs.values);
 
         dim3 blockDims(std::min(BLOCK_SIZE, static_cast<uint32_t>(internal->maxThreadsPerBlock())));
         dim3 gridDims(std::min(uint32_t(CEIL_DIV(output_size, blockDims.x)), static_cast<uint32_t>(internal->gridSizeX())));
@@ -208,10 +209,10 @@ struct DeviceImpl::Opaque {
         for (size_t i = 0; i < output_size; i += step) {
             kernel_func<<<gridDims, blockDims, 0, stream>>>(
                 output_size, info.getNdim(), info.isOutputContiguous(),
-                d_input_contiguous, d_input_broadcasted,
-                d_output_shape, d_input_shapes,
-                d_output_strides, d_input_strides,
-                output, reinterpret_cast<const void **>(d_inputs_arr),
+                input_contiguous, input_broadcasted,
+                output_shape, input_shapes,
+                output_strides, input_strides,
+                output, input_ptrs,
                 i, std::forward<Args>(args)...);
         }
 
@@ -222,6 +223,9 @@ struct DeviceImpl::Opaque {
 template <typename... Args>
 utils::Result<DeviceImpl *> DeviceImpl::create(Args &&...args) {
     auto opaque = std::make_shared<Opaque>(std::forward<Args>(args)...);
+    if (opaque->init_status != INFINI_STATUS_SUCCESS) {
+        return opaque->init_status;
+    }
     return utils::Result<DeviceImpl *>(new DeviceImpl(opaque));
 }
 

src/infiniop/elementwise/moore/elementwise_moore_api.h

@@ -38,22 +38,22 @@ class DeviceImpl final {
         Args &&...args);
 };
 } // namespace op::elementwise::moore
-#define CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC)          \
-                                                                                              \
-    auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC);    \
-    CHECK_RESULT(info_result);                                                                \
-    auto info = info_result.take();                                                           \
-    auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *);       \
-                                                                                              \
-    auto device_impl_result = op::elementwise::moore::DeviceImpl::create(HANDLE->internal()); \
-    CHECK_RESULT(device_impl_result);                                                         \
-                                                                                              \
-    *desc_ptr = new Descriptor(                                                               \
-        DTYPE,                                                                                \
-        std::move(info),                                                                      \
-        std::move(device_impl_result.take()),                                                 \
-        workspace_size,                                                                       \
-        HANDLE->device,                                                                       \
+#define CREATE_ELEMENTWISE_MOORE_DESCRIPTOR(HANDLE, DTYPE, OUT_DESC, INPUT_DESC_VEC)                \
+                                                                                                    \
+    auto info_result = op::elementwise::ElementwiseInfo::create(OUT_DESC, INPUT_DESC_VEC);          \
+    CHECK_RESULT(info_result);                                                                      \
+    auto info = info_result.take();                                                                 \
+    auto workspace_size = 0;                                                                        \
+                                                                                                    \
+    auto device_impl_result = op::elementwise::moore::DeviceImpl::create(HANDLE->internal(), info); \
+    CHECK_RESULT(device_impl_result);                                                               \
+                                                                                                    \
+    *desc_ptr = new Descriptor(                                                                     \
+        DTYPE,                                                                                      \
+        std::move(info),                                                                            \
+        std::move(device_impl_result.take()),                                                       \
+        workspace_size,                                                                             \
+        HANDLE->device,                                                                             \
         HANDLE->device_id);
 
 #endif // __INFINIOP_ELEMENTWISE_MOORE_API_H__

src/infiniop/ops/hardtanh/moore/hardtanh_moore.mu

@@ -94,9 +94,9 @@ infiniStatus_t Descriptor::create(
     auto info_result = op::elementwise::ElementwiseInfo::create(out_desc, input_desc_vec);
     CHECK_RESULT(info_result);
     auto info = info_result.take();
-    auto workspace_size = info.getMetaMemSize() + info.getInputSize() * sizeof(void *);
+    auto workspace_size = 0;
 
-    auto device_impl_result = op::elementwise::moore::DeviceImpl::create(handle->internal());
+    auto device_impl_result = op::elementwise::moore::DeviceImpl::create(handle->internal(), info);
     CHECK_RESULT(device_impl_result);
 
     *desc_ptr = new Descriptor(