[Neuron] Enable `torch.compile` compatibility with Neuron device by JingyaHuang · Pull Request #13485 · huggingface/diffusers

JingyaHuang · 2026-04-15T16:57:37Z

What does this PR do?

Following #13289

Adding torch.compile support for Neuron to diffusers.

Neuron supports torch.compile with the following mechanism to improve performance:

torch.compile(model, backend="neuron", fullgraph=True)                                                                                                 
           │                                                
           ▼                                                                                                                                             
     TorchDynamo                   traces Python → FX graph (aten ops)
           │                                                                                                                                             
           ▼                                                
     AOT Autograd                  lifts params as inputs, decomposes ops                                                                                
           │                                                                                                                                             
           ▼
     Neuron backend                Neuron-specific decompositions                                                                                        
     (torch_neuronx)               (int64→int32, SDPA lowering, NKI passes)
           │  StableHLO                                                                                                                                  
           ▼                                                                                                                                             
     neuronx-cc                    compiles HLO → NEFF binary                                                                                            
                                   allocates on-chip SRAM (SB, 28 MB/core)                                                                               
           │  NEFF                                                                                                                                       
           ▼                                                                                                                                             
     NEFF cache                    keyed by StableHLO hash                                                                                               
     (TORCH_NEURONX_NEFF_CACHE_DIR)  miss → recompile; hit → load instantly                                                                              
           │                                                                                                                                             
           ▼                                                                                                                                             
     Executor                      preprocesses inputs (contiguous, int64→int32)                                                                         
     (executor.py)                 calls _C.execute_compiled_graph()

This PR integrates torch.compile on Neuron seamlessly into Diffusers via the following changes:

In src/diffusers/models/transformers/transformer_flux2.py
- Flux2PosEmbed.forward() — add Neuron to the float32 RoPE guard (Neuron doesn't support float64, same pattern as MPS/NPU)
- Flux2SwiGLU.forward() — replace x.chunk(2, dim=-1) with explicit x[..., :half] / x[..., half:] slices (equivalent on all backends; avoids a gather op in the Neuron IR that causes NCC_ITEN404 DataLocalityOpt compilation failure under torch.compile)
In src/diffusers/models/unets/unet_2d_condition.py
- get_time_embed() — use int32 timestep dtype on Neuron (Neuron doesn't support int64)
In src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py
- Keep timesteps on CPU for Neuron (same pattern as XLA);
- cast to float32 before moving to device;
- run scheduler step on CPU output
In src/diffusers/utils/import_utils.py + torch_utils.py
- Register the neuron backend: BACKEND_SYNCHRONIZE, BACKEND_SUPPORTS_TRAINING, BACKEND_DEVICE_COUNT, randn_tensor forces CPU random generation then .to(device) (Neuron has no native RNG)
a test for pixart

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)

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.

Caveats

Somes fixes need to be upsteamed to TorchNeuron, this PR won't be stable with potential changes in the SDK

… into add-neuron-backend

HuggingFaceDocBuilderDev · 2026-04-15T17:06:58Z

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.

…ers into neuron-torch-comppile

sayakpaul · 2026-06-05T02:51:12Z

        # We set the index here to remove DtoH sync, helpful especially during compilation.
        # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
-        self.scheduler.set_begin_index(0)
+        if hasattr(self.scheduler, "set_begin_index"):


This seems like a bigger change that should be propagated to other pipelines that use self.scheduler.set_begin_index(0)?

yes, only a part of self.scheduler have set_begin_index, so it will be good to have the check, I only added for the pipelines related to the PR, do you want me to propagate to others? Or I could open a separate PR for this.

sayakpaul · 2026-06-05T02:52:15Z

+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(0)


If the pipeline initially didn't have it, maybe we don't need it?

The change was suggested in the initial PR for supporting neuron: #13289 (comment).

the line is needed for the compiled path otherwise the first scheduler.step calls index_for_timestep, which does a comparison/nonzero and forces a DtoH sync inside the compiled denoising loop (same motivation as #11696).

sayakpaul

Left some further comments.

…ers into neuron-torch-comppile

JingyaHuang

Thanks @sayakpaul for the review! I addressed the comments you made.

JingyaHuang · 2026-06-15T13:31:06Z

-        x1, x2 = x.chunk(2, dim=-1)
-        x = self.gate_fn(x1) * x2
+        half = x.shape[-1] // 2
+        x = self.gate_fn(x[..., :half]) * x[..., half:]


yeah, here is totally a fix for neuron compilation. neuron compiler lower chunk to split/gather and previously introduced a compilation failure, now it's fixed by neuron's new beta drop. But the explicit static slices x[..., :half] / x[..., half:] lower to clean slice ops, and the neuron compiler compiles it faster.

JingyaHuang · 2026-06-15T14:22:20Z

+        if hasattr(self.scheduler, "set_begin_index"):
+            self.scheduler.set_begin_index(0)


The change was suggested in the initial PR for supporting neuron: #13289 (comment).

the line is needed for the compiled path otherwise the first scheduler.step calls index_for_timestep, which does a comparison/nonzero and forces a DtoH sync inside the compiled denoising loop (same motivation as #11696).

JingyaHuang · 2026-06-17T07:22:51Z

        # We set the index here to remove DtoH sync, helpful especially during compilation.
        # Check out more details here: https://github.com/huggingface/diffusers/pull/11696
-        self.scheduler.set_begin_index(0)
+        if hasattr(self.scheduler, "set_begin_index"):


yes, only a part of self.scheduler have set_begin_index, so it will be good to have the check, I only added for the pipelines related to the PR, do you want me to propagate to others? Or I could open a separate PR for this.

JingyaHuang · 2026-06-17T07:54:02Z

+    def test_flux2_klein_neuron_compile_128(self):
+        import torch_neuronx  # noqa: F401 — registers torch.neuron
+        from torch_neuronx.neuron_dynamo_backend import set_model_name
+        from transformers.utils.output_capturing import install_all_output_capturing_hooks


It was added for the compilation of the text encoder, which is not a component we compile in general, removed it!

yiyixuxu · 2026-06-18T03:28:01Z

thanks a ton for working on this @JingyaHuang !

) * draft:add neuron as a legit backend * feat: neuron-specific changes in the pipeline * tests: eager tests * fix: style * fix:apr_02 beta * fix:pixart * fix: rewrite flux swiglu activation to avoid gather op in neuron IR * test: pixart compile mode on neuron * cleanup & fix style * test: flux2 klein and sdxl-turbo * review: address some comments * revert:put back expected slices * cleanup:move the ts index ahead of denoising + cleanup tests * fix: style * fix style * review:remove the hook --------- Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

JingyaHuang and others added 16 commits March 18, 2026 11:15

draft:add neuron as a legit backend

98f6c8c

Merge branch 'huggingface:main' into add-neuron-backend

c58b8b8

Merge branch 'huggingface:main' into add-neuron-backend

3367409

Merge branch 'main' into add-neuron-backend

0c51734

feat: neuron-specific changes in the pipeline

a76953c

tests: eager tests

2480388

fix: style

929ab72

Merge branch 'huggingface:main' into add-neuron-backend

52cac76

Merge branch 'add-neuron-backend' of github.com:JingyaHuang/diffusers…

28a5086

… into add-neuron-backend

Merge branch 'huggingface:main' into add-neuron-backend

68689e5

Merge branch 'main' into add-neuron-backend

da79308

fix:apr_02 beta

3bb9c7c

Merge branch 'add-neuron-backend' of github.com:JingyaHuang/diffusers…

c4facab

… into add-neuron-backend

Merge branch 'huggingface:main' into add-neuron-backend

1eb5ff9

fix:pixart

cbe8f28

fix: rewrite flux swiglu activation to avoid gather op in neuron IR

16b9606

github-actions Bot added lora models tests utils pipelines examples size/L PR with diff > 200 LOC labels Apr 15, 2026

test: pixart compile mode on neuron

7f13f68

github-actions Bot added size/L PR with diff > 200 LOC and removed size/L PR with diff > 200 LOC labels Apr 15, 2026

JingyaHuang changed the title ~~[Neuron] Enable torch.comppile compatibility with Neuron device~~ [Neuron] Enable torch.compile compatibility with Neuron device Apr 21, 2026

Merge branch 'main' into neuron-torch-comppile

a46cb19

github-actions Bot removed the size/L PR with diff > 200 LOC label Apr 22, 2026

github-actions Bot added size/M PR with diff < 200 LOC and removed size/L PR with diff > 200 LOC labels Jun 3, 2026

JingyaHuang and others added 2 commits June 3, 2026 13:29

Merge branch 'main' into neuron-torch-comppile

b586e14

review: address some comments

cf20ff1

github-actions Bot added size/L PR with diff > 200 LOC and removed size/M PR with diff < 200 LOC labels Jun 3, 2026

revert:put back expected slices

d8d0db3

github-actions Bot added size/M PR with diff < 200 LOC and removed size/L PR with diff > 200 LOC labels Jun 3, 2026

JingyaHuang and others added 4 commits June 4, 2026 15:28

cleanup:move the ts index ahead of denoising + cleanup tests

f536c59

Merge branch 'main' into neuron-torch-comppile

7df50b8

fix: style

7f6a730

Merge branch 'neuron-torch-comppile' of github.com:JingyaHuang/diffus…

f07d960

…ers into neuron-torch-comppile

JingyaHuang requested a review from sayakpaul June 4, 2026 16:52

sayakpaul reviewed Jun 5, 2026

View reviewed changes

JingyaHuang and others added 5 commits June 10, 2026 22:18

fix style

90b94bc

Merge branch 'main' into neuron-torch-comppile

1e5ab32

Merge branch 'main' into neuron-torch-comppile

956af41

review:remove the hook

7f97666

Merge branch 'neuron-torch-comppile' of github.com:JingyaHuang/diffus…

9d8f244

…ers into neuron-torch-comppile

JingyaHuang commented Jun 17, 2026

View reviewed changes

Merge branch 'main' into neuron-torch-comppile

f72d3f8

sayakpaul approved these changes Jun 17, 2026

View reviewed changes

yiyixuxu merged commit 236e5dd into huggingface:main Jun 18, 2026
21 of 22 checks passed

JingyaHuang deleted the neuron-torch-comppile branch June 18, 2026 17:25

Beinsezii mentioned this pull request Jul 1, 2026

Regression (#13485) Broken TorchAO Compat #14097

Open

		if hasattr(self.scheduler, "set_begin_index"):
		self.scheduler.set_begin_index(0)

Uh oh!

Conversation

JingyaHuang commented Apr 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Validation

Who can review?

Caveats

Uh oh!

HuggingFaceDocBuilderDev commented Apr 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

JingyaHuang left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yiyixuxu commented Jun 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

JingyaHuang commented Apr 15, 2026 •

edited

Loading