Kandinsky 6: add context-parallel (_cp_plan) support to Kandinsky6Transformer3DModel #15156

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

On 2× T4, Kandinsky6Transformer3DModel sharded with device_map="balanced" keeps only one GPU busy at a time during denoising. A 200 ms nvidia-smi trace shows the GPUs taking turns on every step; both are busy in only 3 % of samples, and combined utilization is about one GPU's worth. This is expected for layer-wise model parallelism at batch size 1, so a second GPU adds memory but does not reduce per-clip latency. Context parallelism would let both GPUs work on the same clip, but in diffusers 0.42.1 Kandinsky6Transformer3DModel defines no _cp_plan (nor _tp_plan), so enable_parallelism(config=ContextParallelConfig(...)) is not available for it.

Raw trace: nvidia-smi-denoise-1793367281356898082.csv (timestamp, GPU index, utilization.gpu %, memory.used MiB, power.draw W)

Setup:

• Checkpoint: kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers @ 4c487e8b4aeeb9059e5a77905aa2ddb59ec28f81
• Hardware: Kaggle 2× Tesla T4 (16 GB each), DiT in fp16 (T4 has no native bf16)
• Versions: diffusers==0.42.1, transformers==5.22.1, accelerate==1.13.0, torch 2.11.0+cu128
• Sampling: T2V, 432×256, 25 frames, 10 steps (PiflowScheduler, n_grid=10), guidance_scale=1.0, batch size 1, joint video+audio (sample_audio=True)

repo, rev = "kandinskylab/Kandinsky-6.0-Lite-distill-5s-Diffusers", "4c487e8b4aeeb9059e5a77905aa2ddb59ec28f81"
transformer = Kandinsky6Transformer3DModel.from_pretrained(
    repo, subfolder="transformer", revision=rev, torch_dtype=torch.float16,
    device_map="balanced", max_memory={0: "10GiB", 1: "10GiB"})
pipe = Kandinsky6TI2VAPipeline.from_pretrained(
    repo, revision=rev, transformer=transformer, text_encoder=None, text_encoder_2=None,
    tokenizer=None, tokenizer_2=None, torch_dtype={"default": torch.float32})
# No pipe.to(...) / enable_model_cpu_offload(), so the device map is kept.
# Prompt embeddings are precomputed in a separate text-encoding stage and passed in.
output = pipe(prompt_embeds=..., pooled_prompt_embeds=..., num_inference_steps=10,
              guidance_scale=1.0, sample_audio=True, output_type="latent")

Measurements (trace started right before and stopped right after the pipe(...) call; 67 polls ≈ 13.4 s):

Each of the 10 steps (~1.3 s) runs through the cuda:0 blocks and then the cuda:1 blocks.

Describe the solution you'd like.

Add a _cp_plan to Kandinsky6Transformer3DModel so that batch-1 inference can use context parallelism:

from diffusers import ContextParallelConfig
pipe.transformer.enable_parallelism(config=ContextParallelConfig(ulysses_degree=2))  # under torchrun, 2 processes

In 0.42.1, _cp_plan is already defined for Wan, LTX, LTX-2, Flux, Flux2, QwenImage, Cosmos3, ChronoEdit, Helios and MiniMax-H3. LTX-2 is also a joint audio+video model, so its plan may be a good template for Kandinsky 6's dual-stream video/audio blocks. Open questions:

• Should only the visual token stream be sharded, or the audio stream as well (it is much shorter)?
• Do the text/visual cross-attention and the rotary embeddings need special handling in the plan?

Describe alternatives you've considered.

• Current setup: device_map="balanced" on the transformer only. It works and fits memory, but GPUs alternate (data above), so per-clip latency does not improve. Is this the recommended multi-GPU setup for Kandinsky6TI2VAPipeline until context parallelism exists, or is there a better-supported pipeline-level strategy?
• One full pipeline per GPU. At this resolution each GPU uses only ~3.4 GiB, so the whole DiT (5.94 GiB fp16) should fit on one T4 (not tested yet). Two independent clips at once would double throughput, but not reduce per-clip latency. Fit at the full 121-frame 864×480 size is unknown.
• CFG parallelism. Not applicable: the distilled Lite model runs with guidance_scale=1.0, so there is no cond/uncond pair to split.
• Pipeline-level device_map. It places whole components on devices; since stages run one after another, it does not make two GPUs compute at the same time.

Additional context.

• Notebook (Kaggle, 2× T4, version 4): https://www.kaggle.com/code/applehawk/kandinsky-6-lite-diffusers-generation
• Pipeline placement in that notebook: text encoders on cuda:1 (Qwen with partial CPU offload), DiT split by device_map="balanced", video VAE (fp32, tiled) on cuda:1, audio VAE + vocoder (fp32) on cuda:0; each stage is loaded and freed before the next.
• utilization.gpu is the share of time a kernel was running, not how much compute was used; the 50 % threshold is my choice, and the plot shows the alternation regardless of it.
• Scope: a single run at the small Smoke resolution; the full Baseline size and a single-GPU comparison have not been measured yet.