From 6b86d992df39a5256cd6d8c942874871500a67c2 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 20 Oct 2026 18:56:18 +0000 Subject: [PATCH 1/3] fix compilation CI. --- .../test_models_autoencoder_mmaudio.py | 11 +++++- .../autoencoders/test_models_vocoder.py | 25 ++++++++++++- .../test_models_latent_upscaler.py | 14 +++++++- tests/models/testing_utils/quantization.py | 9 +++++ .../test_models_transformer_kandinsky6_sr.py | 35 +++++++++++++++++-- 5 files changed, 89 insertions(+), 5 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py index a12e01ee4049..810fdb6f31d0 100644 --- a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -13,9 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import MMAudioVAE +from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device @@ -93,4 +95,11 @@ class TestMMAudioVAEMemory(MMAudioVAETesterConfig, MemoryTesterMixin): class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin): - pass + @pytest.mark.xfail( + condition=is_torch_version("<", "2.12"), + reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", + raises=torch._dynamo.exc.Unsupported, + strict=True, + ) + def test_torch_compile_recompilation_and_graph_break(self): + super().test_torch_compile_recompilation_and_graph_break() diff --git a/tests/models/autoencoders/test_models_vocoder.py b/tests/models/autoencoders/test_models_vocoder.py index 225cfdbf4386..779d62c6be20 100644 --- a/tests/models/autoencoders/test_models_vocoder.py +++ b/tests/models/autoencoders/test_models_vocoder.py @@ -13,9 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import MMAudioVocoder +from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device @@ -83,4 +85,25 @@ class TestMMAudioVocoderMemory(MMAudioVocoderTesterConfig, MemoryTesterMixin): class TestMMAudioVocoderTorchCompile(MMAudioVocoderTesterConfig, TorchCompileTesterMixin): - pass + @pytest.mark.xfail( + condition=is_torch_version("<", "2.12"), + reason=( + 'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, ' + "which Dynamo cannot trace in a full graph." + ), + raises=torch._dynamo.exc.Unsupported, + strict=True, + ) + def test_torch_compile_recompilation_and_graph_break(self): + super().test_torch_compile_recompilation_and_graph_break() + + @pytest.mark.xfail( + reason=( + "Block-level offloading hooks each `ups[i]` `ModuleList`, but `forward` calls `ups[i][0]` directly, so the " + "hook never onloads the weights." + ), + raises=RuntimeError, + strict=True, + ) + def test_compile_with_group_offloading(self): + super().test_compile_with_group_offloading() diff --git a/tests/models/latent_upscaler/test_models_latent_upscaler.py b/tests/models/latent_upscaler/test_models_latent_upscaler.py index 059041af075d..a384c66da30c 100644 --- a/tests/models/latent_upscaler/test_models_latent_upscaler.py +++ b/tests/models/latent_upscaler/test_models_latent_upscaler.py @@ -13,9 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import Kandinsky6SRLatentUpscalerBank +from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device @@ -93,4 +95,14 @@ class TestKandinsky6SRLatentUpscalerBankMemory(Kandinsky6SRLatentUpscalerBankTes class TestKandinsky6SRLatentUpscalerBankTorchCompile( Kandinsky6SRLatentUpscalerBankTesterConfig, TorchCompileTesterMixin ): - pass + @pytest.mark.xfail( + condition=is_torch_version("<", "2.12"), + reason=( + 'With deterministic algorithms on, torch < 2.12 routes `F.pad(mode="replicate")` through `importlib`, ' + "which Dynamo cannot trace in a full graph." + ), + raises=torch._dynamo.exc.Unsupported, + strict=True, + ) + def test_torch_compile_recompilation_and_graph_break(self): + super().test_torch_compile_recompilation_and_graph_break() diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 66b9b66c3204..c0d429f52d77 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -2093,8 +2093,16 @@ class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin): Use `pytest -m "not bfl"` to skip these tests """ + @staticmethod + def _xfail_if_no_fp8_compile(): + # Both schemes store E4M3 tensors (fp8r weights, nvfp4 block scales) and compute in bf16; Triton cannot + # compile either below compute capability 8.9 (e.g. T4). + if torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9): + pytest.xfail("torch.compile of BFL models needs FP8 support in Triton (compute capability 8.9+)") + @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_torch_compile(self, config_name): + self._xfail_if_no_fp8_compile() if ( config_name == "nvfp4" and torch_device == "cuda" @@ -2109,4 +2117,5 @@ def test_bfl_torch_compile(self, config_name): @pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"]) def test_bfl_torch_compile_with_group_offload(self, config_name): + self._xfail_if_no_fp8_compile() self._test_torch_compile_with_group_offload(BFLConfigMixin.BFL_CONFIGS[config_name]) diff --git a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py index 27d2a56a4a7a..ca0609c8a4d3 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py +++ b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py @@ -17,6 +17,7 @@ import torch from diffusers import Kandinsky6SRTransformer3DModel +from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device @@ -60,10 +61,11 @@ def get_init_dict(self) -> dict: "out_visual_dim": 8, "time_dim": 16, "patch_size": (1, 2, 2), - "model_dim": 24, + "model_dim": 32, "ff_dim": 32, "num_visual_blocks": 2, - "axes_dims": (4, 4, 4), + # head_dim = sum(axes_dims); Inductor lowers flex attention only for head dims of at least 16 + "axes_dims": (4, 6, 6), } def _build_dummy_inputs(self, batch_size: int, num_frames: int, height: int, width: int) -> dict: @@ -106,6 +108,35 @@ def different_shapes_for_compilation(self): def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict: return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width) + @pytest.mark.xfail( + condition=is_torch_version("<", "2.12"), + reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", + raises=torch._dynamo.exc.Unsupported, + strict=True, + ) + def test_torch_compile_recompilation_and_graph_break(self): + super().test_torch_compile_recompilation_and_graph_break() + + @pytest.mark.xfail( + condition=is_torch_version("<", "2.12"), + reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", + raises=torch._dynamo.exc.Unsupported, + strict=True, + ) + def test_torch_compile_repeated_blocks(self): + super().test_torch_compile_repeated_blocks() + + @pytest.mark.xfail( + reason=( + "The token grid is split into 8x8 blocks and `sliding_tile_mask` sizes its table by `max(...)` of the block " + "counts, so the compiled graph specializes on the input shape and recompiles." + ), + raises=(torch._dynamo.exc.Unsupported, torch._dynamo.exc.RecompileError), + strict=True, + ) + def test_compile_on_different_shapes(self): + super().test_compile_on_different_shapes() + @pytest.mark.skipif(torch_device == "cpu", reason="FlexAttention does not support backward on CPU.") class TestKandinsky6SRTransformerTraining(Kandinsky6SRTransformerTesterConfig, TrainingTesterMixin): From 053e1a62f9239f88e51765adefd71f12efb8cbbc Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 22 Oct 2026 04:53:01 +0000 Subject: [PATCH 2/3] fix attentionbackend name dispatch. --- src/diffusers/models/attention_dispatch.py | 3 ++- .../test_models_autoencoder_mmaudio.py | 10 +-------- .../test_models_transformer_kandinsky6_sr.py | 21 +------------------ 3 files changed, 4 insertions(+), 30 deletions(-) diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py index 18e575d8ebbc..3d2ae36bd9a2 100644 --- a/src/diffusers/models/attention_dispatch.py +++ b/src/diffusers/models/attention_dispatch.py @@ -418,7 +418,8 @@ def dispatch_attention_fn( # variable), or we use a custom backend based on whether user is using the `attention_backend` context manager backend_name, backend_fn = _AttentionBackendRegistry.get_active_backend() else: - backend_name = AttentionBackendName(backend) + # Skip the `Enum` call for members: older torch versions cannot always trace it in a full graph. + backend_name = backend if isinstance(backend, AttentionBackendName) else AttentionBackendName(backend) backend_fn = _AttentionBackendRegistry._backends.get(backend_name) kwargs = { diff --git a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py index 4c6ef22da267..45f47322a312 100644 --- a/tests/models/autoencoders/test_models_autoencoder_mmaudio.py +++ b/tests/models/autoencoders/test_models_autoencoder_mmaudio.py @@ -17,7 +17,6 @@ import torch from diffusers import MMAudioVAE -from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, require_accelerator, torch_device @@ -164,11 +163,4 @@ def test_disk_offload_with_safetensors(self, base_model_output, tmp_path): class TestMMAudioVAETorchCompile(MMAudioVAETesterConfig, TorchCompileTesterMixin): - @pytest.mark.xfail( - condition=is_torch_version("<", "2.12"), - reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", - raises=torch._dynamo.exc.Unsupported, - strict=True, - ) - def test_torch_compile_recompilation_and_graph_break(self): - super().test_torch_compile_recompilation_and_graph_break() + pass diff --git a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py index ca0609c8a4d3..1ee07ce2d38d 100644 --- a/tests/models/transformers/test_models_transformer_kandinsky6_sr.py +++ b/tests/models/transformers/test_models_transformer_kandinsky6_sr.py @@ -17,7 +17,6 @@ import torch from diffusers import Kandinsky6SRTransformer3DModel -from diffusers.utils import is_torch_version from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import enable_full_determinism, torch_device @@ -108,30 +107,12 @@ def different_shapes_for_compilation(self): def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict: return self._build_dummy_inputs(batch_size=1, num_frames=2, height=height, width=width) - @pytest.mark.xfail( - condition=is_torch_version("<", "2.12"), - reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", - raises=torch._dynamo.exc.Unsupported, - strict=True, - ) - def test_torch_compile_recompilation_and_graph_break(self): - super().test_torch_compile_recompilation_and_graph_break() - - @pytest.mark.xfail( - condition=is_torch_version("<", "2.12"), - reason="torch < 2.12 cannot trace `AttentionBackendName(backend)` in `dispatch_attention_fn` in a full graph.", - raises=torch._dynamo.exc.Unsupported, - strict=True, - ) - def test_torch_compile_repeated_blocks(self): - super().test_torch_compile_repeated_blocks() - @pytest.mark.xfail( reason=( "The token grid is split into 8x8 blocks and `sliding_tile_mask` sizes its table by `max(...)` of the block " "counts, so the compiled graph specializes on the input shape and recompiles." ), - raises=(torch._dynamo.exc.Unsupported, torch._dynamo.exc.RecompileError), + raises=torch._dynamo.exc.RecompileError, strict=True, ) def test_compile_on_different_shapes(self): From 0cf057604a82e6a91847b58f2705ba6ec68fdab5 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Thu, 22 Oct 2026 04:58:01 +0000 Subject: [PATCH 3/3] skip with a marker. --- tests/models/testing_utils/quantization.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index c0d429f52d77..dc4681f23ecc 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -2074,6 +2074,12 @@ def test_bfl_rejects_unquantized_checkpoint(self): @is_torch_compile @require_accelerator @require_accelerate +# Both schemes store E4M3 tensors (fp8r weights, nvfp4 block scales) and compute in bf16; Triton cannot compile either +# below compute capability 8.9 (e.g. T4). +@pytest.mark.skipif( + torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9), + reason="torch.compile of BFL models needs FP8 support in Triton (compute capability 8.9+)", +) class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin): """ Mixin class for testing torch.compile with BFL (`fp8r` / `nvfp4`) quantized models. @@ -2093,16 +2099,8 @@ class BFLCompileTesterMixin(BFLConfigMixin, QuantizationCompileTesterMixin): Use `pytest -m "not bfl"` to skip these tests """ - @staticmethod - def _xfail_if_no_fp8_compile(): - # Both schemes store E4M3 tensors (fp8r weights, nvfp4 block scales) and compute in bf16; Triton cannot - # compile either below compute capability 8.9 (e.g. T4). - if torch_device == "cuda" and torch.cuda.get_device_capability() < (8, 9): - pytest.xfail("torch.compile of BFL models needs FP8 support in Triton (compute capability 8.9+)") - @pytest.mark.parametrize("config_name", list(BFLConfigMixin.BFL_CONFIGS), ids=list(BFLConfigMixin.BFL_CONFIGS)) def test_bfl_torch_compile(self, config_name): - self._xfail_if_no_fp8_compile() if ( config_name == "nvfp4" and torch_device == "cuda" @@ -2117,5 +2115,4 @@ def test_bfl_torch_compile(self, config_name): @pytest.mark.parametrize("config_name", ["fp8r"], ids=["fp8r"]) def test_bfl_torch_compile_with_group_offload(self, config_name): - self._xfail_if_no_fp8_compile() self._test_torch_compile_with_group_offload(BFLConfigMixin.BFL_CONFIGS[config_name])