Enable cache support (FirstBlockCache) for Cosmos3 Omni

[johnnynunez]

What does this PR do?

Enables the native diffusers cache framework (CacheMixin → FirstBlockCacheConfig, etc.) for Cosmos3 Omni (nvidia/Cosmos3-Nano / nvidia/Cosmos3-Super, Cosmos3OmniPipeline).

Before this PR, pipe.transformer.enable_cache(...) fails because Cosmos3OmniTransformer doesn't inherit CacheMixin, and even with the mixin the FBC hooks can't read the block I/O: Cosmos3's MoT dual-stream decoder layers use forward(und_seq, gen_seq, rotary_emb) -> (und_seq, gen_seq) instead of the hidden_states/encoder_hidden_states contract.

Changes

1. Cosmos3OmniTransformer: add CacheMixin.
2. TransformerBlockRegistry: register Cosmos3VLTextMoTDecoderLayer with hidden_states_argument_name="gen_seq" / encoder_hidden_states_argument_name="und_seq" and return indices (1, 0). Mapping the gen stream to hidden_states is load-bearing: the cache decision must track the denoised tokens — the und (text/understanding) stream is quasi-static across denoising steps and would produce degenerate always-hit decisions.
3. TransformerBlockMetadata: add encoder_hidden_states_argument_name (analogous to the existing hidden_states_argument_name), and make the FBC hooks use both metadata fields instead of hardcoded "hidden_states" / "encoder_hidden_states" strings. Note the existing hidden_states_argument_name field (used by Kandinsky5's visual_embed) was silently ignored by first_block_cache.py — this also fixes that for any model with a custom argument name.
4. Cosmos3OmniPipeline: wrap the cond/uncond transformer calls in cache_context("cond"/"uncond"), following the Wan pipeline pattern, so stateful caches keep separate CFG state (the pipeline runs CFG as two separate transformer calls).

Measured (real hardware)

nvidia/Cosmos3-Nano T2V, RTX PRO 6000 Blackwell (96 GB, sm_120), torch 2.12.1+cu132, 189 frames / 35 steps / UniPC flow_shift=10, fixed seed, warmup excluded, median of 3:

Res	Config	Median	Speedup
480p	baseline (bf16, SDPA)	138.59s	1.00×
480p	FBC threshold=0.10	93.03s	1.49×
480p	FBC threshold=0.10 + torch.compile	81.20s	1.71×
256p	baseline	11.87s	1.00×
256p	FBC threshold=0.10	9.04s	1.31×

Output latents finite (no NaN, mean|v| ≈ 0.67–0.77, in line with dense runs). As expected for step-caching, the sample is different-but-equivalent vs the dense trajectory at the same seed (this is inherent to FBC, not specific to this integration).

Who can review?

@a-r-r-o-w @DN6 (cache hooks / Cosmos)

[github-actions]

Hi @johnnynunez, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.