fix

a-r-r-o-w · a-r-r-o-w · commit e76fc948b0df · 2025-07-16T21:41:09.000+02:00
diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py
@@ -899,6 +899,7 @@ def _templated_context_parallel_attention(
 @_AttentionBackendRegistry.register(
     AttentionBackendName.FLASH,
     constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+    supports_context_parallel=True,
 )
 def _flash_attention(
     query: torch.Tensor,

Original file line number	Diff line number	Diff line change
`@@ -899,6 +899,7 @@ def _templated_context_parallel_attention(`
`899`	`899`	`@_AttentionBackendRegistry.register(`
`900`	`900`	`AttentionBackendName.FLASH,`
`901`	`901`	`constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
	`902`	`+ supports_context_parallel=True,`
`902`	`903`	`)`
`903`	`904`	`def _flash_attention(`
`904`	`905`	`query: torch.Tensor,`