use per batch attention with sliced KV

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

Author: NVShreyas

examples/visual_gen/configs/cosmos3-nano-1gpu.yaml

@@ -13,20 +13,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# 1-GPU Cosmos3 (Nano / Super) with FP8 dynamic quantization.
+# 1-GPU Cosmos3 (Nano / Super).
 # Model: nvidia/Cosmos3-Nano or nvidia/Cosmos3-Super
 # Shared by offline examples (--visual_gen_args) and trtllm-serve.
 #
 # Cosmos3 constraints: VANILLA attention only;
-# no Attention2D / Ring. Use CFG + Ulysses for multi-GPU (see cosmos3-super-4gpu.yaml).
-quant_config:
-  quant_algo: FP8
-  dynamic: true
-  ignore: ["language_model.*", "vae2llm", "llm2vae", "time_embedder.*"]
+# Use CFG + Ulysses for multi-GPU (see cosmos3-super-4gpu.yaml).
 attention_config:
   backend: VANILLA
 parallel_config:
   cfg_size: 1
   ulysses_size: 1
-cuda_graph_config:
-  enable: false

examples/visual_gen/configs/cosmos3-super-4gpu.yaml

@@ -13,16 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# 4-GPU Cosmos3-Super with FP8 dynamic quantization (CFG + Ulysses + parallel VAE).
+# 4-GPU Cosmos3-Super with (CFG + Ulysses + parallel VAE).
 # Launch with 4 processes, e.g. torchrun --nproc_per_node=4 ...
 # Model: nvidia/Cosmos3-Super
 # Shared by offline examples (--visual_gen_args) and trtllm-serve.
 #
 # GPU layout: cfg_size=2 (positive | negative) x ulysses_size=2 (sequence split).
-quant_config:
-  quant_algo: FP8
-  dynamic: true
-  ignore: ["language_model.*", "vae2llm", "llm2vae", "time_embedder.*"]
 attention_config:
   backend: VANILLA
 parallel_config:

examples/visual_gen/models/cosmos3_ti2v.py

@@ -45,7 +45,7 @@
 - ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE
 
 Usage:
-    python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \\
+    python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \
         --prompt "The video opens with a view of a well-lit indoor space featuring a " \\
         "wooden display case with compartments filled with various fruits, " \\
         "including bananas, apples, pears, oranges, and carambolas. " \\
@@ -66,11 +66,38 @@
         "position, leaving the display case and surrounding area unchanged. " \\
         "The video showcases a seamless and efficient automated fruit-picking " \\
         "process, highlighting the precision and efficiency of modern robotics " \\
-        "in a retail setting." \\
+        "in a retail setting." \
         --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml
+
+    python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \
+        --prompt "A low-angle tracking shot follows a man riding a vintage black motorcycle " \\
+        "across a lush green grassy yard. Sunlight filters through overhead trees, casting " \\
+        "dappled shadows across the vibrating chrome exhaust and the rider's leather jacket. " \\
+        "He kicks up small blades of grass as he maneuvers the bike. He gradually decelerates, " \\
+        "the front fork compressing slightly as he brakes to a smooth halt beside another " \\
+        "individual standing in the shade. The camera settles into a medium two-shot, capturing " \\
+        "the rider lifting his visor to speak, his face framed by a matte helmet. The video is " \\
+        "8 seconds long and is of 24 FPS. This video is of 1280x720 resolution. Audio description: " \\
+        "The rhythmic, mechanical chugging of a four-stroke motorcycle engine dominates the " \\
+        "foreground, characterized by a throaty, guttural timbre. Periodic high-pitched revs " \\
+        "punctuate the steady idle as the throttle is twisted. The sound of tires crunching " \\
+        "softly over dry grass and twigs provides a textured background layer. As the vehicle " \\
+        "slows, the engine note drops to a low-frequency rumble before clicking into neutral. " \\
+        "A muffled, mid-range male voice begins speaking, accompanied by the metallic clink of " \\
+        "a helmet visor snapping upward and the faint chirping of distant birds in an open-air " \\
+        "environment." \
+        --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
+        --enable_audio
+
+    python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \
+        --prompt "A cute puppy playing with a ball in a park" \
+        --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \
+        --output_type image \
+        --output_path output.png
 """
 
 import argparse
+import json
 
 from tensorrt_llm import VisualGen, VisualGenArgs
 
@@ -109,6 +136,26 @@ def main():
         default="cosmos3_ti2v_output.mp4",
         help="Path to save the output video",
     )
+    parser.add_argument(
+        "--enable_duration_template", action="store_true", help="Enable duration template in prompt"
+    )
+    parser.add_argument(
+        "--enable_resolution_template",
+        action="store_true",
+        help="Enable resolution template in prompt",
+    )
+    parser.add_argument(
+        "--use_system_prompt", action="store_true", help="Use system prompt in prompt"
+    )
+    parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation")
+    parser.add_argument(
+        "--output_type", type=str, default="video", help="Output type (video, image)"
+    )
+
+    # Guardrails
+    parser.add_argument(
+        "--disable_guardrails", action="store_true", help="NOT RECOMMENDED: Disable guardrails"
+    )
     args = parser.parse_args()
 
     # Engine config from shared YAML (optional); model-specific defaults apply otherwise.
@@ -121,13 +168,25 @@ def main():
     if args.image_path is not None:
         params.image = args.image_path
 
+    negative_prompt = json.load(open("neg_prompt.json"))
+
+    params.extra_params["use_duration_template"] = args.enable_duration_template
+    params.extra_params["use_resolution_template"] = args.enable_resolution_template
+    params.extra_params["use_system_prompt"] = args.use_system_prompt
+    params.extra_params["enable_audio"] = args.enable_audio
+    params.extra_params["use_guardrails"] = not args.disable_guardrails
+    params.extra_params["output_type"] = args.output_type
+
+    params.negative_prompt = json.dumps(negative_prompt)
+
     output = visual_gen.generate(
         inputs=args.prompt,
         params=params,
     )
 
     output.save(args.output_path)
     print(f"Saved: {args.output_path}")
+    print(output.metrics)
 
 
 if __name__ == "__main__":

tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py

@@ -343,6 +343,7 @@ def forward(
         v_und: torch.Tensor,
         freqs_cos: torch.Tensor,
         freqs_sin: torch.Tensor,
+        real_text_lens: Optional[list[int]],
     ) -> torch.Tensor:
         """
         Args:
@@ -366,15 +367,31 @@ def forward(
         q, k = self.apply_qk_norm(q, k)
         q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin)
 
-        k_all = torch.cat([k_und, k], dim=1).contiguous()
-        v_all = torch.cat([v_und, v], dim=1).contiguous()
-
-        out = self._attn_impl(
-            q,
-            k_all,
-            v_all,
-            attention_mask=PredefinedAttentionMask.FULL,
-        )
+        if real_text_lens is not None and batch_size > 1:
+            outs = []
+            for b in range(batch_size):
+                Lb = int(real_text_lens[b])
+                k_all_b = torch.cat([k_und[b : b + 1, :Lb], k[b : b + 1]], dim=1)
+                v_all_b = torch.cat([v_und[b : b + 1, :Lb], v[b : b + 1]], dim=1)
+                outs.append(
+                    self._attn_impl(
+                        q[b : b + 1],
+                        k_all_b,
+                        v_all_b,
+                        attention_mask=PredefinedAttentionMask.FULL,
+                    )
+                )
+            out = torch.cat(outs, dim=0)
+        else:
+            k_all = torch.cat([k_und, k], dim=1).contiguous()
+            v_all = torch.cat([v_und, v], dim=1).contiguous()
+
+            out = self._attn_impl(
+                q,
+                k_all,
+                v_all,
+                attention_mask=PredefinedAttentionMask.FULL,
+            )
 
         return self.to_out[0](out)
 
@@ -485,6 +502,7 @@ def forward(
         k_und: torch.Tensor,
         v_und: torch.Tensor,
         freqs: Tuple[torch.Tensor, torch.Tensor],
+        real_text_lens: Optional[list[int]] = None,
     ) -> torch.Tensor:
         residual = hidden_states
         hidden_states = self.input_layernorm(hidden_states)
@@ -496,6 +514,7 @@ def forward(
             v_und=v_und,
             freqs_cos=cos,
             freqs_sin=sin,
+            real_text_lens=real_text_lens,
         )
         hidden_states = residual + hidden_states
 
@@ -968,6 +987,7 @@ def forward(
         T, H, W = video_shape
         Hp, Wp, _, _ = self._pad_to_patch_size(H, W)
         max_real_len = text_mask.sum(dim=1).max().item()
+        real_text_lens = text_mask.sum(dim=1).tolist()
 
         hidden_gen = self.vae2llm(self.patchify(hidden_states, T, H, W))
 
@@ -1059,7 +1079,9 @@ def forward(
             if not self.sharder.is_active:
                 k_und = k_und[:, :max_real_len]
                 v_und = v_und[:, :max_real_len]
-            hidden_gen = layer(hidden_gen, k_und, v_und, freqs_gen)
+                hidden_gen = layer(hidden_gen, k_und, v_und, freqs_gen, real_text_lens)
+            else:
+                hidden_gen = layer(hidden_gen, k_und, v_und, freqs_gen)
 
         hidden_gen = self.sharder.gather(hidden_gen, dim=1, unpad_to=S_gen)