[TRTLLM-11410][feat] MoT World Model Support

tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py

@@ -99,7 +99,14 @@ def forward(
             f"Invalid v shape: expected [B={q.shape[0]}, H_kv, S_kv, D={self.head_dim}], got {v.shape}"
         )
 
-        return F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=self.scale)
+        return F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            is_causal=is_causal,
+            scale=self.scale,
+            enable_gqa=self.num_heads != self.num_kv_heads,
+        )
 
     @property
     def preferred_layout(self) -> AttentionTensorLayout:

tensorrt_llm/_torch/visual_gen/models/__init__.py

@@ -19,6 +19,7 @@
 
 from ..pipeline import BasePipeline
 from ..pipeline_registry import AutoPipeline, register_pipeline
+from .cosmos3 import Cosmos3OmniMoTPipeline
 from .flux import Flux2Pipeline, FluxPipeline
 from .ltx2 import LTX2Pipeline  # noqa: F401
 from .wan import WanImageToVideoPipeline, WanPipeline
@@ -30,5 +31,6 @@
     "Flux2Pipeline",
     "WanPipeline",
     "WanImageToVideoPipeline",
+    "Cosmos3OmniMoTPipeline",
     "register_pipeline",
 ]

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

@@ -0,0 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .pipeline_cosmos3 import Cosmos3OmniMoTPipeline
+
+__all__ = ["Cosmos3OmniMoTPipeline"]

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

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Per-model default generation parameters for Cosmos3 pipelines.
+
+Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths.
+"""
+
+from typing import Dict
+
+from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema
+
+# ---------------------------------------------------------------------------
+# Constant tables
+# ---------------------------------------------------------------------------
+
+COSMOS3_720P_PARAMS = {
+    "height": 720,
+    "width": 1280,
+    "num_inference_steps": 35,
+    "guidance_scale": 6.0,
+    "max_sequence_length": 1024,
+    "num_frames": 189,
+    "frame_rate": 24.0,
+}
+
+COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = {
+    "use_duration_template": ExtraParamSchema(
+        type="bool",
+        default=True,
+        description="Whether to use the duration template.",
+    ),
+    "use_resolution_template": ExtraParamSchema(
+        type="bool",
+        default=True,
+        description="Whether to use the resolution template.",
+    ),
+    "use_system_prompt": ExtraParamSchema(
+        type="bool",
+        default=False,
+        description="Whether to use the system prompt.",
+    ),
+    "use_guardrails": ExtraParamSchema(
+        type="bool",
+        default=True,
+        description="Whether to use the guardrails.",
+    ),
+}

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

@@ -0,0 +1,67 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from tensorrt_llm.logger import logger
+
+GUARDRAIL_HF_REPO = "nvidia/Cosmos-1.0-Guardrail"
+GUARDRAIL_REVISION = "cf03c0395fac8c4de386c0bdab12cc4fc8d66362"
+
+
+def download_guardrail_checkpoint() -> str:
+    from huggingface_hub import snapshot_download
+    from huggingface_hub.errors import GatedRepoError
+
+    try:
+        return snapshot_download(
+            GUARDRAIL_HF_REPO,
+            revision=GUARDRAIL_REVISION,
+            local_files_only=True,
+        )
+    except FileNotFoundError:
+        logger.warning(f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO}")
+        try:
+            return snapshot_download(
+                GUARDRAIL_HF_REPO,
+                revision=GUARDRAIL_REVISION,
+            )
+        except GatedRepoError:
+            raise ValueError(
+                "Cosmos Guardrail checkpoint not found. "
+                "Please ensure "
+                "a) you have accepted the terms of use (https://huggingface.co/nvidia/Cosmos-1.0-Guardrail) "
+                "b) you have set a valid HF_TOKEN environment variable"
+            )
+
+
+def check_video_safety(video_tensor: torch.Tensor, safety_checker: Any) -> torch.Tensor | None:
+    v = video_tensor.detach().cpu()
+    was_batched = v.dim() == 5
+    if was_batched:
+        v = v[0]
+    frames_np = v.numpy()
+    frames_np = safety_checker.check_video_safety(frames_np)
+    if frames_np is None:
+        return None
+
+    result = torch.from_numpy(frames_np)
+    if was_batched:
+        result = result.unsqueeze(0)
+    return result.to(video_tensor.device)