[fix] fix workspace mismatch when using fa2

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>();
-    uint64_t seed = gen_obj.at("seed").to_number<uint64_t>();
-    uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number<uint64_t>();
-
-    auto state = global_generator_state_registry.get_state_from_id(state_id, seed);
-    state->register_graph(reinterpret_cast<at::cuda::CUDAGraph*>(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>();
+      uint64_t seed = gen_obj.at("seed").to_number<uint64_t>();
+      uint64_t wholegraph_increment = gen_obj.at("wholegraph_increment").to_number<uint64_t>();
+
+      auto state = global_generator_state_registry.get_state_from_id(state_id, seed);
+      state->register_graph(reinterpret_cast<at::cuda::CUDAGraph*>(graph.get()));
+      graph->captured_generator_states_[state] = wholegraph_increment;
+    }
   }
 
   const json::object& allocator_events = root.at("allocator_events").as_object();

csrc/CUDAGraphParallel.cpp

@@ -2047,7 +2047,20 @@ GraphLoadResult finish_one_graph_load_impl(std::shared_ptr<PendingGraphLoads> 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();

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");

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();