Skip to content

Consolidate torch device backend dispatch - #14792

Open
DN6 wants to merge 5 commits into
mainfrom
device_dtype_routing
Open

Consolidate torch device backend dispatch#14792
DN6 wants to merge 5 commits into
mainfrom
device_dtype_routing

Conversation

@DN6

@DN6 DN6 commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Support for torch device backends keeps growing, and each new backend adds a branch at every site that touches the device module. e.g. Adding NPU streams to group offloading (#14785) looks like this:

if torch.cuda.is_available():
    stream = torch.cuda.Stream()
elif hasattr(torch, "xpu") and torch.xpu.is_available():
    stream = torch.Stream()
elif hasattr(torch, "npu") and torch.npu.is_available():
    stream = torch.npu.Stream()

This PR introduces a TorchDeviceBackend object that automatically routes to the appropriate backend module for a given device type. e.g.

backend = TorchDeviceBackend(device)   # torch.cuda for "cuda", torch.xpu for "xpu", torch.npu for "npu", ...

# These methods will now work regardless of the device. We can avoid the need for multiple if/else checks
backend.Stream()                       
backend.empty_cache()                  
backend.max_memory_allocated()

Fixes # (issue)

Before submitting

  • Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
    • Did you read the Coding with AI agents guide?
    • Did you run the self-review skill on the diff?
    • Did you share the final self-review notes in the PR description or a comment?
  • Did you read the contributor guideline?
  • Did you read our philosophy doc? (important for complex PRs)
  • Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?
  • Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@DN6 DN6 added this to the Release 0.41.0 milestone Sep 16, 2026
@github-actions github-actions Bot added the size/L PR with diff > 200 LOC label Sep 16, 2026
DN6
DN6 commented Sep 16, 2026
View reviewed changes
Comment thread src/diffusers/pipelines/pag/pipeline_pag_sana.py
image = latents
else:
latents = latents.to(self.vae.dtype)
torch_accelerator_module = getattr(torch, get_device(), torch.cuda)

@DN6 DN6 Sep 16, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Deadcode since min supported torch version is 2.6

@DN6
DN6 requested review from sayakpaul and yiyixuxu September 16, 2026 12:35
@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Sep 16, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@sayakpaul sayakpaul moved this to In Progress in Diffusers Roadmap Sep 17, 2026
sayakpaul
sayakpaul reviewed Sep 17, 2026
View reviewed changes

@sayakpaul sayakpaul left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clean

Comment thread src/diffusers/hooks/group_offloading.py Outdated
raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.")
if use_stream and onload_device.type == "cpu":
raise ValueError("Using streams for data transfer requires an accelerator onload device, got `cpu`.")
stream = TorchDeviceBackend(onload_device).Stream() if use_stream else None

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if the underlying device doesn't support streams?

@DN6 DN6 Sep 28, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed.

Comment thread src/diffusers/pipelines/sana/pipeline_sana_controlnet.py
Comment on lines -1056 to -1061
torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
oom_error = (
torch.OutOfMemoryError
if is_torch_version(">=", "2.5.0")
else torch_accelerator_module.OutOfMemoryError
)

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Safe because we pin on >=2.6.

Comment thread src/diffusers/utils/torch_utils.py
Comment on lines +224 to +228
def empty_cache(self) -> None:
# Backends without a caching allocator (cpu, neuron) have nothing to clear.
empty_cache = getattr(self.module, "empty_cache", None)
if empty_cache is not None:
empty_cache()

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we also warn in case empty_cache is None?

Comment thread tests/others/test_utils.py Outdated
"""Tests for :class:`diffusers.utils.torch_utils.TorchDeviceBackend` on the CPU backend."""

def test_module_resolves_from_str_and_torch_device(self):
import torch

@sayakpaul sayakpaul Sep 17, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can keep all the imports at the top. They seem harmless.

@DN6
DN6 requested a review from sayakpaul September 28, 2026 04:01
sayakpaul
sayakpaul approved these changes Sep 28, 2026
View reviewed changes
Comment thread src/diffusers/quantizers/gguf/gguf_quantizer.py
else torch.cuda.current_device()
)
model.to(device)
model.to(get_device())

@sayakpaul sayakpaul Sep 28, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are all the sites of .to(device) being changed to .to(get_device() in this PR?

@DN6 DN6 Sep 30, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not everywhere. Just the places where we were inlining the device fallback. Can do a follow up pass to find the other places where this is applicable.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@sayakpaul sayakpaul sayakpaul approved these changes
@yiyixuxu yiyixuxu Awaiting requested review from yiyixuxu

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation hooks lora modular-pipelines pipelines quantization size/L PR with diff > 200 LOC tests utils

Projects

Diffusers Roadmap
Status: In Progress

Milestone

Release 0.41.0

Development

Successfully merging this pull request may close these issues.

3 participants

@DN6 @HuggingFaceDocBuilderDev @sayakpaul