diff --git a/tests/diffusion/attention/test_attention_config.py b/tests/diffusion/attention/test_attention_config.py index 297f760fb..35be5e6f0 100644 --- a/tests/diffusion/attention/test_attention_config.py +++ b/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" diff --git a/vllm_omni/diffusion/attention/layer.py b/vllm_omni/diffusion/attention/layer.py index 4d7c0d362..bbf3481fe 100644 --- a/vllm_omni/diffusion/attention/layer.py +++ b/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, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index e2f91d9aa..5a87e7762 100644 --- a/vllm_omni/diffusion/data.py +++ b/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, ) diff --git a/vllm_omni/diffusion/hooks/sequence_parallel.py b/vllm_omni/diffusion/hooks/sequence_parallel.py index 2aaa3c169..26a7d1370 100644 --- a/vllm_omni/diffusion/hooks/sequence_parallel.py +++ b/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", diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py index e687cb7b2..8163f7854 100644 --- a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py +++ b/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: