[Bugfix][Examples][MammothModa2] Serve T2I via LLM-typed DiT stage (#7199) + benchmark harness and profiler wiring (#7293)

Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -337,7 +338,7 @@ class _RequestConditioning:
     image_mask: torch.Tensor
 
 
-class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
+class MammothModa2DiTPipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
     """
     MammothModa2 DiT + VAE generation stage (non-autoregressive).
 
@@ -424,12 +425,6 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
 
         self._llm_hidden_size = llm_hidden_size
 
-        # Cache-DiT lifecycle: the diffusion runner enables the configured
-        # backend (``cache_backend`` on the deploy YAML stage entry) at startup
-        # and transfers ownership here via the request-scoped protocol;
-        # forward() then reconciles per-request state (step count, CFG parity).
-        # The runner also emits the cache summary when
-        # ``enable_cache_dit_summary`` is set.
         self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
         self._cache_dit_config: DiffusionCacheConfig | None = None
         if str(getattr(od_config, "cache_backend", "") or "").lower() == "cache_dit":
@@ -440,6 +435,14 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 else DiffusionCacheConfig.from_dict(cache_config or {})
             )
 
+        profiler_targets = ["gen_transformer.forward", "gen_vae.decode"]
+        if self.gen_image_condition_refiner is not None:
+            profiler_targets.append("gen_image_condition_refiner.forward")
+        self.setup_diffusion_pipeline_profiler(
+            profiler_targets=profiler_targets,
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
+
     def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
         """Adopt a runner-installed Cache-DiT backend (request-scoped protocol)."""