[BugFix] Fix repeated warnings from get_current_omni_diffusion_config outside init phase (#593)

Signed-off-by: Yipeng Liu <2206457717@qq.com>
This commit is contained in:
Yipeng Liu authored and GitHub committed 2026-01-03 22:43:41 +08:00
1 parent fe509c0c10
commit 4321b42e47
1 file changed
+5 -10
+5 -10
View File
@@ -46,6 +46,7 @@ class Attention(nn.Module):
causal=causal,
num_kv_heads=num_kv_heads,
)
self.backend_pref = None
self.softmax_scale = softmax_scale
self.scatter_idx = scatter_idx
@@ -59,6 +60,7 @@ class Attention(nn.Module):
try:
config = get_current_omni_diffusion_config()
self.backend_pref = config.attention_backend
if config.parallel_config.ring_degree > 1:
self.use_ring = True
try:
@@ -103,19 +105,12 @@ class Attention(nn.Module):
return out
def _run_local_attention(self, query, key, value, attn_metadata):
# Check backend preference from config
try:
config = get_current_omni_diffusion_config()
backend_pref = config.attention_backend
except Exception:
backend_pref = None
if backend_pref == "flash_attn" and query.dtype == torch.float32:
if self.backend_pref == "flash_attn" and query.dtype == torch.float32:
logger.warning(
"Flash Attention does not support float32. Overriding user config "
f"attention_backend='{backend_pref}' to 'sdpa' for dtype={query.dtype}."
f"attention_backend='{self.backend_pref}' to 'sdpa' for dtype={query.dtype}."
)
backend_pref = "sdpa"
self.backend_pref = "sdpa"
if is_npu():
return self.attention(