diff --git a/csrc/CUDAGraph.cpp b/csrc/CUDAGraph.cpp index 25df352..dcc1a01 100644 --- a/csrc/CUDAGraph.cpp +++ b/csrc/CUDAGraph.cpp @@ -1527,16 +1527,24 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool) TORCH_INTERNAL_ASSERT(graph->mempool_id_.first > 0); } - const json::array& generators_array = root.at("generators").as_array(); - for (const auto& gen_val : generators_array) { - const json::object& gen_obj = gen_val.as_object(); - uint64_t state_id = gen_obj.at("id").to_number(); - uint64_t seed = gen_obj.at("seed").to_number(); - uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); - - auto state = global_generator_state_registry.get_state_from_id(state_id, seed); - state->register_graph(reinterpret_cast(graph.get())); - graph->captured_generator_states_[state] = wholegraph_increment; + { + // Suspend the tracked region: the first register_graph lazily + // allocates seed/offset_extragraph_ (see finish_one_graph_load_impl + // in CUDAGraphParallel.cpp for the full rationale); in-region pool + // growth here would push the VMM cursor past the recorded + // start_base_addr and abort the replay below. + foundry::SuspendAllocationRegion suspend_region; + const json::array& generators_array = root.at("generators").as_array(); + for (const auto& gen_val : generators_array) { + const json::object& gen_obj = gen_val.as_object(); + uint64_t state_id = gen_obj.at("id").to_number(); + uint64_t seed = gen_obj.at("seed").to_number(); + uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number(); + + auto state = global_generator_state_registry.get_state_from_id(state_id, seed); + state->register_graph(reinterpret_cast(graph.get())); + graph->captured_generator_states_[state] = wholegraph_increment; + } } const json::object& allocator_events = root.at("allocator_events").as_object(); diff --git a/csrc/CUDAGraphParallel.cpp b/csrc/CUDAGraphParallel.cpp index 6ebfc0f..1d6f2f7 100644 --- a/csrc/CUDAGraphParallel.cpp +++ b/csrc/CUDAGraphParallel.cpp @@ -2047,7 +2047,20 @@ GraphLoadResult finish_one_graph_load_impl(std::shared_ptr pe // Register deferred generators (before allocator replay to match // SAVE mode timing where generators are created before graph capture). + // + // The first registration lazily creates the generator state's + // seed/offset_extragraph_ tensors (two at::empty in + // CUDAGeneratorState::register_graph). Whether that at::empty grows a + // fresh caching-allocator segment depends on the free-pool state, which + // is NOT mirrored between SAVE and LOAD (SAVE's lazy init fires inside + // capture_begin, before the start_base_addr snapshot). If the growth + // happened inside the tracked region here, the VMM cursor would move + // past the recorded start_base_addr and replay_hook_events_from_json + // below would abort ("Memory offset mismatch during replay", observed + // as a 2 MB drift on A100). Suspend the region so any segment growth + // lands outside the tracked cursor. if (pending->registry && !entry.generators_meta.is_null()) { + foundry::SuspendAllocationRegion suspend_region; const boost::json::array& gen_array = entry.generators_meta.as_array(); for (const auto& gen_val : gen_array) { const boost::json::object& gen_obj = gen_val.as_object(); diff --git a/csrc/hook.cpp b/csrc/hook.cpp index d552d08..64007e5 100644 --- a/csrc/hook.cpp +++ b/csrc/hook.cpp @@ -3158,6 +3158,10 @@ void resume_allocation_region() { #endif } +bool allocation_region_enabled() { + return tls_storage.enabled; +} + bool preallocate_region(size_t size) { if (!tls_storage.region_initialized) { fprintf(stderr, "[HOOK] ERROR: Cannot preallocate before allocation region is set\n"); diff --git a/include/hook.h b/include/hook.h index 4d6e132..2ae2b8c 100644 --- a/include/hook.h +++ b/include/hook.h @@ -27,6 +27,30 @@ namespace foundry { void set_allocation_region(void* base, size_t size); void stop_allocation_region(); void resume_allocation_region(); +bool allocation_region_enabled(); + +// RAII: temporarily route allocations to the ordinary CUDA allocator +// instead of the tracked VMM region (no cursor movement). Restores the +// prior enabled state, so it nests safely and is a no-op when no region +// is active on the calling thread. +class SuspendAllocationRegion { + public: + SuspendAllocationRegion() : was_enabled_(allocation_region_enabled()) { + if (was_enabled_) { + stop_allocation_region(); + } + } + ~SuspendAllocationRegion() { + if (was_enabled_) { + resume_allocation_region(); + } + } + SuspendAllocationRegion(const SuspendAllocationRegion&) = delete; + SuspendAllocationRegion& operator=(const SuspendAllocationRegion&) = delete; + + private: + bool was_enabled_; +}; bool preallocate_region(size_t size); void free_preallocated_region(); size_t get_current_alloc_offset();