register _native_npu_attention to _supports_context_parallel

TmacAaron · TmacAaron · commit 002e7ef841b1 · 2026-01-08T08:29:29.000Z
Signed-off-by: yyt &lt;yangyit139@gmail.com&gt;
diff --git a/src/diffusers/models/attention_dispatch.py b/src/diffusers/models/attention_dispatch.py
@@ -2183,6 +2183,7 @@ def _native_math_attention(
 @_AttentionBackendRegistry.register(
     AttentionBackendName._NATIVE_NPU,
     constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
+    supports_context_parallel=True,
 )
 def _native_npu_attention(
     query: torch.Tensor,

Original file line number	Diff line number	Diff line change
`@@ -2183,6 +2183,7 @@ def _native_math_attention(`
`2183`	`2183`	`@_AttentionBackendRegistry.register(`
`2184`	`2184`	`AttentionBackendName._NATIVE_NPU,`
`2185`	`2185`	`constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],`
	`2186`	`+ supports_context_parallel=True,`
`2186`	`2187`	`)`
`2187`	`2188`	`def _native_npu_attention(`
`2188`	`2189`	`query: torch.Tensor,`