Add TensorRT-LLM distilled 4-step guide

[ishovkun]

Summary

• add a TensorRT-LLM walkthrough for the two published DMD2-distilled Cosmos3-Super students, nvidia/Cosmos3-Super-Text2Image-4Step (text-to-image) and nvidia/Cosmos3-Super-Image2Video-4Step (image-to-video), each on a single GPU
• document the distilled request contract: the checkpoint's scheduler config locks the schedule to four stochastic (SDE) steps and classifier-free guidance is baked into the weights, so requests omit num_inference_steps and guidance_scale entirely and the server fills both in; a conflicting value is rejected, not clamped
• leave use_system_prompt unset in every request so the image-to-video student's checkpoint-declared default_use_system_prompt: true applies
• add both trtllm-serve launch commands to the shared setup guide, and index the notebook from the root README
• separately: reword two guardrail blocklist false positives in the shared audiovisual negative prompts, without which no image-to-video request can run at all (see below)

This is the follow-up #310 deferred ("keep the four-step distilled T2I/I2V examples out of this PR pending separate output-quality follow-up"). The examples match the assets every other backend already uses for its distilled sections — robot_draping for text-to-image and car_driving for image-to-video — so the cookbook surface is consistent across Cosmos Framework, Diffusers, vLLM-Omni, SGLang, and now TensorRT-LLM.

The shipped scope is text-to-image and image-to-video only. These students are task-specialized and no other backend advertises text-to-video, video-to-video, or synchronized audio for them either.

The guardrail fix, and why it is in this PR

The shared audiovisual negative prompts contained two substrings that the cosmos_guardrail 0.3.0 Blocklist censors. Any request carrying them fails with HTTP 500 and Text guardrail blocked prompt before any denoising runs:

Text	Censored word	Replacement
color bleeding between elements	bleeding	color smearing between elements
The scene feels lifeless and sterile	lifeless	The scene feels inert and sterile

Neither describes anything the blocklist is meant to catch; both are ordinary image-artifact vocabulary.

The replacement wording is not invented here. TensorRT-LLM already ships a corrected copy of this exact prompt as examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json, and that is the default every Cosmos3 request falls back to when it sends no negative prompt of its own. Upstream's copy and the cookbook's differ in exactly four leaves, and all four are these two substrings — so the cookbook asset is simply stale. Both files are now identical to the TensorRT-LLM default. Kept as separate commits so they can be split out if preferred.

Text-to-image was unaffected because that path sends no negative prompt — which is why only the image-to-video example hit it.

Three related issues found while diagnosing this, not changed here:

1. The two files are shared by the Diffusers, vLLM-Omni, SGLang and base TensorRT-LLM notebooks, so this also unblocks their image-to-video / text-to-video / video-to-video examples.
2. assets/prompts/image2video/humanoid_robot.json is blocked by a different rule (partial match, not the word blocklist). That is the prompt the merged run_with_trt_llm.ipynb uses for i2v_nano and i2v_super. Left alone deliberately: rewording a positive prompt changes generated content.
3. transfer/assets/negative_prompt.json carries the same two substrings. That file belongs to Add TensorRT-LLM Transfer and Action guides #310, which now rewords both, so it is covered there rather than duplicated here.

Validation

GPU inference

Both notebook sections were executed as written with jupyter nbconvert --to notebook --execute against real trtllm-serve VisualGen servers — the shipped notebook's own cells, not a re-implementation. TensorRT-LLM at 505a02ce04, whose tensorrt_llm/_torch/visual_gen, tensorrt_llm/visual_gen and examples/visual_gen trees are byte-identical to main (2dbab44d4f); the only delta is test files. Hardware: 1 x NVIDIA B300.

Text-to-image — Cosmos3-Super-Text2Image-4Step with configs/cosmos3-t2i-1gpu.yaml:

Distilled Cosmos3 checkpoint: fixed 4-step schedule
  [1.0, 0.9375, 0.8333333333333334, 0.625], classifier-free guidance baked in.
Running warmup for Cosmos3OmniMoTPipeline: 1 shape(s) [1024x1024x1], 4 steps
Denoising: 4 steps, guidance=1.0
POST /v1/videos/generations HTTP/1.1" 200 OK

Artifact: 57,617 bytes, decoded with PyAV as H.264 1024x1024, 1 frame.

Image-to-video — Cosmos3-Super-Image2Video-4Step, no config file:

Running warmup for Cosmos3OmniMoTPipeline: 1 shape(s) [720x1280x189], 4 steps
Denoising: 4 steps, guidance=1.0
POST /v1/videos/generations HTTP/1.1" 200 OK

Artifact: 4,890,310 bytes, decoded as H.264 1280x720, 24 fps, 189 frames. The warmup shape confirms the claim that this checkpoint deploys at the default omni shape and needs no config file.

Negative control (both servers): a request that sets guidance_scale=6.0 is rejected rather than clamped —

HTTP 400: This is a distilled Cosmos3 checkpoint; classifier-free guidance is
baked into the weights. guidance_scale must be 1.0 or left unset (got 6.0).

Guardrail fix verified directly against cosmos_guardrail 0.3.0's Blocklist: both reworded negative prompts and both positive prompts used by the notebook now report safe=True. The trigger words were located by diffing the censored output against the original, not guessed. Re-verified after adopting upstream's wording: bleeding and lifeless are present in NVIDIA's blocklist/exact_match word sets (383 and 1430 entries) while smearing and inert are not.

Visual review: the text-to-image frame renders the prompt's robot-arm-draping-satin scene coherently. The image-to-video clip follows its temporal_caption — coastal dashcam view, rockfall beginning a few seconds in, boulders and dust building until the road is partially blocked — with no flicker or geometry snapping across the 189 frames.

Local validation

• parsed the new notebook as JSON
• validated it against the nbformat schema with nbformat==5.11.1 (the CI gate)
• linted it with ruff==0.16.4 check --select=E9,F63,F7,F82 (the CI gate)
• verified unique cell IDs, empty outputs, null execution counts, and that every code cell parses
• re-validated both edited JSON assets with json.tool
• ran git diff --check

Source contract audit

The request contract is read from TensorRT-LLM main (tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py):

• load_scheduler selects FlowMatchEulerDiscreteScheduler for distilled checkpoints and UniPCMultistepScheduler for base ones; an unknown declaration is a load-time error
• Cosmos3SamplingPolicy.validate_request rejects any num_inference_steps other than len(t_list) and any guidance_scale other than 1.0
• generation_default_overrides supplies both values when the request leaves them unset
• set_flow_shift is a structural no-op for distilled checkpoints, so flow_shift does not apply
• both checkpoints declare fixed_step_sampler_config.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625] with stochastic_sampling: true

Dependencies and review state

• TensorRT-LLM support is already merged: [None][feat] Support the DMD2-distilled Cosmos3-Super-Text2Image-4Step checkpoint TensorRT-LLM#16563 (distilled text-to-image) and [None][feat] Support the DMD2-distilled Cosmos3 4-step image-to-video checkpoint TensorRT-LLM#16690 (distilled image-to-video).
• No dependency on Add TensorRT-LLM Transfer and Action guides #310; the two PRs touch different notebooks. If Add TensorRT-LLM Transfer and Action guides #310 merges first there is no expected conflict.

[ConstBob]

Thanks. The distilled 4-step request contract matches TensorRT-LLM (omit num_inference_steps / guidance_scale, reject mismatches, leave use_system_prompt unset; T2I as 1024x1024 one-frame /v1/videos/generations, I2V as 720p x 189 with no config yaml).

Static checks on 0e1a27a: notebook JSON, unique cell IDs, empty outputs, Python parse, assets (robot_draping, car_driving), and the shared AV negative prompts no longer contain bleeding / lifeless.

Did not re-run GPU inference (author already executed both sections as written on 1x B300).

Nits, non-blocking:

• transfer/assets/negative_prompt.json still has both trigger words on this base; called out in the PR.
• #354 edits the same two audiovisual neg_prompt.json files, so the later merge may need a trivial conflict resolve.

LGTM.

[skasetty-Siri]

Please fix the default TensorRT-LLM root in run_distilled_with_trt_llm.ipynb, in the “Start the Server” section. The instructions say to run from the TensorRT-LLM checkout, but ${TRTLLM_ROOT:-$PWD/TensorRT-LLM} adds another TensorRT-LLM directory when the variable is unset. The config file cannot be found.

Please use ${TRTLLM_ROOT:-$PWD} for the stated working directory, or update the instructions consistently, and verify from a fresh shell with TRTLLM_ROOT unset.

This path issue was reproduced during assisted review. GPU inference was not independently rerun.

[ishovkun]

@skasetty-Siri re: your review — confirmed, and reproduced before fixing. Inside checkout, TRTLLM_ROOT unset:

PWD         = .../code/TensorRT-LLM
TRTLLM_ROOT = .../code/TensorRT-LLM/TensorRT-LLM
config      = .../TensorRT-LLM/TensorRT-LLM/examples/visual_gen/configs/cosmos3-t2i-1gpu.yaml
RESULT: NOT FOUND

Origin: copied from run_with_trt_llm.ipynb. Same default there, but that file also says "From the repository root" — one level above a sibling clone, where the default is correct. This notebook kept the default and dropped that line, leaving only "from that checkout". Contradiction.

Took your first option, ${TRTLLM_ROOT:-$PWD} (1c625cd):

• working directory now stated in prose, not implied
• export hoisted out of the per-model blocks → one definition, not one per server
• added up-front path check. trtllm-serve validates --visual_gen_args only after startup → bad place to learn the root was wrong → why this got past me

Same default was wrong in two more places. Both fixed:

• cookbooks/cosmos3/README.md — setup block does git clone → cd TensorRT-LLM → docker run --workdir "$PWD", so already inside the checkout at the export. Fix also repairs the pre-existing Nano and Super serve commands.
• run_with_trt_llm.ipynb (314f856) — broken default in both blocks, same contradiction.

No $PWD/TensorRT-LLM left in the repo. Every export TRTLLM_ROOT= line is identical now.

Verified as asked. Fresh shell, TRTLLM_ROOT unset, cwd = checkout. Extracted the export lines and every --visual_gen_args path from the shipped files rather than retyping them:

Source	Config paths	Result
run_distilled_with_trt_llm.ipynb	1 (t2i-1gpu)	all resolve
run_with_trt_llm.ipynb	2 (nano-1gpu, super-4gpu)	all resolve
cookbooks/cosmos3/README.md	3 (nano-1gpu, super-4gpu, t2i-1gpu)	all resolve

The fix also adds a guard so a wrong root is caught before trtllm-serve starts, rather than surfacing mid-startup.

[skasetty-Siri]

Rechecked 314f856. The startup-path finding is resolved in the distilled notebook, shared setup, and base audiovisual notebook. Assisted offline checks and latest notebook CI passed. No remaining blocking findings from my review. GPU inference was not independently rerun.

[MaciejBalaNV]

Approving after a check from @skasetty-Siri