Graph-Safe RNG State Exchange for Tensor Parallelism (pytorch#114068)

See pytorch#113541

The PR allows for registering and controlling multiple RNG states using indices, ensuring cudagraph-safe operations, and includes both C++ and Python API changes to support this functionality.

cc  @eellison @anijain2305 @jansel @ezyang @ptrblck @csarofeen @mcarilli
Pull Request resolved: pytorch#114068
Approved by: https://github.com/ezyang, https://github.com/eqy, https://github.com/xuzhao9

Author: eee4017

aten/src/ATen/core/Generator.cpp

@@ -13,4 +13,12 @@ at::Tensor Generator::get_state() const {
   return at::Tensor::wrap_tensor_impl(this->impl_->get_state());
 }
 
+void Generator::graphsafe_set_state(const Generator& new_state) {
+  this->impl_->graphsafe_set_state(new_state.getIntrusivePtr());
+}
+
+Generator Generator::graphsafe_get_state() const {
+  return Generator(this->impl_->graphsafe_get_state());
+}
+
 } // namespace at

aten/src/ATen/core/Generator.h

@@ -107,6 +107,10 @@ struct TORCH_API Generator {
 
   at::Tensor get_state() const;
 
+  void graphsafe_set_state(const Generator& new_state);
+
+  Generator graphsafe_get_state() const;
+
   std::mutex& mutex() {
     return impl_->mutex_;
   }