Merge branch 'main' into qwen-image-2.1

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -1,13 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from contextlib import contextmanager
 from dataclasses import dataclass
+from types import SimpleNamespace
 
+import cache_dit
 import pytest
 import torch
+from cache_dit import BlockAdapter
 from torch import nn
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig
+from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, RequestScopedCacheDiTRuntime
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import (
+    AttentionConfig,
+    DiffusionCacheConfig,
+    DiffusionOutput,
+    OmniDiffusionConfig,
+    TransformerConfig,
+)
+from vllm_omni.diffusion.models.interface import adopt_request_scoped_cache_dit
 from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
 from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     MammothModa2DiTPipeline,
@@ -17,6 +31,7 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -110,6 +125,10 @@ def _pipeline_shell() -> MammothModa2DiTPipeline:
     pipeline.device = torch.device("cpu")
     pipeline.config = _build_mammoth_config(_od_config())
     pipeline._llm_hidden_size = 8
+    # Cache-DiT runtime state that __init__ normally creates (disabled here,
+    # matching a pipeline built without a cache_backend).
+    pipeline._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipeline)
+    pipeline._cache_dit_config = None
     return pipeline
 
 
@@ -118,12 +137,14 @@ def _batch(
     request_id: str = "req-7",
     prompt: object | None = None,
     sampling: OmniDiffusionSamplingParams | None = None,
+    height: int = 32,
+    width: int = 48,
 ) -> DiffusionRequestBatch:
     if prompt is None:
         prompt = {
             "prompt": "",
-            "height": 32,
-            "width": 48,
+            "height": height,
+            "width": width,
             "additional_information": {
                 "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8),
                 "full_token_ids": [10, 11, 100, 101],
@@ -390,3 +411,162 @@ def test_forward_rejects_missing_visual_tokens_before_model_access() -> None:
     }
     with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"):
         _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt))
+
+
+@contextmanager
+def _force_torch_sdpa():
+    """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3)."""
+    od_config = SimpleNamespace(
+        diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+        parallel_config=SimpleNamespace(ring_degree=1),
+    )
+    with set_current_diffusion_config(od_config):
+        yield
+
+
+def _constructible_raw_config() -> dict:
+    """Checkpoint config whose DiT/VAE are small enough to actually run on CPU."""
+    return {
+        "model_type": "mammothmoda2",
+        "llm_config": {
+            "model_type": "mammothmoda2_qwen2_5_vl",
+            "text_config": {
+                "model_type": "mammothmoda2_qwen2_5_vl_text",
+                "hidden_size": 8,
+                "gen_vocab_start_index": 100,
+            },
+        },
+        "gen_vae_config": {"block_out_channels": [8, 8], "norm_num_groups": 8},
+        "gen_dit_config": {
+            "patch_size": 2,
+            "in_channels": 4,
+            "hidden_size": 96,
+            "num_layers": 4,
+            "num_refiner_layers": 1,
+            "num_attention_heads": 2,
+            "num_kv_heads": 2,
+            "multiple_of": 8,
+            "axes_dim_rope": [16, 16, 16],
+            "axes_lens": [300, 512, 512],
+            "text_feat_dim": 16,
+        },
+        # Top-level axes feed the pipeline-owned rotary table (see pipeline __init__).
+        "gen_axes_dim_rope": [16, 16, 16],
+        "gen_axes_lens": [300, 512, 512],
+    }
+
+
+def _cache_dit_od_config() -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(_constructible_raw_config()),
+        cache_backend="cache_dit",
+        cache_config=DiffusionCacheConfig(),
+    )
+
+
+@pytest.mark.skipif(
+    current_omni_platform.is_rocm(),
+    reason="vLLM ROCm custom ops lack CPU fallback",
+)
+def test_forward_transitions_request_scoped_cache_dit_across_requests(request: pytest.FixtureRequest) -> None:
+    """Cache-DiT request lifecycle on one production-built pipeline instance.
+
+    Mirrors the diffusion runner startup (``get_cache_backend`` → ``enable`` →
+    ``adopt_request_scoped_cache_dit``) on a tiny pipeline built through the
+    real ``__init__``, then drives three requests through the real ``forward()``:
+
+    1. a partial-CFG request installs the request-scoped cache context:
+       cond/uncond forwards stay paired on every denoise step (parity
+       accounting), even outside ``cfg_range``, and blocks record cached steps;
+    2. a guidance=1.0 request runs with hooks disabled (``prepare(None)``) and
+       its output is bit-identical to a reference captured before the backend
+       was ever enabled, proving a clean uninstall with no stale module state;
+    3. a CFG request with a different step count and image size re-installs the
+       hooks through a fresh context, so no stale shape or step count survives.
+    """
+    torch.manual_seed(0)  # Deterministic random init → deterministic cache hits.
+    with _force_torch_sdpa():
+        pipeline = MammothModa2DiTPipeline(od_config=_cache_dit_od_config())
+    pipeline.eval()
+    assert pipeline._cache_dit_config is not None
+    request.addfinalizer(pipeline._cache_dit_runtime.disable)
+
+    # The default prompt carries 2 text + 2 visual condition tokens, so the
+    # conditional pass sees 4 tokens and the uncond pass (empty prompt) sees 0.
+    cond_tokens, uncond_tokens = 4, 0
+    seq_lens: list[int] = []
+
+    def record_call(_module, _args, kwargs):
+        text = kwargs.get("text_hidden_states")
+        seq_lens.append(int(text.shape[1]) if text is not None else -1)
+
+    hook = pipeline.gen_transformer.register_forward_pre_hook(record_call, with_kwargs=True)
+
+    unguided_sampling = OmniDiffusionSamplingParams(
+        height=32, width=48, seed=4321, guidance_scale=1.0, num_inference_steps=8
+    )
+    partial_cfg_sampling = OmniDiffusionSamplingParams(
+        height=32,
+        width=48,
+        seed=1234,
+        guidance_scale=4.0,
+        num_inference_steps=8,
+        extra_args={"cfg_range": [0.2, 0.8]},
+    )
+    changed_cfg_sampling = OmniDiffusionSamplingParams(
+        height=48,
+        width=64,
+        seed=777,
+        guidance_scale=6.0,
+        num_inference_steps=10,
+        extra_args={"cfg_range": [0.1, 0.9]},
+    )
+
+    # Reference run for phase 2: identical request while hooks were never
+    # installed on this pipeline.
+    assert not pipeline.is_cache_dit_enabled()
+    reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    assert seq_lens == [cond_tokens] * 8
+    assert torch.isfinite(reference.float()).all()
+
+    # Runner startup: enable once, then the pipeline owns all transitions.
+    backend = get_cache_backend("cache_dit", DiffusionCacheConfig())
+    assert isinstance(backend, CacheDiTBackend)
+    backend.enable(pipeline)
+    assert adopt_request_scoped_cache_dit(pipeline, backend)
+    assert pipeline.is_cache_dit_enabled()
+
+    # Phase 1: partial CFG keeps cond/uncond paired on every step, including
+    # the steps whose normalized index falls outside cfg_range [0.2, 0.8].
+    seq_lens.clear()
+    pipeline.forward(_batch(sampling=partial_cfg_sampling))
+    assert seq_lens == [cond_tokens, uncond_tokens] * 8
+    assert pipeline.is_cache_dit_enabled()
+    assert BlockAdapter.is_cached(pipeline.gen_transformer)
+    hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
+    assert hits.cached_steps or hits.cfg_cached_steps
+
+    # Phase 2: no-CFG request disables the hooks and must reproduce the
+    # never-enabled reference bit-for-bit.
+    seq_lens.clear()
+    unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    assert seq_lens == [cond_tokens] * 8
+    assert not pipeline.is_cache_dit_enabled()
+    assert not BlockAdapter.is_cached(pipeline.gen_transformer)
+    assert torch.equal(unguided, reference)
+
+    # Phase 3: CFG with a new step count and image size re-installs the hooks
+    # on a fresh context; output shape follows the new request geometry.
+    seq_lens.clear()
+    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
+    assert seq_lens == [cond_tokens, uncond_tokens] * 10
+    assert pipeline.is_cache_dit_enabled()
+    assert BlockAdapter.is_cached(pipeline.gen_transformer)
+    hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0]
+    assert hits.cached_steps or hits.cfg_cached_steps
+    assert regenerated.shape == (1, 3, 12, 16)
+    assert torch.isfinite(regenerated.float()).all()
+
+    hook.remove()