diffusers/tests/models/test_attention_dispatch.py

# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
import os

import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn.functional as F

from diffusers.models._modeling_parallel import ContextParallelConfig, ParallelConfig
from diffusers.models.attention_dispatch import attention_backend as attention_backend_ctx
from diffusers.models.attention_dispatch import dispatch_attention_fn

from ..testing_utils import (
    is_attention,
    is_context_parallel,
    is_kernels_available,
    is_torch_compile,
    require_torch_accelerator,
    require_torch_multi_accelerator,
    torch_device,
)
from .testing_utils.parallelism import DEVICE_CONFIG, _find_free_port


# Max allowed relative error between the context parallel gradients and the single-process reference.
GRAD_RTOL = 2e-2


def _attention_backward_parity_worker(rank, world_size, master_port, cp_dict, attention_backend, return_dict):
    """Op-level worker: check `dispatch_attention_fn` gradients against a single-process reference.

    This guards the ring-attention backward pass, which historically produced silently wrong
    gradients (see https://github.com/huggingface/diffusers/issues/14265) because it recomputed
    every ring iteration against the iteration-0 KV chunk.
    """
    try:
        os.environ["MASTER_ADDR"] = "localhost"
        os.environ["MASTER_PORT"] = str(master_port)
        os.environ["RANK"] = str(rank)
        os.environ["WORLD_SIZE"] = str(world_size)

        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
        device_type = torch_device.split(":")[0]
        device_module = device_config["module"]

        dist.init_process_group(backend=device_config["backend"], rank=rank, world_size=world_size)
        device_module.set_device(rank)
        device = torch.device(f"{device_type}:{rank}")

        # Identical inputs on every rank so each rank can compute the same full-sequence reference.
        B, S, H, D = 1, 128, 4, 64
        torch.manual_seed(777)
        q = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16)
        k = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16)
        v = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16)
        grad_out = torch.randn(B, S, H, D, device=device, dtype=torch.bfloat16)

        # Reference: full-sequence fp32 SDPA gradients on a single process.
        q_ref, k_ref, v_ref = (t.float().requires_grad_(True) for t in (q, k, v))
        ref_out = F.scaled_dot_product_attention(
            q_ref.transpose(1, 2), k_ref.transpose(1, 2), v_ref.transpose(1, 2)
        ).transpose(1, 2)
        ref_dq, ref_dk, ref_dv = torch.autograd.grad(ref_out, (q_ref, k_ref, v_ref), grad_out.float())

        mesh = dist.device_mesh.init_device_mesh(
            device_type,
            (cp_dict.get("ring_degree", 1), cp_dict.get("ulysses_degree", 1)),
            mesh_dim_names=("ring", "ulysses"),
        )
        cp_config = ContextParallelConfig(**cp_dict)
        cp_config.setup(rank, world_size, device, mesh)
        parallel_config = ParallelConfig(context_parallel_config=cp_config)

        # Each rank runs its sequence shard through the templated CP attention path.
        shard = slice(rank * S // world_size, (rank + 1) * S // world_size)
        qs, ks, vs = (t.detach()[:, shard].clone().requires_grad_(True) for t in (q, k, v))
        with attention_backend_ctx(attention_backend):
            out = dispatch_attention_fn(qs, ks, vs, parallel_config=parallel_config)
        out.backward(grad_out[:, shard])

        rel_errs = {}
        for name, got, ref in (("dq", qs.grad, ref_dq), ("dk", ks.grad, ref_dk), ("dv", vs.grad, ref_dv)):
            ref_shard = ref[:, shard].to(got.dtype)
            rel_errs[name] = ((got - ref_shard).norm() / ref_shard.norm()).item()

        if rank == 0:
            return_dict["status"] = "success"
            return_dict["rel_errs"] = dict(rel_errs)

    except Exception as e:
        if rank == 0:
            return_dict["status"] = "error"
            return_dict["error"] = repr(e)
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()


@is_attention
@is_context_parallel
@require_torch_multi_accelerator
class TestContextParallelAttentionBackward:
    """Op-level tests for the context parallel backward of `dispatch_attention_fn`, model independent."""

    @pytest.mark.parametrize("cp_type", ["ulysses_degree", "ring_degree"])
    @pytest.mark.parametrize(
        "attention_backend",
        [
            "_native_flash",
            pytest.param(
                "_native_cudnn",
                marks=pytest.mark.skipif(torch_device != "cuda", reason="CUDA is required for _native_cudnn backend."),
            ),
            pytest.param(
                "flash_hub",
                marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
            ),
            pytest.param(
                "flash_varlen_hub",
                marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
            ),
            pytest.param(
                "_flash_3_hub",
                marks=[
                    pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
                    pytest.mark.skipif(torch_device != "cuda", reason="CUDA is required for _flash_3_hub backend."),
                ],
            ),
            pytest.param(
                "_flash_3_varlen_hub",
                marks=[
                    pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
                    pytest.mark.skipif(
                        torch_device != "cuda", reason="CUDA is required for _flash_3_varlen_hub backend."
                    ),
                ],
            ),
        ],
    )
    def test_attn_backend_backward_parity(self, cp_type, attention_backend):
        """Ring and Ulysses attention gradients must match a single-GPU reference.

        Regression test for https://github.com/huggingface/diffusers/issues/14265, where the ring
        backward silently used the iteration-0 KV chunk for every ring iteration.
        """
        if not torch.distributed.is_available():
            pytest.skip("torch.distributed is not available.")

        world_size = 2
        cp_dict = {cp_type: world_size}
        master_port = _find_free_port()

        manager = mp.Manager()
        return_dict = manager.dict()
        mp.spawn(
            _attention_backward_parity_worker,
            args=(world_size, master_port, cp_dict, attention_backend, return_dict),
            nprocs=world_size,
            join=True,
        )

        assert return_dict.get("status") == "success", (
            f"Context parallel backward parity run failed: {return_dict.get('error', 'Unknown error')}"
        )
        rel_errs = return_dict["rel_errs"]
        for name, rel in rel_errs.items():
            assert rel < GRAD_RTOL, (
                f"{attention_backend} {cp_type} gradient `{name}` rel_err={rel:.3e} "
                f"exceeds tol={GRAD_RTOL:.1e}: {rel_errs}"
            )


@is_attention
class TestVarlenAttentionCompile:
    """Op-level tests for the varlen backends under `torch.compile` with dynamic shapes, model independent."""

    @pytest.mark.parametrize(
        "attention_backend",
        [
            pytest.param(
                "flash_varlen_hub",
                marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
            ),
            pytest.param(
                "_flash_3_varlen_hub",
                marks=pytest.mark.skipif(not is_kernels_available(), reason="`kernels` is not available."),
            ),
        ],
    )
    @is_torch_compile
    @require_torch_accelerator
    def test_dynamic_shapes(self, attention_backend):
        def make_qkv(seq_len):
            torch.manual_seed(0)
            return tuple(torch.randn(2, seq_len, 4, 64, device=torch_device, dtype=torch.bfloat16) for _ in range(3))

        with contextlib.ExitStack() as stack, torch.no_grad():
            try:
                stack.enter_context(attention_backend_ctx(attention_backend))
                dispatch_attention_fn(*make_qkv(128))
            except Exception as e:
                pytest.skip(f"Skipping test for backend '{attention_backend}': {e}")

            torch.compiler.reset()
            compiled_attention = torch.compile(dispatch_attention_fn, dynamic=True, fullgraph=True)
            try:
                with (
                    torch._inductor.utils.fresh_inductor_cache(),
                    torch._dynamo.config.patch(error_on_recompile=True),
                ):
                    for seq_len in (128, 256, 256):
                        q, k, v = make_qkv(seq_len)
                        torch.testing.assert_close(
                            compiled_attention(q, k, v), dispatch_attention_fn(q, k, v), atol=1e-3, rtol=1e-3
                        )
            finally:
                torch.compiler.reset()