Add low-memory relay mode for T2I inference

[laj52]

Draft PR Notes: T2I Relay Memory Mode

Summary

This adds an opt-in memory_mode=relay path for text-to-image inference with KV-cache. The default remains memory_mode=parallel.

Relay lowers peak GPU memory by loading the UND tower for prefix/KV prefill, offloading it, loading the GEN tower for denoising, discarding Lance weights, and then loading the VAE for decode. The generated output remains bit-identical to parallel on the same hardware with the same seed and settings.

API

• MEMORY_MODE=parallel: default, fastest, existing behavior.
• MEMORY_MODE=relay: low-memory T2I path. Requires TASK_NAME=t2i and USE_KVCACHE=true.
• RELAY_MEMORY_LOG=true: optional CUDA allocator phase summaries.

No automatic memory policy is introduced in this PR.

Validation Evidence

RunPod A40, PyTorch 2.8.0+cu128, CUDA 12.8.

Prompt: A cat holds a poster with rainbow text "STOP"

Settings: task=t2i, num_frames=1, validation_num_timesteps=30, validation_timestep_shift=3.5, validation_data_seed=42, cfg_text_scale=4.0, use_KVcache=true.

Resolution	Mode	Peak GPU MiB	Time s	SHA256	Bit-identical
256res	parallel	29539	379.54	64f22bb7bb2a2221a17632d9066422b53a80782de016e58167d67a53f01a7e75	yes
256res	relay	6753	458.35	64f22bb7bb2a2221a17632d9066422b53a80782de016e58167d67a53f01a7e75	yes
512res	parallel	29535	313.49	c3e4bc78be423e80584cff6f524ad74291f2f0c2b2bf99a795f092a448dc336d	yes
512res	relay	6769	447.87	c3e4bc78be423e80584cff6f524ad74291f2f0c2b2bf99a795f092a448dc336d	yes
768res	parallel	29531	392.57	305d6378c799c1841efa4ffa8dfbc08dab8642b4c5c418b3c3494eeb8f1c8d69	yes
768res	relay	6809	500.06	305d6378c799c1841efa4ffa8dfbc08dab8642b4c5c418b3c3494eeb8f1c8d69	yes

Additional local smoke: relay T2I 256 completed on an RTX 3060 Ti 8GB without observed thrashing or VRAM swapping beyond the intentional relay transfers. Cross-hardware hashes are not expected to match because CUDA kernels and GPU architectures can produce small numeric differences.

Limitations

• T2I only.
• Requires KV-cache.
• Slower than parallel.
• Video, image editing, video editing, and understanding tasks are unchanged.

Reviewer Reproduction

python tools/compare_memory_modes_t2i.py \
  --model-path downloads/Lance_3B \
  --resolutions image_256res image_512res image_768res \
  --relay-memory-log