[Bug]: Cosmos3 sequence-parallel path attends zero-padded text K/V when CFG prompts have unequal lengths (cfg_size=1, ulysses/CP > 1)

System Info

• Found by code inspection of main at commit 3503e3f9bbf2673e019da6737ab008c58567ae32 (head of main on 2026-09-04). Not yet reproduced numerically; the mechanism is deterministic from the code, see the walk-through below.
• Affects the VisualGen Cosmos3 transformer (tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py) whenever sequence parallelism is active (ulysses_size > 1 and/or Attention2D / ring CP) and cfg_size == 1 with guidance_scale > 1.
• Hardware-independent.

Who can help?

Authors of #14012 (MoT / Cosmos3 sequence-parallel text K/V sharding) and #14827 (Cosmos3 audio output, which added the per-sample text-length loop on the single-GPU path).

Information

• The official example scripts (examples/visual_gen/models/cosmos3/cosmos3.py)
• My own modified scripts

Tasks

• An officially supported task in the examples folder (Cosmos3 text-to-video with CFG)
• My own task or dataset

Reproduction

Configuration that triggers it (2 GPUs is enough, Cosmos3-Nano fits):

parallel_config:
  cfg_size: 1        # <- CFG branches run as one batched B=2 call
  ulysses_size: 2    # <- any seq-parallel size > 1 (Attention2D / ring CP behave the same)

Run text-to-video with the default guidance_scale (6.0), the default video negative prompt, a fixed seed, and a short positive prompt (for example 15 words). Run the identical request on 1 GPU. Compare latents or decoded frames.

Code walk-through (all links at commit 3503e3f9bbf2673e019da6737ab008c58567ae32):

1. The prompt encoder right-pads both prompts to max_sequence_length (4096) and returns a 1/0 mask: pipeline_cosmos3.py#L1014-L1018.
2. With cfg_size == 1, the base pipeline concatenates the negative and positive prompts into one B=2 batch (text_ids / text_mask included): pipeline.py#L1053-L1060. The default video negative prompt is ~500 words of JSON, so the two samples have very different real lengths.
3. The text LLM zeroes pad rows before every layer, so pad-slot K/V are exactly zero (no projection bias): transformer_cosmos3.py#L1034.
4. Single-GPU path (exact): text K/V is trimmed to the batch-max real length and, because real_text_lens is passed, the cross-attention slices each sample to its own length in a per-sample loop: #L1728-L1738 and #L688-L703.
5. Sequence-parallel path (not exact): the cache is cut to max_real_len rounded up to sharder.size, the tail is zeroed, and the layer is called without real_text_lens: #L1585-L1600 and #L1739-L1746. The cross-attention therefore takes the batched branch (#L704-L714) with PredefinedAttentionMask.FULL and no key-padding mask.

The comment at L1586-1587 says "At most size-1 extra positions, negligible softmax dilution". That holds for B=1, but with the batched B=2 CFG call the shorter sample sees max_real_len - L_b extra zero slots, which is hundreds of tokens with the default negative prompt, not size - 1.

Expected behavior

For the same seed and prompts, a ulysses_size=2, cfg_size=1 run should match the single-GPU run up to floating-point reduction noise. Video tokens should attend only the real text tokens of their own CFG branch, which is what the single-GPU path and the RoPE position layout (positions start after the real text length, #L1250) already assume.

actual behavior

On the sequence-parallel path the shorter CFG branch (normally the positive prompt) additionally attends max_real_len - L_b zero-K/V slots plus the rounding tail. A zero key has logit 0, so each such slot adds exp(0) = 1 to the softmax denominator instead of being excluded. Multi-GPU results therefore differ from single-GPU results by construction, and the deviation grows with the length gap between the two prompts. The visual magnitude has not been measured yet; since the same softmax also spans all video keys, the effect may be modest, but it is systematic and prompt-dependent.

The shipped multi-GPU recipes (cosmos3-super-4gpu.yaml, cosmos3_t2v_bf16_gb200_nvl72.yaml) use cfg_size=2, which gives B=1 per CFG rank and avoids the issue except for the size-1 rounding tail. Any cfg_size=1 sequence-parallel run (for example a 2-GPU box with ulysses_size=2) is affected.

additional notes

Workaround: use cfg_size=2 for multi-GPU runs, or run on a single GPU.

Possible fixes, from smallest to cleanest:

1. Keep the text K/V for the shorter sample masked rather than zeroed on the sharded path: pass a key_padding_mask built from real_text_lens through _attn_impl. Note the Ulysses all-to-all gathers per-rank chunks, so the global key order is rank-major (text_r0, video_r0, text_r1, video_r1, ...); the valid keys are not a prefix and the FA4 seqused_k translation cannot be used as-is, only a mask-capable inner backend (VANILLA SDPA) would be exact.
2. Do not sequence-shard the text K/V at all. It is small (at most max_real_len tokens). Keep the trimmed full text K/V on every rank, select the local head shard to match the post-all-to-all head layout, and append it after the video K/V. The valid keys of sample b then form the prefix [0, S_video + L_b), so a per-sample length (seqused_k in FA4, or a key-padding mask in VANILLA) makes both CFG branches exact in a single launch and also removes the rounding tail. This would also let the single-GPU path drop its per-sample Python loop.

Start with examples/visual_gen/models/cosmos3/cosmos3.py and trace the single-GPU and sequence-parallel paths in tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py, especially text K/V handling and real_text_lens. Run the two-GPU cfg_size=1, ulysses_size=2 reproduction against a one-GPU run with fixed prompts and seed. Done means the shorter CFG branch attends only its real text tokens and results match up to floating-point reduction noise.