Sage attention fails with ring-attention.

[songh11]

Describe the bug

When using Sage attention (SAGE or SAGE_HUB) combined with ring attention (context parallel),
the forward pass fails because sageattn CUDA kernels internally call tensor.data_ptr(),
which assumes storage_offset == 0.
In ring attention, funcol.all_gather_tensor(...).chunk(world_size) produces tensor views
with non-zero storage_offset. When these views are passed to sageattn, the kernel
reads from the wrong memory address.

Reproduction

import torch
import torch.distributed as dist
from diffusers import FluxPipeline
from diffusers.models._modeling_parallel import ContextParallelConfig
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--ring",    type=int, default=1)
    parser.add_argument("--ulysses", type=int, default=2)
    parser.add_argument("--sage", action="store_true")
    args = parser.parse_args()

    # 1. init dist — torchrun injects RANK / LOCAL_RANK / WORLD_SIZE automatically
    dist.init_process_group(backend="nccl")
    local_rank = dist.get_rank() % torch.cuda.device_count()
    torch.cuda.set_device(local_rank)

    # 2. load model
    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev/", torch_dtype=torch.bfloat16)
    pipe.to(f"cuda:{local_rank}") #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power

    # 3. enable sage
    if args.sage:
        pipe.transformer.set_attention_backend("sage")

    # 4. build parallel config via optkit and enable CP on the transformer
    parallelism_config = ContextParallelConfig(
        ring_degree=args.ring,
        ulysses_degree=args.ulysses,
    )
    pipe.transformer.enable_parallelism(config=parallelism_config)

    prompt = "A cat holding a sign that says hello world"
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]
    image.save("flux-dev.png")

    dist.destroy_process_group()

if __name__ == "__main__":
    main()

torchrun --nproc_per_node=2 ./demo.py --ring 2 --ulysses 1 --sage

Logs

[rank1]: Traceback (most recent call last):
[rank1]:   File "/root/workspace/optkit/examples/flux/./flux_context_parallel.py", line 37, in <module>
[rank1]:     image = pipe(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank1]:     return func(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/pipelines/flux/pipeline_flux.py", line 949, in __call__
[rank1]:     noise_pred = self.transformer(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/hooks/hooks.py", line 189, in new_forward
[rank1]:     output = function_reference.forward(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/utils/peft_utils.py", line 315, in wrapper
[rank1]:     result = forward_fn(self, *args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/transformers/transformer_flux.py", line 726, in forward
[rank1]:     encoder_hidden_states, hidden_states = block(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/transformers/transformer_flux.py", line 453, in forward
[rank1]:     attention_outputs = self.attn(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]:     return self._call_impl(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]:     return forward_call(*args, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/transformers/transformer_flux.py", line 352, in forward
[rank1]:     return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/transformers/transformer_flux.py", line 118, in __call__
[rank1]:     hidden_states = dispatch_attention_fn(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/attention_dispatch.py", line 432, in dispatch_attention_fn
[rank1]:     return backend_fn(**kwargs)
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/attention_dispatch.py", line 3262, in _sage_attention
[rank1]:     out = _templated_context_parallel_attention(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/attention_dispatch.py", line 2269, in _templated_context_parallel_attention
[rank1]:     return TemplatedRingAttention.apply(
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line 576, in apply
[rank1]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank1]:   File "/usr/local/lib/python3.10/dist-packages/diffusers/models/attention_dispatch.py", line 1935, in forward
[rank1]:     lse = prev_lse - torch.nn.functional.logsigmoid(prev_lse - lse)
[rank1]: torch.AcceleratorError: CUDA error: an illegal memory access was encountered

System Info

• 🤗 Diffusers version: 0.37.1
• Platform: Linux-5.10.134-15.al8.x86_64-x86_64-with-glibc2.35
• Running on Google Colab?: No
• Python version: 3.10.12
• PyTorch version (GPU?): 2.8.0+cu128 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 1.11.0
• Transformers version: 5.5.4
• Accelerate version: 1.13.0
• PEFT version: 0.19.1
• Bitsandbytes version: not installed
• Safetensors version: 0.7.0
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 4090, 49140 MiB
  NVIDIA GeForce RTX 4090, 49140 MiB
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

No response

[DN6]

@songh11 I'm unable to reproduce this using your script? Would you be able to run it in a fresh env with the sage_hub backend to see if the error persists?

[songh11]

@songh11 I'm unable to reproduce this using your script? Would you be able to run it in a fresh env with the sage_hub backend to see if the error persists?

Thanks for the quick response!

I’ve tested this in a fresh environment using both sage and sage_hub backends on a dual-RTX 5090 setup, and the error still persists. For the test, I was using the latest diffusers source code (installed from main).

I’m not sure about your specific test environment, but since my setup relies on PCIe for inter-GPU communication, it’s possible the issue is related to the PCIe communication overhead.

System Info

• 🤗 Diffusers version: 0.38.0.dev0
• Platform: Linux-5.10.134-15.al8.x86_64-x86_64-with-glibc2.35
• Running on Google Colab?: No
• Python version: 3.10.12
• PyTorch version (GPU?): 2.8.0+cu128 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Huggingface_hub version: 1.11.0
• Transformers version: 5.6.2
• Accelerate version: 1.13.0
• PEFT version: 0.19.1
• Bitsandbytes version: not installed
• Safetensors version: 0.8.0-rc.0
• xFormers version: not installed
• Accelerator: NVIDIA GeForce RTX 5090, 32607 MiB
  NVIDIA GeForce RTX 5090, 32607 MiB
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Error logs

Details

[rank0]: Traceback (most recent call last):
[rank0]:   File "/root/workspace/./demo.py", line 59, in <module>
[rank0]:     main()
[rank0]:   File "/root/workspace/./demo.py", line 45, in main
[rank0]:     image = pipe(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/pipelines/flux/pipeline_flux.py", line 946, in __call__
[rank0]:     noise_pred = self.transformer(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/hooks/hooks.py", line 189, in new_forward
[rank0]:     output = function_reference.forward(*args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/utils/peft_utils.py", line 315, in wrapper
[rank0]:     result = forward_fn(self, *args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 726, in forward
[rank0]:     encoder_hidden_states, hidden_states = block(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 453, in forward
[rank0]:     attention_outputs = self.attn(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 352, in forward
[rank0]:     return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/transformers/transformer_flux.py", line 118, in __call__
[rank0]:     hidden_states = dispatch_attention_fn(
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/attention_dispatch.py", line 440, in dispatch_attention_fn
[rank0]:     return backend_fn(**kwargs)
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/attention_dispatch.py", line 3333, in _sage_attention_hub
[rank0]:     out = _templated_context_parallel_attention(
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/attention_dispatch.py", line 2257, in _templated_context_parallel_attention
[rank0]:     return TemplatedRingAttention.apply(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line 576, in apply
[rank0]:     return super().apply(*args, **kwargs)  # type: ignore[misc]
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/attention_dispatch.py", line 1898, in forward
[rank0]:     out, lse = forward_op(
[rank0]:   File "/root/workspace/diffusers/src/diffusers/models/attention_dispatch.py", line 1493, in _sage_attention_hub_forward_op
[rank0]:     out = func(
[rank0]:   File "/root/.cache/huggingface/hub/models--kernels-community--sage-attention/snapshots/3b5365a429cb54d51e222d88e1488e949993f7a8/build/torch28-cxx11-cu128-x86_64-linux/core.py", line 155, in sageattn
[rank0]:     return sageattn_qk_int8_pv_fp8_cuda(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 929, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/root/.cache/huggingface/hub/models--kernels-community--sage-attention/snapshots/3b5365a429cb54d51e222d88e1488e949993f7a8/build/torch28-cxx11-cu128-x86_64-linux/core.py", line 710, in sageattn_qk_int8_pv_fp8_cuda
[rank0]:     lse / 1.44269504 + lse_correction * sm_scale
[rank0]: torch.AcceleratorError: CUDA error: no kernel image is available for execution on the device
[rank0]: CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
[rank0]: For debugging consider passing CUDA_LAUNCH_BLOCKING=1
[rank0]: Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
</details>