Add Cosmos3 Nano/Super Reasoner examples with TensorRT-LLM, fix vLLM reasoner examples

README.md

@@ -31,6 +31,7 @@
     - [Generator with SGLang](#generator-with-sglang)
     - [Reasoner with Transformers](#reasoner-with-transformers)
     - [Reasoner with vLLM](#reasoner-with-vllm)
+    - [Reasoner with TensorRT-LLM](#reasoner-with-tensorrt-llm)
     - [Reasoner with NIM](#reasoner-with-nim)
   - [Troubleshooting](#troubleshooting)
     - [Which CUDA version should I use?](#which-cuda-version-should-i-use)
@@ -69,7 +70,7 @@ Cosmos 3 exposes two runtime surfaces:
 - **World understanding:** Analyze videos and images for captions, temporal events, next actions, spatial grounding, physical plausibility, and causal outcomes.
 - **World generation:** Produce images, videos, synchronized sound, and action-conditioned rollouts from text, image, video, or action inputs.
 - **Action modeling:** Predict policy actions, inverse dynamics, and forward dynamics for robotics, camera motion, egocentric motion, and autonomous-driving settings.
-- **Research and production paths:** Use Diffusers and Transformers for Python-first development, vLLM-Omni, vLLM, or SGLang for OpenAI-compatible serving, and NIM containers for turnkey Reasoner serving or Generator deployment for text-to-video and image-to-video generation.
+- **Research and production paths:** Use Diffusers and Transformers for Python-first development, vLLM-Omni, vLLM, TensorRT-LLM, or SGLang for OpenAI-compatible serving, and NIM containers for turnkey Reasoner serving or Generator deployment for text-to-video and image-to-video generation.
 - **Post-training recipes:** Adapt vision, action, and reasoner workflows with Cosmos Framework training recipes and task-specific evaluation [Coming Soon].
 
 ### Model Architecture
@@ -922,7 +923,7 @@ Then reuse the `model.generate` and `batch_decode` block from the image example.
 For `nvidia/Cosmos3-Super`, change `model_id` to `nvidia/Cosmos3-Super`.
 `device_map="auto"` can shard the model across multiple GPUs when Accelerate is
 installed. For an OpenAI-compatible server, use
-[Reasoner with vLLM](#reasoner-with-vllm) or [Reasoner with NIM](#reasoner-with-nim).
+[Reasoner with vLLM](#reasoner-with-vllm), [Reasoner with TensorRT-LLM](#reasoner-with-tensorrt-llm), or [Reasoner with NIM](#reasoner-with-nim).
 
 </details>
 
@@ -969,6 +970,25 @@ Configuration notes:
 
 </details>
 
+#### Reasoner with TensorRT-LLM
+
+<details>
+<summary>Expand TensorRT-LLM Reasoner setup</summary>
+
+Use the TensorRT-LLM release container or install it from source to serve
+Cosmos3-Nano on one GPU or Cosmos3-Super on four GPUs through the
+OpenAI-compatible `/v1/chat/completions` API. The server accepts image and video
+inputs and loads only the Reasoner path for these requests.
+
+See the shared
+[TensorRT-LLM Reasoner setup](cookbooks/cosmos3/README.md#tensorrt-llm-reasoner) for
+the container or local installation options, Nano and Super launch commands,
+and server endpoint. The
+[Reasoner notebook](cookbooks/cosmos3/reasoner/run_with_tensorrt_llm.ipynb)
+contains the complete image and video request walkthrough.
+
+</details>
+
 #### Reasoner with NIM
 
 <details>
@@ -1109,7 +1129,7 @@ The Cosmos Framework requires `uv >= 0.11.3` (enforced via its `pyproject.toml`)
 | Generator broader production/API serving | vLLM-Omni / SGLang | API path for image, video, sound, and action outputs |
 | Generator turnkey deployment | NIM | Prebuilt NGC container for T2V/I2V video generation only; uses `/v1/infer` and returns JSON `b64_video` |
 | Reasoner research or model development | Transformers | Python-first path for prompts, processors, and model behavior |
-| Reasoner production inference | vLLM | OpenAI-compatible endpoint for text outputs from text and vision inputs |
+| Reasoner production inference | vLLM / TensorRT-LLM | OpenAI-compatible endpoints for text outputs from image and video inputs |
 | Reasoner turnkey deployment | NIM | Prebuilt, optimized OpenAI-compatible container — no vLLM/CUDA setup |
 | Runnable setup, training, or evaluation | Cosmos Framework | Full workflow docs for setup, inference, omni-model training, and evaluation |
 
@@ -1135,6 +1155,7 @@ We are building examples that show Cosmos 3 Super/Nano/Edge capabilities end to
 | Transfer with vLLM-Omni | Generator | Video transfer: edge, blur, depth, segmentation, and world-scenario controls with captions, against an OpenAI-compatible vLLM-Omni server. | [Notebook](cookbooks/cosmos3/generator/transfer/run_video_transfer_with_vllm_omni.ipynb) | [![Render with nbviewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.org/github/nvidia/cosmos/blob/main/cookbooks/cosmos3/generator/transfer/run_video_transfer_with_vllm_omni.ipynb) |
 | Reasoner with Cosmos Framework | Reasoner | Text and image reasoning: detailed captioning, robot task planning, 2D grounding, describe-anything, and action-trajectory prompts, through the `cosmos_framework.scripts.inference` entrypoint. | [Notebook](cookbooks/cosmos3/reasoner/run_with_cosmos_framework.ipynb) | [![Render with nbviewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.org/github/nvidia/cosmos/blob/main/cookbooks/cosmos3/reasoner/run_with_cosmos_framework.ipynb) |
 | Reasoner with vLLM | Reasoner | Image and video reasoning: captioning, temporal localization, embodied reasoning, common-sense reasoning, 2D grounding, describe-anything, action CoT, driving scenes, physical-plausibility, and situation understanding, against an OpenAI-compatible vLLM server (Cosmos3-Super on 4 GPUs by default; switch to Nano per the cookbook README). | [Notebook](cookbooks/cosmos3/reasoner/run_with_vllm.ipynb) | [![Render with nbviewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.org/github/nvidia/cosmos/blob/main/cookbooks/cosmos3/reasoner/run_with_vllm.ipynb) |
+| Reasoner with TensorRT-LLM | Reasoner | Image and video reasoning through an OpenAI-compatible TensorRT-LLM server, with one-GPU Cosmos3-Nano and four-GPU Cosmos3-Super launch options. | [Notebook](cookbooks/cosmos3/reasoner/run_with_tensorrt_llm.ipynb) | [![Render with nbviewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.org/github/nvidia/cosmos/blob/main/cookbooks/cosmos3/reasoner/run_with_tensorrt_llm.ipynb) |
 | Reasoner with NIM | Reasoner | The same image and video reasoning examples as the vLLM notebook, run against the prebuilt, OpenAI-compatible [Cosmos 3 Reasoner NIM](https://catalog.ngc.nvidia.com/orgs/nim/teams/nvidia/containers/cosmos3-reasoner) container; local media is sent as base64 data URIs. | [Notebook](cookbooks/cosmos3/reasoner/run_with_nim.ipynb) | [![Render with nbviewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.org/github/nvidia/cosmos/blob/main/cookbooks/cosmos3/reasoner/run_with_nim.ipynb) |
 
 ### Inference Benchmarks

cookbooks/cosmos3/README.md

@@ -8,7 +8,8 @@ backend you want to run and follow that one section.
 | --- | --- | --- |
 | [Cosmos Framework](#cosmos-framework) | Native PyTorch inference, launched with `torchrun` | Reasoner, Generator (Audiovisual, Action, **Transfer**) |
 | [Diffusers](#diffusers) | Direct generation with `Cosmos3OmniPipeline` | Generator (Audiovisual) |
-| [TensorRT-LLM](#tensorrt-llm) | OpenAI-compatible VisualGen server (image/video generation) | Generator (Audiovisual) |
+| [TensorRT-LLM Generator](#tensorrt-llm-generator) | OpenAI-compatible VisualGen server (image/video generation) | Generator (Audiovisual) |
+| [TensorRT-LLM Reasoner](#tensorrt-llm-reasoner) | OpenAI-compatible image/video reasoning server | Reasoner |
 | [Transformers](#transformers) | Hugging Face Transformers inference | Reasoner |
 | [vLLM](#vllm) | OpenAI-compatible reasoning server (image/video understanding) | Reasoner |
 | [vLLM-Omni](#vllm-omni) | OpenAI-compatible generation server (image/video/audio/action/transfer) | Generator (Audiovisual, Action, **Transfer**) |
@@ -165,7 +166,7 @@ uv pip install --torch-backend=cu130 \
   transformers
 ```
 
-## TensorRT-LLM
+## TensorRT-LLM Generator
 
 OpenAI-compatible **VisualGen** server for Generator audiovisual text-to-image,
 text-to-video, and image-to-video examples. Cosmos3 support was added in TensorRT-LLM PR
@@ -252,6 +253,54 @@ so use a TensorRT-LLM build that includes the Cosmos3 VisualGen API schema.
 The notebook sets request-level `max_sequence_length=2048` for longer structured
 JSON prompts.
 
+## TensorRT-LLM Reasoner
+
+OpenAI-compatible **reasoning** server for image and video understanding. Run
+TensorRT-LLM in the prebuilt
+[`nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc22`](https://catalog.ngc.nvidia.com/orgs/nvidia/tensorrt-llm/containers/release)
+container, or follow the
+[TensorRT-LLM Generator source-build instructions](#tensorrt-llm-generator) to
+install TensorRT-LLM from source.
+Authenticate with Hugging Face before loading gated checkpoints.
+
+Install headless OpenCV in the environment that runs the server. Cosmos3
+Reasoner uses it to process vision inputs:
+
+```bash
+python -m pip install opencv-python-headless
+```
+
+### Start the server
+
+Run one of these commands inside the release container or activated local
+environment.
+
+**Cosmos3-Nano** (single GPU, port 8001):
+
+```bash
+CUDA_VISIBLE_DEVICES=0 \
+trtllm-serve nvidia/Cosmos3-Nano \
+  --host 0.0.0.0 \
+  --port 8001 \
+  --max_num_tokens 32768
+```
+
+**Cosmos3-Super** (four GPUs, port 8001):
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1,2,3 \
+trtllm-serve nvidia/Cosmos3-Super \
+  --host 0.0.0.0 \
+  --port 8001 \
+  --tensor_parallel_size 4 \
+  --max_num_tokens 32768
+```
+
+The server exposes `/health` and the OpenAI-compatible API at
+`http://localhost:8001/v1`. See the
+[Reasoner TensorRT-LLM notebook](reasoner/run_with_tensorrt_llm.ipynb) for image
+and video requests.
+
 ## Transformers
 
 Local Python inference for the Cosmos3 Reasoner. This backend uses the

cookbooks/cosmos3/reasoner/README.md

@@ -16,17 +16,17 @@ how much GPU memory you pay depends on the backend:
   the generator: video VAE, audio tokenizer, and diffusion expert). It therefore
   uses the most memory — even for text-only reasoner inference. `model_mode=reasoner`
   only changes the execution path; it does **not** unload the generator weights.
-- The **vLLM**, **Transformers**, and **NIM** backends load **only the Reasoner (VLM)
-  weights**, so they use much less memory.
+- The **vLLM**, **TensorRT-LLM**, **Transformers**, and **NIM** backends load
+  **only the Reasoner (VLM) weights**, so they use much less memory.
 
 **If you need to save GPU memory for reasoner-only inference, use vLLM,
-Transformers, or NIM.** Reach for the Cosmos Framework backend only when you also
-need generation.
+TensorRT-LLM, Transformers, or NIM.** Reach for the Cosmos Framework backend only
+when you also need generation.
 
 For example, on a single GPU the **Cosmos3-Nano** Reasoner needs roughly **~34 GB**
-with the Cosmos Framework backend versus **~16–17 GB** with vLLM / Transformers / NIM
-— about half. Larger tiers (e.g. Cosmos3-Super) use more in absolute terms, but the
-same full-omni-vs-reasoner-only gap applies.
+with the Cosmos Framework backend versus **~16–17 GB** with vLLM / TensorRT-LLM /
+Transformers / NIM — about half. Larger tiers (e.g. Cosmos3-Super) use more in
+absolute terms, but the same full-omni-vs-reasoner-only gap applies.
 
 ## Reasoner Prompt Guide
 
@@ -303,3 +303,56 @@ and loads `Cosmos3OmniForConditionalGeneration` in process. A small
 then runs the image and video examples shown above. To scale from **Nano** to
 **Super**, change only `model_id` in the load cell and re-run; `device_map="auto"`
 shards Super across multiple GPUs.
+
+## Run with TensorRT-LLM
+
+### Quickstart
+
+Set up the environment and start the server with the shared
+[TensorRT-LLM Reasoner setup](../README.md#tensorrt-llm-reasoner). It covers the
+prebuilt release container or a local source installation, Hugging Face
+authentication, and launch commands for **Cosmos3-Nano** on one GPU or
+**Cosmos3-Super** on four GPUs.
+
+The [`run_with_tensorrt_llm.ipynb`](./run_with_tensorrt_llm.ipynb) notebook
+publishes the selected server on host port **8001** and waits for `/health`
+before running the request examples. Point standalone clients at
+`http://localhost:8001/v1`.
+
+Once the server is ready, query it with the OpenAI client:
+
+```python
+from pathlib import Path
+import openai
+
+image_path = Path("assets/robot_153.jpg").resolve()
+image_url = image_path.as_uri()
+
+client = openai.OpenAI(api_key="EMPTY", base_url="http://localhost:8001/v1")
+
+response = client.chat.completions.create(
+    model=client.models.list().data[0].id,
+    messages=[
+        {
+            "role": "user",
+            "content": [
+                {"type": "image_url", "image_url": {"url": image_url}},
+                {"type": "text", "text": "Caption the image in detail."},
+            ],
+        }
+    ],
+    max_tokens=4096,
+    seed=0,
+)
+
+print(response.choices[0].message.content)
+```
+
+### Notebook walkthrough
+
+[`run_with_tensorrt_llm.ipynb`](./run_with_tensorrt_llm.ipynb) launches either
+**Cosmos3-Nano** or **Cosmos3-Super**, waits for the server to become healthy,
+and walks through image and video examples covering detailed captioning,
+temporal localization, embodied reasoning, common-sense reasoning, 2D
+grounding, describe-anything, action CoT trajectories, driving scenes,
+physical plausibility, and situation understanding.