Enable `torch.Generator` to support pytorch/xla generator implementation #161369

iwknow · 2025-08-24T05:10:41Z

Currently, the implementation of torch.Generator only support "cpu" and "cuda" device type. https://github.com/pytorch/pytorch/blob/main/torch/csrc/Generator.cpp#L55-L61

This change enables torch.Generator to support more device type by allowing any device backend to register their own generator factory through a Generator Registry. This is similar to what "DeviceGuardImpl registry" does today.

Key Changes:

New registry API:

Added GeneratorRegistry.h and GeneratorRegistry.cpp in c10/core/impl.
API supports registerGenerator(DeviceType, GeneratorFactory), unregisterGenerator(DeviceType), and getGeneratorFactory(DeviceType).
Uses c10::DeviceType as the key and stores a factory function returning c10::intrusive_ptrc10::GeneratorImpl.

Python/C++ integration:

The registry is consulted in the torch.Generator constructor path for non-CPU/CUDA devices.
If a factory is registered for the requested device, it constructs the appropriate generator; otherwise, raises an error.

Backend extensibility:

Out-of-tree backends (e.g., torch_xla, torch-directml, torch_npu) can now register their custom generator implementation at module load via a static registrar object.
Example usage:

C++
namespace {
  struct Registrar {
    Registrar() {
      at::detail::registerGenerator(c10::DeviceType::XLA, &CreateXlaGenerator);
    }
  } registrar_instance;
}

This allows torch.Generator(device='xla') to return an XlaGeneratorImpl when the torch_xla extension is imported.

…f-tree backends.

pytorch-bot · 2025-08-24T05:10:45Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/161369

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

ROCm MI2xx CI/CD workflows failing due to : download from https://api.github.com/repos/pytorch/pytorch timed out.

✅ No Failures

As of commit d838990 with merge base 76f81b5 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

linux-foundation-easycla · 2025-08-24T05:10:47Z

The committers listed above are authorized under a signed CLA.

✅ login: iwknow / name: Frank Liu (c57e4f1, ee5ab76, 87a5059, d838990, 32c2b34)

iwknow · 2025-08-24T22:57:54Z

related issue: pytorch/xla#9159

EikanWang

Regarding out-of-tree accelerators, I think PrivateUse1 based mechanism should be the recommended option - https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/detail/PrivateUse1HooksInterface.h#L76-L77

iwknow · 2025-08-25T21:22:42Z

hi @EikanWang Thanks for reviewing!
i don't think PrivateUse1 is viable for XLA use case if my understanding is correct.

for backends like XLA, it is a first‑class backend in PyTorch with its own DispatchKey::XLA and DeviceType::XLA. It is "off-tree" simply because it is developed and maintained in a different package. This is fundamentally different from the PrivateUse1 use case where the entire backend is registered under PrivateUse1, including use "DispatchKey::PrivateUse1".
Or, if you are suggesting to keep XLA and only make RNG kernels explicitly consume a PrivateUse1 generator. That wouldn't work neither. torch.device('xla') is already bound to DeviceType::XLA. You can’t rebind the name “xla” to PrivateUse1 in the same process/build. Registering a PrivateUse1 backend under the “xla” name would conflict with the existing device.

FFFrog · 2025-08-26T02:08:16Z

@EikanWang Thanks for the mention.

@iwknow PyTorch supports registering generators for out-of-tree backends by inheriting from the AcceleratorHooksInterface class. You can refer to the below link for more information.

Therefore, it seems to me that you might need to add new files named XLAHooksInterface.h and XLAHooksInterface.cpp to internally implement XLAHooksInterface in PyTorch Repo, and then create new XLAHooks.h and XLAHooks.cpp files in the XLA Repo that internally inherit from XLAHooksInterface implements the XLAHooks custom class and registers it in PyTorch.

FFFrog · 2025-08-26T02:15:53Z

torch/csrc/Generator.cpp

+  } else if (c10::impl::hasGenerator(device_type)) {
+    self->cdata = at::Generator(c10::impl::makeGenerator(device));
+  } else {
+    throw std::runtime_error("No generator available for device type: " +
+                             c10::toString(device_type));


So, wo do not need to modify this one, all the accelerators except CPU will follow the same AcceleratorHooksInterface logic.

iwknow · 2025-08-26T16:55:13Z

@EikanWang Thanks for the mention.

@iwknow PyTorch supports registering generators for out-of-tree backends by inheriting from the AcceleratorHooksInterface class. You can refer to the below link for more information.

https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/detail/XPUHooksInterface.h

https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/detail/XPUHooksInterface.cpp

https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/xpu/detail/XPUHooks.h

https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/xpu/detail/XPUHooks.cpp

Therefore, it seems to me that you might need to add new files named XLAHooksInterface.h and XLAHooksInterface.cpp to internally implement XLAHooksInterface in PyTorch Repo, and then create new XLAHooks.h and XLAHooks.cpp files in the XLA Repo that internally inherit from XLAHooksInterface implements the XLAHooks custom class and registers it in PyTorch.

This is actually a brilliant idea, i will take a closer look.

… from off-tree backends." This reverts commit 32c2b34.

iwknow · 2025-08-30T04:51:48Z

Hi @FFFrog
I've updated the PR based on your suggestion. please review. thank you very much!

FFFrog · 2025-08-30T07:33:50Z

aten/src/ATen/detail/XLAHooksInterface.h

+  virtual DeviceIndex getNumDevices() const {
+    return 0;
+  }


Why do we need this one? Can deviceCount from base class AcceleratorHooksInterface help?

good catch. i guess it is copy/paste from getNumGPUs XPUHooksInterface.h. removed to use DeviceIndex deviceCount from AcceleratorHooksInterface

FFFrog · 2025-08-30T07:36:30Z

It would be better if we change the code below into return detail::getXLAHooks().hasXLA(); as well.

pytorch/aten/src/ATen/Context.h

Lines 182 to 183 in f44ad54

    
           static bool hasXLA() { 
        
             return c10::impl::hasDeviceGuardImpl(c10::DeviceType::XLA);

iwknow · 2025-08-31T03:46:28Z

@FFFrog Thanks for reviewing. PR updated.

FFFrog

Thanks, LGTM, let's trigger CI first.

@iwknow, I don't have valid approval permissions, so we'll have to get approval from @albanD.

Hey, @albanD, please take a look at this, thanks.

Also, it's necessary to find a better way to implement Hooks registration, otherwise every time we add a new backend to PyTorch, we'll need to add a Hooks interface for it.

iwknow · 2025-09-03T19:55:26Z

@albanD, a kindly ping.

iwknow · 2025-09-09T16:36:12Z

Hi @FFFrog, can you please nominate another approver as it seems that albanD is unresponsive.

FFFrog · 2025-09-10T04:51:46Z

Hi @FFFrog, can you please nominate another approver as it seems that albanD is unresponsive.

Sure, Hey @albanD @ezyang @malfet, could you please help to take a look at this one? Thank you.

albanD

That sounds ok to add to make pytorch/XLA work. Do you have a sister PR on the xla side to make it work we should look at at the same time?
Also I would like one of the pytorch/XLA maintainer to comment here to make sure this is good on their end.

@FFFrog xla is a bit of a weird backend for historical reason. We don't plan on adding any new other backend that is not going through PrivateUse1 going forward indeed.

iwknow · 2025-09-11T20:44:12Z

That sounds ok to add to make pytorch/XLA work. Do you have a sister PR on the xla side to make it work we should look at at the same time? Also I would like one of the pytorch/XLA maintainer to comment here to make sure this is good on their end.

@FFFrog xla is a bit of a weird backend for historical reason. We don't plan on adding any new other backend that is not going through PrivateUse1 going forward indeed.

Sister PR for xla: yes, i do have the code change for XLA. it's basically a concrete implementation of the interface. but it will be in pytorch/XLA package. would you like to review?

qihqi has been reviewing the generator change. I think he can take a look at this change as well. qihqi, please leave a comment if this change looks good to you and is what torch/XLA team want.

FFFrog · 2025-09-12T00:55:27Z

@FFFrog xla is a bit of a weird backend for historical reason. We don't plan on adding any new other backend that is not going through PrivateUse1 going forward indeed.

Thank you, I got it.

iwknow · 2025-09-16T16:17:02Z

@zhanyong-wan can you please take a look or nominate someone from XLA side to see if this change is good from the XLA side. thanks!

iwknow · 2025-09-25T20:13:24Z

@qihqi ping again.

Enable torch.Generator to support generator implementations from of…

32c2b34

…f-tree backends.

pytorchbot added the open source label Aug 24, 2025

EikanWang reviewed Aug 25, 2025

View reviewed changes

EikanWang requested review from guangyey, FFFrog and albanD August 25, 2025 18:12

FFFrog reviewed Aug 26, 2025

View reviewed changes

soulitzer added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Aug 26, 2025

iwknow and others added 3 commits August 29, 2025 21:46

Merge branch 'pytorch:main' into main

ee5ab76

Revert "Enable torch.Generator to support generator implementations…

c57e4f1

… from off-tree backends." This reverts commit 32c2b34.

Add XLAHooksInterface.

87a5059

FFFrog reviewed Aug 30, 2025

View reviewed changes

address feedbacks.

d838990

FFFrog added the topic: not user facing topic category label Sep 1, 2025

FFFrog approved these changes Sep 1, 2025

View reviewed changes

albanD reviewed Sep 11, 2025

View reviewed changes

albanD changed the title ~~Enable torch.Generator to support off-tree generator implementation~~ Enable torch.Generator to support pytorch/xla generator implementation Sep 12, 2025

iwknow mentioned this pull request Oct 2, 2025

Add XLAHooksInterface to PyTorch for XLA pytorch/xla#9667

Open

Enable torch.Generator to support pytorch/xla generator implementation #161369

Are you sure you want to change the base?

Enable torch.Generator to support pytorch/xla generator implementation #161369

Uh oh!

Conversation

iwknow commented Aug 24, 2025

Key Changes:

New registry API:

Python/C++ integration:

Backend extensibility:

Uh oh!

pytorch-bot bot commented Aug 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/161369

❗ 1 Active SEVs

✅ No Failures

Uh oh!

linux-foundation-easycla bot commented Aug 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

iwknow commented Aug 24, 2025

Uh oh!

EikanWang left a comment

Choose a reason for hiding this comment

Uh oh!

iwknow commented Aug 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

FFFrog commented Aug 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

FFFrog Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

iwknow commented Aug 26, 2025

Uh oh!

iwknow commented Aug 30, 2025

Uh oh!

FFFrog Aug 30, 2025

Choose a reason for hiding this comment

Uh oh!

iwknow Aug 31, 2025

Choose a reason for hiding this comment

Uh oh!

FFFrog commented Aug 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

iwknow commented Aug 31, 2025

Uh oh!

FFFrog left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

iwknow commented Sep 3, 2025

Uh oh!

iwknow commented Sep 9, 2025

Uh oh!

FFFrog commented Sep 10, 2025

Uh oh!

albanD left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

iwknow commented Sep 11, 2025

Uh oh!

FFFrog commented Sep 12, 2025

Uh oh!

iwknow commented Sep 16, 2025

Uh oh!

iwknow commented Sep 25, 2025

Uh oh!

Uh oh!

Enable `torch.Generator` to support pytorch/xla generator implementation #161369

Enable `torch.Generator` to support pytorch/xla generator implementation #161369

pytorch-bot bot commented Aug 24, 2025 •

edited

Loading

linux-foundation-easycla bot commented Aug 24, 2025 •

edited

Loading

iwknow commented Aug 25, 2025 •

edited

Loading

FFFrog commented Aug 26, 2025 •

edited

Loading

FFFrog commented Aug 30, 2025 •

edited

Loading

FFFrog left a comment •

edited

Loading

albanD left a comment •

edited

Loading