[BugFix] Rename attention_config to diffusion_attention_config (#3489)

Signed-off-by: gcanlin <canlinguosdu@gmail.com>

tests/diffusion/attention/test_attention_config.py

@@ -285,14 +285,14 @@ class TestOmniDiffusionConfigAttentionParsing:
 
     def test_diffusion_attention_backend_sets_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="SAGE_ATTN")
-        assert isinstance(config.attention_config, AttentionConfig)
-        assert config.attention_config.default is not None
-        assert config.attention_config.default.backend == "SAGE_ATTN"
+        assert isinstance(config.diffusion_attention_config, AttentionConfig)
+        assert config.diffusion_attention_config.default is not None
+        assert config.diffusion_attention_config.default.backend == "SAGE_ATTN"
 
     def test_diffusion_attention_backend_auto_means_platform_default(self):
         config = OmniDiffusionConfig.from_kwargs(diffusion_attention_backend="auto")
-        assert isinstance(config.attention_config, AttentionConfig)
-        assert config.attention_config.default is None
+        assert isinstance(config.diffusion_attention_config, AttentionConfig)
+        assert config.diffusion_attention_config.default is None
 
     def test_diffusion_attention_backend_and_default_are_mutually_exclusive(self):
         with pytest.raises(ValueError):
@@ -303,27 +303,19 @@ class TestOmniDiffusionConfigAttentionParsing:
 
     def test_dict_diffusion_attention_config(self):
         config = OmniDiffusionConfig(
-            attention_config={
+            diffusion_attention_config={
                 "default": {"backend": "FLASH_ATTN"},
                 "per_role": {"self": "SPARSE_BLOCK"},
             }
         )
-        assert config.attention_config.default.backend == "FLASH_ATTN"
-        assert config.attention_config.per_role["self"].backend == "SPARSE_BLOCK"
+        assert config.diffusion_attention_config.default.backend == "FLASH_ATTN"
+        assert config.diffusion_attention_config.per_role["self"].backend == "SPARSE_BLOCK"
 
-    def test_legacy_attention_config_name_maps_in_from_kwargs(self):
-        config = OmniDiffusionConfig.from_kwargs(
-            attention_config={
-                "default": {"backend": "FLASH_ATTN"},
-            }
-        )
-        assert config.attention_config.default.backend == "FLASH_ATTN"
-
-    def test_no_attention_config_defaults_to_empty(self):
+    def test_no_diffusion_attention_config_defaults_to_empty(self):
         config = OmniDiffusionConfig()
-        assert isinstance(config.attention_config, AttentionConfig)
-        assert config.attention_config.default is None
-        assert config.attention_config.per_role == {}
+        assert isinstance(config.diffusion_attention_config, AttentionConfig)
+        assert config.diffusion_attention_config.default is None
+        assert config.diffusion_attention_config.per_role == {}
 
 
 class TestCurrentDiffusionConfig:
@@ -397,7 +389,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         )
 
         od_config = SimpleNamespace(
-            attention_config=AttentionConfig(
+            diffusion_attention_config=AttentionConfig(
                 default=AttentionSpec(backend="FLASH_ATTN"),
                 per_role={"cross": AttentionSpec(backend="TORCH_SDPA", extra={"block_size": 128})},
             ),
@@ -418,7 +410,7 @@ class TestAttentionInitUsesCurrentDiffusionConfig:
         assert captured["role"] == "cross"
         assert captured["role_category"] == "cross"
         assert captured["head_size"] == 64
-        assert captured["attention_config"] is od_config.attention_config
+        assert captured["attention_config"] is od_config.diffusion_attention_config
         assert attn.backend_pref == "TORCH_SDPA"
         assert attn.attention.kwargs["backend_kwargs"] == {"block_size": 128}
         assert attn.attention.kwargs["qkv_layout"] == "BSND"

vllm_omni/diffusion/attention/layer.py

@@ -57,7 +57,7 @@ class Attention(nn.Module):
         self.backend_pref = None
 
         config = get_current_diffusion_config_or_none()
-        attention_config = config.attention_config if config is not None else None
+        attention_config = config.diffusion_attention_config if config is not None else None
 
         attn_backend_cls, spec = get_attn_backend_for_role(
             role=role,

vllm_omni/diffusion/data.py

@@ -367,7 +367,7 @@ class OmniDiffusionConfig:
     tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
 
     # Attention
-    attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
+    diffusion_attention_config: "AttentionConfig" = field(default_factory=lambda: AttentionConfig())
 
     # Running mode
     # mode: ExecutionMode = ExecutionMode.INFERENCE
@@ -655,7 +655,7 @@ class OmniDiffusionConfig:
 
         # Match vLLM's config flow: parse entrypoint shorthands before the
         # config object is built, and keep a single runtime truth source.
-        self.attention_config = build_attention_config(self.attention_config)
+        self.diffusion_attention_config = build_attention_config(self.diffusion_attention_config)
 
         if self.max_cpu_loras is None:
             self.max_cpu_loras = 1
@@ -798,20 +798,12 @@ class OmniDiffusionConfig:
         else:
             kwargs.pop("quantization", None)
 
-        # Map "diffusion_attention_config" to "attention_config" so the
-        # engine_args key (namespaced to avoid collision with vLLM's own
-        # attention_config) is mapped to the dataclass field.
-        if "diffusion_attention_config" in kwargs and "attention_config" not in kwargs:
-            kwargs["attention_config"] = kwargs.pop("diffusion_attention_config")
-        else:
-            kwargs.pop("diffusion_attention_config", None)
-
         # Handle "diffusion_attention_backend" shorthand: merge into
-        # attention_config before field filtering.
+        # diffusion_attention_config before field filtering.
         diffusion_attn_backend = kwargs.pop("diffusion_attention_backend", None)
         if diffusion_attn_backend is not None:
-            existing = kwargs.get("attention_config")
-            kwargs["attention_config"] = parse_attention_config(
+            existing = kwargs.get("diffusion_attention_config")
+            kwargs["diffusion_attention_config"] = parse_attention_config(
                 existing,
                 attention_backend=diffusion_attn_backend,
             )

vllm_omni/diffusion/hooks/sequence_parallel.py

@@ -399,7 +399,7 @@ class SequenceParallelSplitHook(ModelHook):
         if is_forward_context_available():
             od_config = get_forward_context().omni_diffusion_config
             if od_config is not None:
-                attention_config = od_config.attention_config
+                attention_config = od_config.diffusion_attention_config
 
         attn_backend, _ = get_attn_backend_for_role(
             role="self",

vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py

@@ -210,7 +210,7 @@ class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin):
             logging.info("No transformer found in diffusers pipeline. Skipping attention backend setting.")
             return
 
-        default_spec = self.od_config.attention_config.default
+        default_spec = self.od_config.diffusion_attention_config.default
         attention_backend_config = default_spec.backend if default_spec is not None else None
         attention_backend_attempts: list[str] = []
         match attention_backend_config: