Cosmos3 Omni transformer recomputes the text-prefix keys and values on every denoising step #15286

Is your feature request related to a problem? Please describe.

Cosmos3OmniTransformer (used by Cosmos3OmniPipeline) runs the text prefix (the und pathway, the first und_len tokens) through every decoder layer on every transformer call:

• forward fills the und rows with embed_tokens(input_ids). The timestep embeddings are added only to the vision, sound and action tokens, which are placed after und_len.
• In Cosmos3AttnProcessor.__call__ the und queries attend only to the und keys and values (causal SDPA), and Cosmos3VLTextMoTDecoderLayer.forward applies the und residual and MLP to the und tokens only. The gen tokens read the und pathway only through torch.cat([k_und_for_gen, k_gen]) and torch.cat([v_und, v_gen]), and the predictions are read only at the vision, sound and action indexes.
• Cosmos3OmniPipeline.__call__ builds the cond and uncond inputs once (cond_packed_static, uncond_packed_static) and passes the same input_ids and position_ids on every step.

So within one pipeline call, the und keys and values of each layer are the same on every step of a CFG branch, and they are computed again on every transformer call. nvidia/Cosmos3-Edge has 28 decoder layers; with the model card i2v prompt and negative prompt the text prefix is 1912 tokens (cond) and 3519 tokens (uncond), and with guidance 6.0 and 20 steps the pipeline calls the transformer 40 times, so per pipeline call the und pathway runs 40 x 28 times (6720 linear and 1120 SDPA calls).

Describe the solution you'd like.

Within one pipeline call, the und keys and values of each layer should be computed once per CFG branch and reused on the later steps, with the same output as now.

To check that this is worth doing, I measured main and a local branch with the change with the script below (one process per checkout, median of 3 calls after one warm-up call, card seed 0): full pipeline call including the safety checker, eager, nvidia/Cosmos3-Edge (revision 5b59028) with the model card "Image to video generation" prompt, negative prompt, input image and settings (480x832 (height x width), 121 frames, 20 steps, guidance 6.0, UniPC flow_shift=12.0, enable_safety_checker=True, pipe.to("cuda")), bf16, batch 1, default attention backend:

Calls: main 55.325, 55.385, 55.237 s; local branch 53.076, 53.017, 53.050 s. Peak memory over the 3 timed calls: 12.39 GiB allocated / 16.25 GiB reserved on both.

Environment: NVIDIA GeForce RTX 5090 (Blackwell, compute capability 12.0 / sm_120; Windows 11 10.0.26200, driver 621.01), Python 3.12.16, torch 2.14.1+cu130, torchvision 0.29.1+cu130, transformers 5.16.1, tokenizers 0.23.2, accelerate 1.9.0, huggingface_hub 1.32.0, safetensors 0.8.0, numpy 2.5.4, Pillow 13.0.0, cosmos_guardrail 0.3.2 (package versions read in the same environment after the runs), diffusers main 43e3002 from a source checkout. cosmos_guardrail 0.3.2 and its dependencies are installed in a separate directory on PYTHONPATH; the safety checker weights are nvidia/Cosmos-1.0-Guardrail revision cf03c03. Environment variables set for the runs: HF_HUB_OFFLINE=1, HF_HOME (local cache path), PYTHONIOENCODING=utf-8, PYTHONUNBUFFERED=1, PYTHONPATH=<checkout>/src;<cosmos_guardrail directory>, PATH with a ruff tool directory prepended; HF_ENABLE_PARALLEL_LOADING, DIFFUSERS_ATTN_BACKEND and CUDA_VISIBLE_DEVICES unset.

The script follows the card call and adds output_type="pt" so that the outputs can be compared value by value; it reads the card's json assets as utf-8 (negative_prompt.json has non-ASCII characters).

# Cosmos3-Edge image-to-video pipeline timing at the model card settings, one process per checkout:
#   PYTHONPATH=<diffusers checkout A>/src python time_cosmos3edge.py a.pt
#   PYTHONPATH=<diffusers checkout B>/src python time_cosmos3edge.py b.pt
#   python time_cosmos3edge.py --compare a.pt b.pt
# Options: --calls N (timed calls after one warm-up call, default 3)
# Settings: the "Image to video generation" block of the nvidia/Cosmos3-Edge model card, safety checker on (needs
# the cosmos_guardrail package). Additions: output_type "pt" (raw floats for the comparison); the json assets are
# read as utf-8 (negative_prompt.json has non-ASCII characters).
import argparse
import hashlib
import json
import time

import torch


p = argparse.ArgumentParser()
p.add_argument("out")
p.add_argument("other", nargs="?")
p.add_argument("--compare", action="store_true")
p.add_argument("--calls", type=int, default=3)
args = p.parse_args()

if args.compare:
    a, b = torch.load(args.out), torch.load(args.other)
    print("differing values:", int((a != b).sum()), "of", a.numel())
    raise SystemExit

from huggingface_hub import hf_hub_download  # noqa: E402

import diffusers  # noqa: E402
from diffusers import Cosmos3OmniPipeline, UniPCMultistepScheduler  # noqa: E402
from diffusers.utils import load_image  # noqa: E402


print("diffusers", diffusers.__version__, "from", diffusers.__file__)
repo = "nvidia/Cosmos3-Edge"


def asset(name):
    with open(hf_hub_download(repo, f"assets/{name}"), encoding="utf-8") as f:
        return json.dumps(json.load(f))


prompt = asset("example_i2v_prompt.json")
negative_prompt = asset("negative_prompt.json")
image = load_image(hf_hub_download(repo, "assets/example_i2v_input.jpg"))
print(
    "prompt md5",
    hashlib.md5(prompt.encode()).hexdigest(),
    "negative md5",
    hashlib.md5(negative_prompt.encode()).hexdigest(),
)

pipe = Cosmos3OmniPipeline.from_pretrained(repo, torch_dtype=torch.bfloat16, enable_safety_checker=True)
pipe.to("cuda")
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=12.0, use_karras_sigmas=False)
pipe.set_progress_bar_config(disable=True)


def call():
    return pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image,
        num_frames=121,
        height=480,
        width=832,
        fps=24.0,
        num_inference_steps=20,
        guidance_scale=6.0,
        generator=torch.Generator(device="cuda").manual_seed(0),
        add_resolution_template=False,
        add_duration_template=False,
        output_type="pt",
    ).video


call()  # warm-up
torch.cuda.reset_peak_memory_stats()
times = []
for _ in range(args.calls):
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    video = call()
    torch.cuda.synchronize()
    times.append(time.perf_counter() - t0)
median = sorted(times)[len(times) // 2]
print(f"median {median:.3f} s over {len(times)} calls ({', '.join(f'{t:.3f}' for t in times)})")
print(
    f"peak memory over the {len(times)} calls: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GiB allocated, "
    f"{torch.cuda.max_memory_reserved() / 1024**3:.2f} GiB reserved"
)
print("video", tuple(video.shape), video.dtype)
torch.save(video.float().cpu(), args.out)

Describe alternatives you've considered.

TextKVCacheConfig (hooks/text_kv_cache.py) caches text keys and values across steps, but it supports only NucleusMoEImageTransformerBlock and does not apply to Cosmos3OmniTransformer.

Additional context.

None.