fix: handle non-zero storage_offset in sage attention for ring compatibility

Author: songh11

src/diffusers/models/attention_dispatch.py

@@ -1451,6 +1451,12 @@ def _sage_attention_forward_op(
     if enable_gqa:
         raise ValueError("`enable_gqa` is not yet supported for Sage attention.")
 
+    # In ring attention, key/value views from chunked AsyncCollectiveTensor may have invalid data_ptr.
+    if key.storage_offset() > 0:
+        key = key.clone()
+    if value.storage_offset() > 0:
+        value = value.clone()
+
     out = sageattn(
         q=query,
         k=key,
@@ -1489,6 +1495,12 @@ def _sage_attention_hub_forward_op(
     if enable_gqa:
         raise ValueError("`enable_gqa` is not yet supported for Sage attention.")
 
+    # In ring attention, key/value views from chunked AsyncCollectiveTensor may have invalid data_ptr.
+    if key.storage_offset() > 0:
+        key = key.clone()
+    if value.storage_offset() > 0:
+        value = value.clone()
+
     func = _HUB_KERNELS_REGISTRY[AttentionBackendName.SAGE_HUB].kernel_fn
     out = func(
         q=query,