[Neuron] Add AWS Neuron (Trainium/Inferentia) as an officially supported device

[JingyaHuang]

What does this PR do?

This PR adds AWS Neuron (Trainium/Inferentia) as an officially supported compute backend in Diffusers, on par with existing backends like CUDA, MPS, XPU, and MLU.

Changes

• import_utils.py — adds is_torch_neuronx_available() detection, following the existing pattern for optional backends.
• torch_utils.py — registers "neuron" in all backend dispatch tables (BACKEND_SUPPORTS_TRAINING, BACKEND_EMPTY_CACHE, BACKEND_DEVICE_COUNT, BACKEND_MANUAL_SEED, etc.) and adds a randn_tensor workaround since Neuron/XLA does not support creating random tensors directly on device (falls back to CPU).
• utils/init.py — exports is_torch_neuronx_available.
• pipeline_utils.py — adds two new DiffusionPipeline methods:
  • enable_neuron_compile(model_names, cache_dir, fullgraph) — wraps pipeline nn.Module components with torch.compile(backend="neuron") for whole-graph NEFF compilation. Supports optional NEFF caching via TORCH_NEURONX_NEFF_CACHE_DIR.
  • neuron_warmup(*args, **kwargs) — runs a single dummy forward pass to trigger upfront neuronx-cc compilation before timed inference.

Usage

• Eager mode

import torch                                                                                                             
import torch_neuronx  # noqa: F401 — registers torch.neuron                                                            
                                                                                                                           
from diffusers import AutoPipelineForText2Image                                                                          
                                                                                                                           
# Load and move to Neuron device                                                                                         
pipe = AutoPipelineForText2Image.from_pretrained(                                                                        
    "stabilityai/sdxl-turbo",                                                                                            
    torch_dtype=torch.bfloat16,                           
    variant="fp16",                                                                                                      
)
pipe = pipe.to(torch.neuron.current_device())                                                                            
                                                                                                                         
# Warmup                                                                   
pipe(prompt="warmup", height=512, width=512, num_inference_steps=1, guidance_scale=0.0)                                                                                                                        
                                                          
# Inference                                                                                               
image = pipe(                                             
    prompt="a golden retriever surfing a wave, photorealistic",                                                          
    height=512,
    width=512,                                                                                                           
    num_inference_steps=1, 
    guidance_scale=0.0,                                                                    
).images[0]                                                                                                              
                                                                                                                         
image.save("output.png") 

Validation

So far we validated the following models, the idea is to manually validated a part of representative models, and ensure the model coverage with pur CIs afterward.

• pixart
• sdxl
• flux2-klein-4B (able to run within one NeuronCore w/o. tp under eager mode for resolution 1024x1024)

Next Steps

• Enable torch.compile on Neuron device
• Add tensor parallel support for memory-bound devices like neuron
• Tackle the compatibility of diffusers+nki kernels lib to boost the performance on neuron.
• Support sequence parallel on Neuron for video models.

[HuggingFaceDocBuilderDev]

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]

Nice thank you!

[sayakpaul]

@claude-2-serge could you do a review please?

[github-actions]

🤗 Serge says:

This PR adds AWS Neuron (Trainium/Inferentia) as a supported backend. The core plumbing (import detection, backend dispatch tables, randn_tensor workaround, dtype fallbacks) follows existing patterns. However, there are several issues ranging from dead code in tests to an inconsistency between the pixart and SDXL neuron checks, and a missing BACKEND_SUPPORTS_TRAINING entry in the test dispatch tables.

Correctness

• Dead atol in integration tests — Both Flux2KleinPipelineIntegrationTests and StableDiffusionXLTurboPipelineIntegrationTests compute atol but never use it in any assertion. The _ = atol line is a no-op. This looks like an incomplete implementation — either the tolerance should be used in an assertLessEqual against an expected slice, or the lines should be removed.

• Inconsistent neuron detection in pixart vs SDXL — The pixart pipeline uses is_torch_neuronx_available() (a package-availability check) to decide timestep_device, while the SDXL pipeline correctly uses device.type == "neuron" (a runtime device check). If torch_neuronx is installed but the pipeline runs on CUDA, the pixart pipeline would unnecessarily force timesteps to CPU. The SDXL approach is more correct.

• BACKEND_SUPPORTS_TRAINING not patched for neuron in tests/testing_utils.py — The _neuron_device key is added to all other dispatch tables but not to BACKEND_SUPPORTS_TRAINING. This means backend_supports_training(torch_device) falls back to "default": True, contradicting torch_utils.py where "neuron": False.

PR Description Mismatch

• The description claims enable_neuron_compile and neuron_warmup methods are added to pipeline_utils.py, but neither exists in the diff or the codebase. Either these were dropped from the PR or the description is stale.

Tests — Environment Variable Leakage

• setUp in Flux2KleinPipelineIntegrationTests sets TORCH_NEURONX_NEFF_CACHE_DIR and TORCH_NEURONX_ENABLE_NKI_SDPA but tearDown never cleans them up. Minor for slow tests but still a leak.

Suggestions / Dead Code Analysis

• The image_slice variable is computed in both new integration tests but only used for a std() check — no expected-slice comparison is performed. Combined with the dead atol, these tests only verify shape and range, not numerical correctness. This is weak for integration tests.

35 LLM turns · 42 tool calls · 201.4s · 926751 in / 8552 out tokens

[JingyaHuang]

Thanks @sayakpaul, I addressed the review comments, shall I merge?

[yiyixuxu]

thanks! i left some questions/comments

[yiyixuxu]

@claude-2-serge can you trace the code for how device is defined here and let me know if there is any code path that _exeuction_devide would not return torch.device type?

[sayakpaul]

Follow-up inline comments are something we're working with @tarekziade. Should land pretty soon.

[JingyaHuang]

addressed the comments! @sayakpaul @yiyixuxu

[JingyaHuang]

The failing CI checks are unrelated to this PR, ok for us to merge? @sayakpaul @yiyixuxu