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.

✅ No Failures

As of commit 880d298 with merge base 8951df0 ():
💚 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 (e5598d4, ee5ab76, 32c2b34, 87a5059, 880d298, c57e4f1, d65a1b8, d838990, da1498b)

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-25T20:13:24Z

@qihqi ping again.

iwknow · 2025-10-14T17:39:37Z

@FFFrog @albanD can you please merge this change. thanks!

FFFrog · 2025-10-15T01:43:23Z

@FFFrog @albanD can you please merge this change. thanks!

Maintainer approval is required before merging (you can merge on your own, see this link for more information)

ysiraichi · 2025-10-15T13:45:30Z

@iwknow Could you open a PR in PyTorch/XLA with the XLAHooks implementation, pinning this PR? I think it's good measure to check that PyTorch/XLA tests pass, given that there is no PyTorch/XLA CI in PyTorch anymore.

iwknow · 2025-10-15T17:22:58Z

@iwknow Could you open a PR in PyTorch/XLA with the XLAHooks implementation, pinning this PR? I think it's good measure to check that PyTorch/XLA tests pass, given that there is no PyTorch/XLA CI in PyTorch anymore.

Will do shortly. meanwhile, @guangyey please review and approve.

guangyey · 2025-10-16T03:41:21Z

Hi @iwknow, I don't have permission to merge this PR. You need to get approval from @albanD

aten/src/ATen/Context.h

Also, rename one method in DeviceAccelerator to align with the interface.

iwknow · 2025-10-20T21:04:51Z

@ysiraichi FYI, this is the XLA hooks implementation: pytorch/xla#9683

albanD · 2025-10-20T21:32:04Z

Thanks for the linked PR, I'm not sure if something else is needed on the xla side, but happy to merge this one if it's enough.

albanD

Accepting here since @qihqi acccepted as well and it doesn't impact other components.

albanD · 2025-10-20T21:33:39Z

aten/src/ATen/DeviceAccelerator.cpp

  DETECT_AND_ASSIGN_ACCELERATOR_COMP(HIP)
  DETECT_AND_ASSIGN_ACCELERATOR_COMP(MPS)
  DETECT_AND_ASSIGN_ACCELERATOR_COMP(HPU)
+  DETECT_AND_ASSIGN_ACCELERATOR_COMP(XLA)


Adding xla here has quite deep implications FYI. It will change the behavior of many higher level systems wrt the xla device.

should we add then? it seems a right thing to do to my knowledge. if you feel we shouldn't make this change, i can revert

@albanD Could you give some examples on the behaviors that might change? Without this line, are we still going to see significant behavior changes w.r.t. PyTorch/XLA?

All the code that checks the current accelerator will now pick up xla when it wasn't before.
https://github.com/search?q=repo%3Apytorch%2Fpytorch+current_accelerator&type=code are the ones in core but there are many more out-of-core now.

Autograd will start to do stream handling for you, autocast will also have different behavior, pinned memory will start trying to use your host allocator, and out of core repos should use generic code instead of xla specific code.

Yes this concern is only about this one line (and the update below in this file). The rest of the changes are lower impact indeed.

after taking a closer look at getAcceleratorHooksInterface(

pytorch/aten/src/ATen/Context.h

Line 72 in 2fde10d

const AcceleratorHooksInterface& getAcceleratorHooksInterface(

), at::getAccelerator only makes difference when opt_device_type is doesn't have value. when i search the usage of getAcceleratorHooksInterface (https://github.com/search?q=repo%3Apytorch%2Fpytorch%20getAcceleratorHooksInterface&type=code), it turns out that it is only used in a few places and torch/csrc/autograd/engine.cpp is the only place that opt_device_type. The random number generator checks and passes opt_device_type. Therefore, it getAcceleratorHooksInterface works properly even at::getAccelerator doesn't recognize XLA.

Considering the scope of evaluating the impact of changing is at::getAccelerator is way beyond the scope of this change, i created a separate issue #166054 to track the change of at::getAccelerator and revert the related change from this PR.

iwknow · 2025-10-22T21:33:02Z

@albanD please merge if there is no other question. thanks!

FFFrog · 2025-10-23T11:34:08Z

@pytorchbot merge

pytorchmergebot · 2025-10-23T11:35:59Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

iwknow · 2025-10-23T17:36:27Z

@FFFrog i am confused by the merging status. is the merge failed? i see Closed with unmerged commits, but it seems that changes are committed to the mainline https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/detail/XLAHooksInterface.h

FFFrog · 2025-10-24T02:40:11Z

@FFFrog i am confused by the merging status. is the merge failed? i see Closed with unmerged commits, but it seems that changes are committed to the mainline https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/detail/XLAHooksInterface.h

Don't worry about the status of the PR, the changes have already been merged into the trunk. :D

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 FFFrog, albanD and guangyey 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

iwknow mentioned this pull request Oct 2, 2025

Add XLAHooksInterface to PyTorch for XLA pytorch/xla#9667

Open

qihqi approved these changes Oct 14, 2025

View reviewed changes

guangyey reviewed Oct 16, 2025

View reviewed changes

aten/src/ATen/Context.h Show resolved Hide resolved

iwknow and others added 3 commits October 19, 2025 15:00

Merge branch 'pytorch:main' into main

e5598d4

Update DeviceAccelerators to include XLA.

da1498b

Also, rename one method in DeviceAccelerator to align with the interface.

fix compile error

d65a1b8

iwknow mentioned this pull request Oct 20, 2025

implement XLAHooks and register it to PyTorch when loaded. pytorch/xla#9683

Open

albanD approved these changes Oct 20, 2025

View reviewed changes

revert changes in getAccelerator

880d298

iwknow mentioned this pull request Oct 22, 2025

Add XLA to at::accelerator::getAccelerator and at::accelerator::isAccelerator #166054

Open

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 23, 2025

pytorchmergebot added the merging label Oct 23, 2025

pytorchmergebot added the Merged label Oct 23, 2025

pytorchmergebot closed this in defb6a8 Oct 23, 2025

pytorchmergebot removed the merging label Oct 23, 2025

clee2000 mentioned this pull request Oct 24, 2025

Add XLAHooksInterface to bazel file #166179

Open

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

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

✅ 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!

Choose a reason for hiding this comment

Uh oh!

iwknow commented Aug 26, 2025

Uh oh!

iwknow commented Aug 30, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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 25, 2025

Uh oh!

iwknow commented Oct 14, 2025

Uh oh!

FFFrog commented Oct 15, 2025

Uh oh!

ysiraichi commented Oct 15, 2025

Uh oh!

iwknow commented Oct 15, 2025

Uh oh!

guangyey commented Oct 16, 2025

Uh oh!

Uh oh!

iwknow commented Oct 20, 2025

Uh oh!

albanD commented Oct 20, 2025

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