From 5c2d56426fbbc46bd45e154f20603242d451a535 Mon Sep 17 00:00:00 2001 From: Kingbo_998 Date: Tue, 4 Aug 2026 16:54:25 +0800 Subject: [PATCH 1/4] fix graph of ascend --- include/infinicore/graph/graph.hpp | 39 +++- include/infinirt.h | 9 + src/infinicore/graph/graph.cc | 170 +++++++++++++++++- .../ascend/mha_kvcache_flashattn_ascend.cc | 127 ++++++++++--- src/infinicore/ops/mha_kvcache/mha_kvcache.cc | 5 + src/infinirt/ascend/infinirt_ascend.cc | 28 +++ src/infinirt/ascend/infinirt_ascend.h | 8 + src/infinirt/infinirt.cc | 40 +++++ 8 files changed, 391 insertions(+), 35 deletions(-) diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index be368f92d..e79de5e3d 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -1,6 +1,8 @@ #pragma once +#include #include +#include #include #include "../tensor.hpp" @@ -17,20 +19,33 @@ class GraphTensor : public Tensor { class GraphOperator { public: virtual void run() const = 0; + virtual bool requires_task_update() const { + return false; + } virtual ~GraphOperator() = default; }; class DispatchableGraphOperator : public GraphOperator { public: void run() const override; + bool requires_task_update() const override { + return requires_task_update_; + } ~DispatchableGraphOperator() override; protected: + void enable_task_update() { + requires_task_update_ = true; + } + using run_schema = void (*)(void *); using cleanup_schema = void (*)(void **); - void *planned_meta_; - run_schema runner_; - cleanup_schema deleter_; + void *planned_meta_ = nullptr; + run_schema runner_ = nullptr; + cleanup_schema deleter_ = nullptr; + +private: + bool requires_task_update_ = false; }; class Graph { @@ -39,6 +54,10 @@ class Graph { ~Graph(); void run() const; + void update_host_int_array( + const Tensor &device_tensor, + const int32_t *host_values, + size_t count); protected: void add_operator(std::shared_ptr op); @@ -50,7 +69,21 @@ class Graph { private: struct DeviceGraph; std::unique_ptr device_graph_; + std::unordered_map> host_int_arrays_; }; + +bool is_task_updating(); +bool is_task_group_capturing(); +void begin_task_group_capture(); +void end_task_group_capture(); +void begin_task_update(); +void end_task_update(); +void stage_task_update_host_int_array( + const Tensor &device_tensor, + const int32_t *host_values, + size_t count); +const std::vector *lookup_task_update_host_int_array( + const Tensor &tensor); } // namespace infinicore::graph #define INFINICORE_GRAPH_OP_CLASS(__OP_NAME__, ...) \ diff --git a/include/infinirt.h b/include/infinirt.h index 4ada848f5..778d5b3db 100644 --- a/include/infinirt.h +++ b/include/infinirt.h @@ -9,6 +9,7 @@ typedef void *infinirtEvent_t; typedef void *infinirtGraph_t; typedef void *infinirtGraphNode_t; typedef void *infinirtGraphExec_t; +typedef void *infinirtGraphTaskGroup_t; // Bitmask describing which fields of an infinirtDeviceResourceSnapshot_t // have been populated by the backend. Backends without a particular @@ -130,5 +131,13 @@ __INFINI_C __export infiniStatus_t infinirtGraphInstantiate( size_t buffer_size); __INFINI_C __export infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exec); __INFINI_C __export infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtGraphTaskGroupBegin(infinirtStream_t stream); +__INFINI_C __export infiniStatus_t infinirtGraphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle); +__INFINI_C __export infiniStatus_t infinirtGraphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle); +__INFINI_C __export infiniStatus_t infinirtGraphTaskUpdateEnd(infinirtStream_t stream); #endif // __INFINIRT_API_H__ diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 4552f2d67..5f5231ee4 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -5,6 +5,66 @@ #include namespace infinicore::graph { +namespace { + +using HostIntArrayMap = std::unordered_map>; + +thread_local const HostIntArrayMap *current_host_int_arrays = nullptr; +thread_local bool task_update_active = false; +thread_local HostIntArrayMap staged_host_int_arrays; +thread_local infinirtGraphTaskGroup_t *capture_task_group_handle = nullptr; +thread_local infinirtGraphTaskGroup_t task_update_handle = nullptr; + +class HostIntArrayScope { +public: + HostIntArrayScope( + const HostIntArrayMap *host_int_arrays, + bool is_task_update) { + current_host_int_arrays = host_int_arrays; + task_update_active = is_task_update; + } + + ~HostIntArrayScope() { + task_update_active = false; + current_host_int_arrays = nullptr; + } + + HostIntArrayScope(const HostIntArrayScope &) = delete; + HostIntArrayScope &operator=(const HostIntArrayScope &) = delete; +}; + +class TaskGroupCaptureScope { +public: + explicit TaskGroupCaptureScope(infinirtGraphTaskGroup_t *handle) { + INFINICORE_ASSERT(capture_task_group_handle == nullptr); + capture_task_group_handle = handle; + } + + ~TaskGroupCaptureScope() { + capture_task_group_handle = nullptr; + } + + TaskGroupCaptureScope(const TaskGroupCaptureScope &) = delete; + TaskGroupCaptureScope &operator=(const TaskGroupCaptureScope &) = delete; +}; + +class TaskUpdateHandleScope { +public: + explicit TaskUpdateHandleScope(infinirtGraphTaskGroup_t handle) { + INFINICORE_ASSERT(task_update_handle == nullptr); + INFINICORE_ASSERT(handle != nullptr); + task_update_handle = handle; + } + + ~TaskUpdateHandleScope() { + task_update_handle = nullptr; + } + + TaskUpdateHandleScope(const TaskUpdateHandleScope &) = delete; + TaskUpdateHandleScope &operator=(const TaskUpdateHandleScope &) = delete; +}; + +} // namespace /* ========================= * GraphTensor @@ -32,10 +92,16 @@ DispatchableGraphOperator::~DispatchableGraphOperator() { * ========================= */ struct Graph::DeviceGraph { + struct UpdatableTask { + std::shared_ptr op; + infinirtGraphTaskGroup_t handle; + }; + infinirtGraph_t graph; infinirtGraphExec_t exec; infinirtGraphNode_t node; std::vector log_buffer; + std::vector updatable_tasks; DeviceGraph() : graph(nullptr), exec(nullptr), node(nullptr) { log_buffer.resize(4 * 1024); @@ -55,11 +121,18 @@ struct Graph::DeviceGraph { } }; -Graph::Graph() { +Graph::Graph() + : host_int_arrays_(std::move(staged_host_int_arrays)) { + staged_host_int_arrays.clear(); } void Graph::run() const { if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { + HostIntArrayScope update_scope(&host_int_arrays_, true); + for (const auto &task : device_graph_->updatable_tasks) { + TaskUpdateHandleScope task_update_scope(task.handle); + task.op->run(); + } device_graph_.get()->launch(); } else { for (auto &op : op_list_) { @@ -72,6 +145,80 @@ void Graph::add_operator(std::shared_ptr op) { op_list_.push_back(op); } +void Graph::update_host_int_array( + const Tensor &device_tensor, + const int32_t *host_values, + size_t count) { + INFINICORE_ASSERT(device_tensor); + INFINICORE_ASSERT(host_values != nullptr || count == 0); + + auto &values = host_int_arrays_[device_tensor->data()]; + values.resize(count); + for (size_t i = 0; i < count; ++i) { + values[i] = static_cast(host_values[i]); + } +} + +bool is_task_updating() { + return task_update_active; +} + +bool is_task_group_capturing() { + return capture_task_group_handle != nullptr; +} + +void begin_task_group_capture() { + INFINICORE_ASSERT(capture_task_group_handle != nullptr); + INFINICORE_ASSERT(*capture_task_group_handle == nullptr); + INFINICORE_CHECK_ERROR( + infinirtGraphTaskGroupBegin(context::getStream())); +} + +void end_task_group_capture() { + INFINICORE_ASSERT(capture_task_group_handle != nullptr); + INFINICORE_ASSERT(*capture_task_group_handle == nullptr); + INFINICORE_CHECK_ERROR(infinirtGraphTaskGroupEnd( + context::getStream(), capture_task_group_handle)); +} + +void begin_task_update() { + INFINICORE_ASSERT(task_update_handle != nullptr); + INFINICORE_CHECK_ERROR(infinirtGraphTaskUpdateBegin( + context::getStream(), task_update_handle)); +} + +void end_task_update() { + INFINICORE_ASSERT(task_update_handle != nullptr); + INFINICORE_CHECK_ERROR( + infinirtGraphTaskUpdateEnd(context::getStream())); +} + +void stage_task_update_host_int_array( + const Tensor &device_tensor, + const int32_t *host_values, + size_t count) { + INFINICORE_ASSERT(device_tensor); + INFINICORE_ASSERT(host_values != nullptr || count == 0); + + auto &values = staged_host_int_arrays[device_tensor->data()]; + values.resize(count); + for (size_t i = 0; i < count; ++i) { + values[i] = static_cast(host_values[i]); + } +} + +const std::vector *lookup_task_update_host_int_array( + const Tensor &tensor) { + if (current_host_int_arrays == nullptr || !tensor) { + return nullptr; + } + auto it = current_host_int_arrays->find(tensor->data()); + if (it == current_host_int_arrays->end()) { + return nullptr; + } + return &it->second; +} + void Graph::instantiate() { // Reset device graph device_graph_ = std::make_unique(); @@ -89,8 +236,25 @@ void Graph::instantiate() { return; } - // Run and record - this->run(); + // Run and record. Operators with dynamic host-side arguments are captured + // as individual ModelRI task groups so those arguments can be updated at + // replay time. + device_graph_->updatable_tasks.clear(); + HostIntArrayScope capture_scope(&host_int_arrays_, false); + for (const auto &op : op_list_) { + if (!op->requires_task_update()) { + op->run(); + continue; + } + + infinirtGraphTaskGroup_t handle = nullptr; + { + TaskGroupCaptureScope task_group_scope(&handle); + op->run(); + } + INFINICORE_ASSERT(handle != nullptr); + device_graph_->updatable_tasks.push_back({op, handle}); + } if (infinirtStreamEndCapture( context::getStream(), diff --git a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc index ead1ca4e2..5ce5d7803 100644 --- a/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc +++ b/src/infinicore/ops/mha_kvcache/ascend/mha_kvcache_flashattn_ascend.cc @@ -42,8 +42,64 @@ struct PlannedMeta { graph::GraphTensor out, q, k_cache, v_cache, seqlens_k, block_table; std::optional alibi_slopes; float scale; + Tensor out_work, q_work, k_work, v_work, block_table_work; }; +static Tensor persistent_contiguous_work_tensor(const Tensor &tensor) { + if (tensor->is_contiguous()) { + return Tensor(tensor); + } + return Tensor::empty( + tensor->shape(), tensor->dtype(), tensor->device()); +} + +static std::vector +get_actual_seq_lengths_k(const PlannedMeta *p, int64_t batch_size) { + if (const auto *bound = graph::lookup_task_update_host_int_array(p->seqlens_k)) { + if (bound->size() != static_cast(batch_size)) { + throw std::runtime_error( + "[mha_kvcache/ascend] bound actualSeqLengthsKv size does not " + "match the captured batch size"); + } + return *bound; + } + + if (graph::is_task_updating()) { + throw std::runtime_error( + "[mha_kvcache/ascend] missing host actualSeqLengthsKv binding " + "during graph task update"); + } + + auto seqlens_k_shape = p->seqlens_k->shape(); + if (seqlens_k_shape.size() != 1 + || seqlens_k_shape[0] != static_cast(batch_size)) { + throw std::runtime_error( + "[mha_kvcache/ascend] seqlens_k must be a 1D tensor whose length " + "matches the batch size"); + } + + std::vector seqlens_k_host(batch_size); + auto copy_ret = aclrtMemcpy( + seqlens_k_host.data(), + batch_size * sizeof(int32_t), + reinterpret_cast(p->seqlens_k->data()), + batch_size * sizeof(int32_t), + ACL_MEMCPY_DEVICE_TO_HOST); + if (copy_ret != ACL_SUCCESS) { + throw std::runtime_error( + std::string( + "[mha_kvcache/ascend] copy seqlens_k to host failed: ") + + std::to_string(copy_ret)); + } + + std::vector actual_seq_k_vec; + actual_seq_k_vec.reserve(batch_size); + for (int64_t i = 0; i < batch_size; ++i) { + actual_seq_k_vec.push_back(seqlens_k_host[i]); + } + return actual_seq_k_vec; +} + void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, const Tensor &v_cache, const Tensor &seqlens_k, const Tensor &block_table, std::optional alibi_slopes, @@ -57,7 +113,12 @@ void *plan(Tensor out, const Tensor &q, const Tensor &k_cache, alibi_slopes ? std::optional( graph::GraphTensor(*alibi_slopes)) : std::nullopt, - scale}; + scale, + persistent_contiguous_work_tensor(out), + persistent_contiguous_work_tensor(q), + persistent_contiguous_work_tensor(k_cache), + persistent_contiguous_work_tensor(v_cache), + persistent_contiguous_work_tensor(block_table)}; } void run(void *planned_meta) { @@ -98,38 +159,32 @@ void run(void *planned_meta) { "[mha_kvcache/ascend] k_cache and v_cache shapes are incompatible"); } - Tensor q_work = p->q->is_contiguous() ? Tensor(p->q) : p->q->contiguous(); - Tensor k_work = p->k_cache->is_contiguous() ? Tensor(p->k_cache) - : p->k_cache->contiguous(); - Tensor v_work = p->v_cache->is_contiguous() ? Tensor(p->v_cache) - : p->v_cache->contiguous(); - Tensor bt_work = p->block_table->is_contiguous() - ? Tensor(p->block_table) - : p->block_table->contiguous(); - Tensor out_work = p->out->is_contiguous() ? Tensor(p->out) : p->out->contiguous(); + const bool task_updating = graph::is_task_updating(); + if (!task_updating) { + if (!p->q->is_contiguous()) { + p->q_work->copy_from(p->q); + } + if (!p->k_cache->is_contiguous()) { + p->k_work->copy_from(p->k_cache); + } + if (!p->v_cache->is_contiguous()) { + p->v_work->copy_from(p->v_cache); + } + if (!p->block_table->is_contiguous()) { + p->block_table_work->copy_from(p->block_table); + } + } + Tensor q_work = p->q_work; + Tensor k_work = p->k_work; + Tensor v_work = p->v_work; + Tensor bt_work = p->block_table_work; + Tensor out_work = p->out_work; aclDataType q_dtype = to_acl_dtype(q_work->dtype()); - // Read seqlens_k to host - auto seqlens_k_shape = p->seqlens_k->shape(); - int64_t seqlens_k_len = seqlens_k_shape[0]; - std::vector seqlens_k_host(seqlens_k_len); - auto copy_ret = aclrtMemcpy(seqlens_k_host.data(), seqlens_k_len * sizeof(int32_t), - reinterpret_cast(p->seqlens_k->data()), - seqlens_k_len * sizeof(int32_t), ACL_MEMCPY_DEVICE_TO_HOST); - if (copy_ret != ACL_SUCCESS) { - throw std::runtime_error( - std::string("[mha_kvcache/ascend] copy seqlens_k to host failed: ") + std::to_string(copy_ret)); - } - - // Build actual_seq vectors std::vector actual_seq_q_vec(batch_size, 1); // decode: always 1 query token - std::vector actual_seq_k_vec; - actual_seq_k_vec.reserve(batch_size); - for (int64_t i = 0; i < batch_size; ++i) { - actual_seq_k_vec.push_back(seqlens_k_host[i]); - } + auto actual_seq_k_vec = get_actual_seq_lengths_k(p, batch_size); // BNSD [batch, num_heads, 1, head_size], viewed from BSND memory. std::vector q_dims = {batch_size, num_heads, 1, head_size}; @@ -248,8 +303,22 @@ void run(void *planned_meta) { .buf; } + const bool capture_task_group = graph::is_task_group_capturing(); + if (capture_task_group) { + graph::begin_task_group_capture(); + } + const bool update_task_group = graph::is_task_updating(); + if (update_task_group) { + graph::begin_task_update(); + } ret = aclnnFusedInferAttentionScoreV4(workspace, workspace_size, executor, stream); + if (update_task_group) { + graph::end_task_update(); + } + if (capture_task_group) { + graph::end_task_group_capture(); + } // Release aclTensor/aclTensorList/aclIntArray resources aclDestroyTensor(query_acl); @@ -269,7 +338,7 @@ void run(void *planned_meta) { } // Copy back if out was not contiguous - if (!p->out->is_contiguous()) { + if (!task_updating && !p->out->is_contiguous()) { p->out->copy_from(out_work); } } diff --git a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc index 0c5b3ae8c..b6d441520 100644 --- a/src/infinicore/ops/mha_kvcache/mha_kvcache.cc +++ b/src/infinicore/ops/mha_kvcache/mha_kvcache.cc @@ -16,6 +16,11 @@ MhaKVCache::MhaKVCache(Tensor out, INFINICORE_ASSERT_TENSORS_SAME_DEVICE(out, q, k_cache, v_cache, seqlens_k, block_table); INFINICORE_GRAPH_OP_DISPATCH(out->device().getType(), out, q, k_cache, v_cache, seqlens_k, block_table, alibi_slopes, scale); +#if defined(ENABLE_ASCEND_FLASH_ATTN) + if (out->device().getType() == Device::Type::ASCEND) { + enable_task_update(); + } +#endif } void MhaKVCache::execute(Tensor out, diff --git a/src/infinirt/ascend/infinirt_ascend.cc b/src/infinirt/ascend/infinirt_ascend.cc index 252b9b8e4..4217196d2 100644 --- a/src/infinirt/ascend/infinirt_ascend.cc +++ b/src/infinirt/ascend/infinirt_ascend.cc @@ -215,6 +215,34 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_SUCCESS; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + CHECK_ACLRT(aclmdlRICaptureTaskGrpBegin((aclrtStream)stream)); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + aclrtTaskGrp task_group = nullptr; + CHECK_ACLRT(aclmdlRICaptureTaskGrpEnd((aclrtStream)stream, &task_group)); + *handle = reinterpret_cast(task_group); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + CHECK_ACLRT(aclmdlRICaptureTaskUpdateBegin( + (aclrtStream)stream, + reinterpret_cast(handle))); + return INFINI_STATUS_SUCCESS; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + CHECK_ACLRT(aclmdlRICaptureTaskUpdateEnd((aclrtStream)stream)); + return INFINI_STATUS_SUCCESS; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/ascend/infinirt_ascend.h b/src/infinirt/ascend/infinirt_ascend.h index d7b2d7513..7b4d5976f 100644 --- a/src/infinirt/ascend/infinirt_ascend.h +++ b/src/infinirt/ascend/infinirt_ascend.h @@ -6,6 +6,14 @@ namespace infinirt::ascend { #ifdef ENABLE_ASCEND_API infiniStatus_t init(); INFINIRT_DEVICE_API_IMPL +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream); +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle); +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle); +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream); #else INFINIRT_DEVICE_API_NOOP #endif diff --git a/src/infinirt/infinirt.cc b/src/infinirt/infinirt.cc index 72614865b..cf485d2b6 100644 --- a/src/infinirt/infinirt.cc +++ b/src/infinirt/infinirt.cc @@ -249,3 +249,43 @@ __INFINI_C infiniStatus_t infinirtGraphExecDestroy(infinirtGraphExec_t graph_exe __INFINI_C infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) { INFINIRT_CALL_DEVICE_API(graphLuanch, (graph_exec, stream)); } + +__INFINI_C infiniStatus_t infinirtGraphTaskGroupBegin(infinirtStream_t stream) { +#ifdef ENABLE_ASCEND_API + if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { + return infinirt::ascend::graphTaskGroupBegin(stream); + } +#endif + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +__INFINI_C infiniStatus_t infinirtGraphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { +#ifdef ENABLE_ASCEND_API + if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { + return infinirt::ascend::graphTaskGroupEnd(stream, handle); + } +#endif + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +__INFINI_C infiniStatus_t infinirtGraphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { +#ifdef ENABLE_ASCEND_API + if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { + return infinirt::ascend::graphTaskUpdateBegin(stream, handle); + } +#endif + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +__INFINI_C infiniStatus_t infinirtGraphTaskUpdateEnd(infinirtStream_t stream) { +#ifdef ENABLE_ASCEND_API + if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { + return infinirt::ascend::graphTaskUpdateEnd(stream); + } +#endif + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} From 63f59862519115bdf0e81b0b35416e9d37f4d230 Mon Sep 17 00:00:00 2001 From: Kingbo_998 Date: Tue, 4 Aug 2026 17:17:55 +0800 Subject: [PATCH 2/4] fix conflict --- include/infinicore/graph/graph.hpp | 5 +++++ src/infinicore/graph/graph.cc | 24 ++++++++++++++++++++++-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index e79de5e3d..d8d310347 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -19,6 +19,9 @@ class GraphTensor : public Tensor { class GraphOperator { public: virtual void run() const = 0; + virtual bool is_device_graph_capture_safe() const { + return true; + } virtual bool requires_task_update() const { return false; } @@ -68,6 +71,8 @@ class Graph { private: struct DeviceGraph; + struct Segment; + std::vector> segments_; std::unique_ptr device_graph_; std::unordered_map> host_int_arrays_; }; diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 5f5231ee4..0092c50a7 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -121,6 +121,25 @@ struct Graph::DeviceGraph { } }; +struct Graph::Segment { + bool capture_safe; + std::vector> ops; + std::unique_ptr device_graph; + + explicit Segment(bool capture_safe_) : capture_safe(capture_safe_) { + } + + void run() const { + if (device_graph) { + device_graph->launch(); + return; + } + for (const auto &op : ops) { + op->run(); + } + } +}; + Graph::Graph() : host_int_arrays_(std::move(staged_host_int_arrays)) { staged_host_int_arrays.clear(); @@ -135,8 +154,9 @@ void Graph::run() const { } device_graph_.get()->launch(); } else { - for (auto &op : op_list_) { - op->run(); + if (segments_.empty()) { + for (const auto &op : op_list_) { + op->run(); } } } From 1b7fb6e67fa266842c497a55ff087496d8896202 Mon Sep 17 00:00:00 2001 From: Kingbo_998 Date: Tue, 4 Aug 2026 20:30:09 +0800 Subject: [PATCH 3/4] fix conflict of graph --- include/infinicore/graph/graph.hpp | 1 - src/infinicore/graph/graph.cc | 133 +++++++++++++++++++---------- 2 files changed, 88 insertions(+), 46 deletions(-) diff --git a/include/infinicore/graph/graph.hpp b/include/infinicore/graph/graph.hpp index d8d310347..57447633b 100644 --- a/include/infinicore/graph/graph.hpp +++ b/include/infinicore/graph/graph.hpp @@ -73,7 +73,6 @@ class Graph { struct DeviceGraph; struct Segment; std::vector> segments_; - std::unique_ptr device_graph_; std::unordered_map> host_int_arrays_; }; diff --git a/src/infinicore/graph/graph.cc b/src/infinicore/graph/graph.cc index 0092c50a7..2bf0c7936 100644 --- a/src/infinicore/graph/graph.cc +++ b/src/infinicore/graph/graph.cc @@ -2,6 +2,7 @@ #include "../utils.hpp" #include "infinicore/context/context.hpp" +#include #include namespace infinicore::graph { @@ -146,18 +147,25 @@ Graph::Graph() } void Graph::run() const { - if (device_graph_ != nullptr && device_graph_.get()->exec != nullptr) { + if (segments_.empty()) { + for (const auto &op : op_list_) { + op->run(); + } + return; + } + + for (const auto &segment : segments_) { + if (!segment->device_graph) { + segment->run(); + continue; + } + HostIntArrayScope update_scope(&host_int_arrays_, true); - for (const auto &task : device_graph_->updatable_tasks) { + for (const auto &task : segment->device_graph->updatable_tasks) { TaskUpdateHandleScope task_update_scope(task.handle); task.op->run(); } - device_graph_.get()->launch(); - } else { - if (segments_.empty()) { - for (const auto &op : op_list_) { - op->run(); - } + segment->device_graph->launch(); } } @@ -240,61 +248,91 @@ const std::vector *lookup_task_update_host_int_array( } void Graph::instantiate() { - // Reset device graph - device_graph_ = std::make_unique(); + segments_.clear(); - // warmup + // Warm the complete op list before splitting it into replay segments. for (size_t iter = 0; iter < 5; ++iter) { this->run(); } infinicore::context::syncStream(); - if (infinirtStreamBeginCapture( - context::getStream(), - INFINIRT_STREAM_CAPTURE_MODE_RELAXED) - != INFINI_STATUS_SUCCESS) { + // Diagnostic escape hatch: keep GraphTensor/operator replay semantics but + // bypass device-graph capture, including segmented PP capture. + if (std::getenv("INFINICORE_DISABLE_DEVICE_GRAPH_SEGMENTS") != nullptr) { + spdlog::info("device graph segments disabled; replaying recorded operators"); return; } - // Run and record. Operators with dynamic host-side arguments are captured - // as individual ModelRI task groups so those arguments can be updated at - // replay time. - device_graph_->updatable_tasks.clear(); - HostIntArrayScope capture_scope(&host_int_arrays_, false); for (const auto &op : op_list_) { - if (!op->requires_task_update()) { - op->run(); + const bool capture_safe = op->is_device_graph_capture_safe(); + if (segments_.empty() || segments_.back()->capture_safe != capture_safe) { + segments_.push_back(std::make_unique(capture_safe)); + } + segments_.back()->ops.push_back(op); + } + + for (auto &segment : segments_) { + if (!segment->capture_safe) { + // Replay non-capturable operators once between captured segments so + // later capture observes the same stream-ordered dependencies. + segment->run(); continue; } - infinirtGraphTaskGroup_t handle = nullptr; - { - TaskGroupCaptureScope task_group_scope(&handle); - op->run(); + segment->device_graph = std::make_unique(); + auto &device_graph = *segment->device_graph; + if (infinirtStreamBeginCapture( + context::getStream(), + INFINIRT_STREAM_CAPTURE_MODE_RELAXED) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error("failed to begin device graph capture"); } - INFINICORE_ASSERT(handle != nullptr); - device_graph_->updatable_tasks.push_back({op, handle}); - } - if (infinirtStreamEndCapture( - context::getStream(), - &device_graph_.get()->graph) - != INFINI_STATUS_SUCCESS) { - return; + device_graph.updatable_tasks.clear(); + HostIntArrayScope capture_scope(&host_int_arrays_, false); + for (const auto &op : segment->ops) { + if (!op->requires_task_update()) { + op->run(); + continue; + } + + infinirtGraphTaskGroup_t handle = nullptr; + { + TaskGroupCaptureScope task_group_scope(&handle); + op->run(); + } + INFINICORE_ASSERT(handle != nullptr); + device_graph.updatable_tasks.push_back({op, handle}); + } + + if (infinirtStreamEndCapture( + context::getStream(), + &device_graph.graph) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error("failed to end device graph capture"); + } + + if (infinirtGraphInstantiate( + &device_graph.exec, + device_graph.graph, + &device_graph.node, + device_graph.log_buffer.data(), + device_graph.log_buffer.size()) + != INFINI_STATUS_SUCCESS) { + throw std::runtime_error( + "failed to instantiate device graph: " + + std::string(device_graph.log_buffer.data())); + } } - if (infinirtGraphInstantiate( - &device_graph_.get()->exec, - device_graph_.get()->graph, - &device_graph_.get()->node, - device_graph_.get()->log_buffer.data(), - device_graph_.get()->log_buffer.size()) - != INFINI_STATUS_SUCCESS) { - static bool warned_once = false; - if (!warned_once) { - warned_once = true; - spdlog::warn("Fail to instantiate device graph: {}", std::string(device_graph_.get()->log_buffer.data())); + if (std::getenv("INFINICORE_GRAPH_DEBUG") != nullptr) { + size_t host_segments = 0; + for (const auto &segment : segments_) { + host_segments += segment->capture_safe ? 0 : 1; } + spdlog::info( + "segmented graph: operators={}, segments={}, host_segments={}", + op_list_.size(), segments_.size(), host_segments); } } @@ -334,4 +372,9 @@ std::shared_ptr GraphManager::stop_recording() { return std::exchange(graph_, nullptr); } +void GraphManager::cancel_recording() { + recording_ = false; + graph_.reset(); +} + } // namespace infinicore::graph From 131d99dfa59409242d059877713fbec3346157c9 Mon Sep 17 00:00:00 2001 From: Kingbo_998 Date: Thu, 6 Aug 2026 11:00:48 +0800 Subject: [PATCH 4/4] adapt to common operation of runtime --- src/infinirt/ascend/infinirt_ascend.h | 8 -------- src/infinirt/bang/infinirt_bang.cc | 20 ++++++++++++++++++ src/infinirt/cpu/infinirt_cpu.cc | 20 ++++++++++++++++++ src/infinirt/cuda/infinirt_cuda.cu | 20 ++++++++++++++++++ src/infinirt/infinirt.cc | 28 ++++---------------------- src/infinirt/infinirt_impl.h | 6 +++++- src/infinirt/kunlun/infinirt_kunlun.cc | 20 ++++++++++++++++++ src/infinirt/metax/infinirt_metax.cc | 20 ++++++++++++++++++ src/infinirt/moore/infinirt_moore.cc | 20 ++++++++++++++++++ 9 files changed, 129 insertions(+), 33 deletions(-) diff --git a/src/infinirt/ascend/infinirt_ascend.h b/src/infinirt/ascend/infinirt_ascend.h index 7b4d5976f..d7b2d7513 100644 --- a/src/infinirt/ascend/infinirt_ascend.h +++ b/src/infinirt/ascend/infinirt_ascend.h @@ -6,14 +6,6 @@ namespace infinirt::ascend { #ifdef ENABLE_ASCEND_API infiniStatus_t init(); INFINIRT_DEVICE_API_IMPL -infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream); -infiniStatus_t graphTaskGroupEnd( - infinirtStream_t stream, - infinirtGraphTaskGroup_t *handle); -infiniStatus_t graphTaskUpdateBegin( - infinirtStream_t stream, - infinirtGraphTaskGroup_t handle); -infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream); #else INFINIRT_DEVICE_API_NOOP #endif diff --git a/src/infinirt/bang/infinirt_bang.cc b/src/infinirt/bang/infinirt_bang.cc index ff40e58cc..14eb91e16 100644 --- a/src/infinirt/bang/infinirt_bang.cc +++ b/src/infinirt/bang/infinirt_bang.cc @@ -182,6 +182,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/cpu/infinirt_cpu.cc b/src/infinirt/cpu/infinirt_cpu.cc index 9036e484f..3a03be06e 100644 --- a/src/infinirt/cpu/infinirt_cpu.cc +++ b/src/infinirt/cpu/infinirt_cpu.cc @@ -196,4 +196,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + } // namespace infinirt::cpu diff --git a/src/infinirt/cuda/infinirt_cuda.cu b/src/infinirt/cuda/infinirt_cuda.cu index f9fc14c42..d33269e7a 100644 --- a/src/infinirt/cuda/infinirt_cuda.cu +++ b/src/infinirt/cuda/infinirt_cuda.cu @@ -567,4 +567,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre CHECK_CUDART(cudaGraphLaunch((cudaGraphExec_t)graph_exec, (cudaStream_t)stream)); return INFINI_STATUS_SUCCESS; } + +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} } diff --git a/src/infinirt/infinirt.cc b/src/infinirt/infinirt.cc index cf485d2b6..102757ff2 100644 --- a/src/infinirt/infinirt.cc +++ b/src/infinirt/infinirt.cc @@ -251,41 +251,21 @@ __INFINI_C infiniStatus_t infinirtGraphLuanch(infinirtGraphExec_t graph_exec, in } __INFINI_C infiniStatus_t infinirtGraphTaskGroupBegin(infinirtStream_t stream) { -#ifdef ENABLE_ASCEND_API - if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { - return infinirt::ascend::graphTaskGroupBegin(stream); - } -#endif - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + INFINIRT_CALL_DEVICE_API(graphTaskGroupBegin, (stream)); } __INFINI_C infiniStatus_t infinirtGraphTaskGroupEnd( infinirtStream_t stream, infinirtGraphTaskGroup_t *handle) { -#ifdef ENABLE_ASCEND_API - if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { - return infinirt::ascend::graphTaskGroupEnd(stream, handle); - } -#endif - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + INFINIRT_CALL_DEVICE_API(graphTaskGroupEnd, (stream, handle)); } __INFINI_C infiniStatus_t infinirtGraphTaskUpdateBegin( infinirtStream_t stream, infinirtGraphTaskGroup_t handle) { -#ifdef ENABLE_ASCEND_API - if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { - return infinirt::ascend::graphTaskUpdateBegin(stream, handle); - } -#endif - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + INFINIRT_CALL_DEVICE_API(graphTaskUpdateBegin, (stream, handle)); } __INFINI_C infiniStatus_t infinirtGraphTaskUpdateEnd(infinirtStream_t stream) { -#ifdef ENABLE_ASCEND_API - if (CURRENT_DEVICE_TYPE == INFINI_DEVICE_ASCEND) { - return infinirt::ascend::graphTaskUpdateEnd(stream); - } -#endif - return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; + INFINIRT_CALL_DEVICE_API(graphTaskUpdateEnd, (stream)); } diff --git a/src/infinirt/infinirt_impl.h b/src/infinirt/infinirt_impl.h index 6c26e0b02..ff17801a9 100644 --- a/src/infinirt/infinirt_impl.h +++ b/src/infinirt/infinirt_impl.h @@ -47,7 +47,11 @@ char *log_buffer, \ size_t buffer_size) IMPL; \ INLINE infiniStatus_t graphExecDestroy(infinirtGraphExec_t graph_exec) IMPL; \ - INLINE infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) IMPL; + INLINE infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stream) IMPL; \ + INLINE infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) IMPL; \ + INLINE infiniStatus_t graphTaskGroupEnd(infinirtStream_t stream, infinirtGraphTaskGroup_t *handle) IMPL; \ + INLINE infiniStatus_t graphTaskUpdateBegin(infinirtStream_t stream, infinirtGraphTaskGroup_t handle) IMPL; \ + INLINE infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) IMPL; #define INFINIRT_DEVICE_API_IMPL INFINIRT_DEVICE_API(, , ) #define INFINIRT_DEVICE_API_NOOP INFINIRT_DEVICE_API( \ diff --git a/src/infinirt/kunlun/infinirt_kunlun.cc b/src/infinirt/kunlun/infinirt_kunlun.cc index 02c46bb33..37a2d655a 100644 --- a/src/infinirt/kunlun/infinirt_kunlun.cc +++ b/src/infinirt/kunlun/infinirt_kunlun.cc @@ -190,6 +190,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } diff --git a/src/infinirt/metax/infinirt_metax.cc b/src/infinirt/metax/infinirt_metax.cc index d9ea00b3f..878240b4d 100644 --- a/src/infinirt/metax/infinirt_metax.cc +++ b/src/infinirt/metax/infinirt_metax.cc @@ -269,4 +269,24 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_SUCCESS; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + } // namespace infinirt::metax diff --git a/src/infinirt/moore/infinirt_moore.cc b/src/infinirt/moore/infinirt_moore.cc index 301cbe77c..54ceee218 100644 --- a/src/infinirt/moore/infinirt_moore.cc +++ b/src/infinirt/moore/infinirt_moore.cc @@ -178,6 +178,26 @@ infiniStatus_t graphLuanch(infinirtGraphExec_t graph_exec, infinirtStream_t stre return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; } +infiniStatus_t graphTaskGroupBegin(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskGroupEnd( + infinirtStream_t stream, + infinirtGraphTaskGroup_t *handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateBegin( + infinirtStream_t stream, + infinirtGraphTaskGroup_t handle) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + +infiniStatus_t graphTaskUpdateEnd(infinirtStream_t stream) { + return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; +} + infiniStatus_t getMemInfo(int device_id, size_t *free_bytes, size_t *total_bytes) { return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED; }