[TRTLLM-11320][refactor] Refactor VisualGenArgs API and registry - #14175
Merged
Conversation
zhenhuaw-me
added
breaking-api-change
zhenhuaw-me
changed the title
zhenhuaw-me
removed
the
breaking-api-change
zhenhuaw-me
force-pushed
the
update-api-step5
branch
from
May 15, 2026 06:19
ab10957 to
097372a
Compare
zhenhuaw-me
added
breaking-api-change
zhenhuaw-me
added a commit
to zhenhuaw-me/TensorRT-LLM
that referenced
this pull request
May 15, 2026
…efactor Apply the agreed code changes from review feedback on PR NVIDIA#14175. Public API renames: - VisualGenArgs.checkpoint_path -> VisualGenArgs.model (field + all callsites + YAML keys + test/example fixtures + docs). - diffusion_args -> visual_gen_args across executor.py, serve.py, visual_gen.py, bench/benchmark/visual_gen.py, examples, and tests. - --visual_gen_args flag dest renamed to visual_gen_args (Python identifier); --extra_visual_gen_options stays as a CLI alias only. Schema cleanup: - Remove runtime state fields (ret_steps, cutoff_steps, num_steps, _cnt) from TeaCacheConfig. TeaCacheBackend now owns warmup/cutoff bounds on the TeaCacheHook instance and derives them per-generation from num_inference_steps in refresh(). - ParallelConfig docstring trimmed; per-field description text now carries the per-axis semantics (cfg_size / ulysses_size / attn2d_size) directly. - CompilationConfig docstring trimmed; quant_config gets a description=. Pipeline registry updates: - LTX2Pipeline.hf_ids: Lightricks/LTX-Video -> Lightricks/LTX-2. - LTX2Pipeline.defaults: only text_encoder_path (default google/gemma-3-12b-it). - LTX2TwoStagesPipeline now registers its own hf_ids + defaults (text_encoder_path + spatial_upsampler_path=None + distilled_lora_path=None); docstring describes stage 1 / stage 2 directly. - WanPipeline / WanImageToVideoPipeline doc strings normalized. Internal config module: - tensorrt_llm._torch.visual_gen.config no longer re-exports public config classes through __all__; internal callers now import VisualGenArgs, PipelineComponent, AttentionConfig, SageAttentionConfig, and TeaCacheConfig from tensorrt_llm.visual_gen.args directly. - examples/visual_gen/serve/configs/ltx2.yml: drop pipeline_config block (the path now defaults from the registry entry). PipelineLoader hardening: - PipelineLoader(args, ...) now requires args; the optional/None branch was unused in production (executor always passes a real instance). Comment scrub: - Repo-wide: drop §13.x / §7.5 / §13.1 / Pre-AC-N references and "previously X / kept for one release / legacy alias" framing across files touched by this PR. Signed-off-by: Zhenhua Wang <zhenhuaw@nvidia.com>
zhenhuaw-me
changed the title
zhenhuaw-me
force-pushed
the
update-api-step5
branch
2 times, most recently
from
May 19, 2026 14:45
e0448d0 to
b8d8a6f
Compare
zhenhuaw-me
marked this pull request as ready for review
May 19, 2026 14:46
zhenhuaw-me
commented
May 19, 2026
Contributor
📝 WalkthroughWalkthroughMigrates VisualGen to a new public Pydantic args schema ( ChangesVisualGen public args API, loaders, attention, and ecosystem updates
Estimated code review effort Possibly related PRs
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md (1)
141-145:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate remaining stale
server_config.parallelreferences toserver_config.parallel_config.The README still has two old-key references, which conflicts with the refactored schema and can cause invalid config authoring.
Suggested patch
-- `hardware.gpus_per_node` must match the GPU count derived from - `server_config.parallel` +- `hardware.gpus_per_node` must match the GPU count derived from + `server_config.parallel_config` ... -3. **Verify `hardware.gpus_per_node`** matches the GPU count implied by - `server_config.parallel` (cfg × ulysses). +3. **Verify `hardware.gpus_per_node`** matches the GPU count implied by + `server_config.parallel_config` (cfg × ulysses).Also applies to: 236-238
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md` around lines 141 - 145, The README references the old key server_config.parallel; update those occurrences to server_config.parallel_config (including the block mentioning visual_gen_args_path merging and the note about hardware.gpus_per_node matching GPU counts derived from server_config.parallel), and ensure any guidance that derives GPU counts or examples use server_config.parallel_config; also confirm the note about client_configs[].generation_mode remains accurate after the key rename.tensorrt_llm/serve/scripts/benchmark_visual_gen.py (1)
573-579:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winUpdate stale CLI help text to the new flag name.
The
--num-gpushelp still says the value is inferred from--extra-visual-gen-options; it should reference--visual-gen-argsto match the renamed argument.Suggested fix
parser.add_argument( "--num-gpus", type=int, default=None, help="Number of GPUs used by the server. Overrides the value inferred " - "from --extra-visual-gen-options. Defaults to 1 if neither is given.", + "from --visual-gen-args. Defaults to 1 if neither is given.", )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/scripts/benchmark_visual_gen.py` around lines 573 - 579, Update the stale help string for the parser.add_argument call that defines the "--num-gpus" option: change the reference from "--extra-visual-gen-options" to the new flag name "--visual-gen-args" so the help text correctly matches the renamed argument (look for the parser.add_argument(... "--num-gpus" ...) entry and edit its help parameter).tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py (1)
269-275:⚠️ Potential issue | 🟠 Major | ⚡ Quick winUse
checkpoint_path=when calling_assert_pipeline_matches_hf().The helper still accepts
checkpoint_path, so this renamed keyword now raisesTypeErrorinstead of running the correctness check.💡 Proposed fix
_assert_pipeline_matches_hf( - model=WAN21_I2V_480P_PATH, + checkpoint_path=WAN21_I2V_480P_PATH, height=480, width=832, num_frames=33, guidance_scale=5.0, model_label="Wan2.1-I2V-14B-480P",🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py` around lines 269 - 275, Call the helper with the renamed keyword by replacing the model= argument with checkpoint_path= when invoking _assert_pipeline_matches_hf; specifically change the call that currently passes model=WAN21_I2V_480P_PATH to checkpoint_path=WAN21_I2V_480P_PATH (keeping the other args height, width, num_frames, guidance_scale, model_label unchanged) so the helper receives the expected parameter name.tests/unittest/_torch/visual_gen/test_attention_perf.py (1)
201-205:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift
uses_sageis only checking requested config, not the executed kernel path.
TrtllmAttentionnow carriesquant_attention_config, so this helper returnsTrueeven if runtime dispatch falls back to plain TRTLLM attention. The new assertions can therefore pass without actually covering SageAttention execution. Please assert on an observable execution signal instead (e.g. backend instrumentation/counters or a mocked dispatch point), not on the presence of the config field. Based on learnings: do not expect public LLM/attention object attributes to reveal which path executed; if a test needs to prove a specific path ran, it needs production-side instrumentation or another observable signal.Also applies to: 447-449, 960-962, 989-990, 1030-1031
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_attention_perf.py` around lines 201 - 205, The helper _is_sage_attention_enabled currently checks only for the presence of quant_attention_config on the attention object which can be set even when runtime dispatch falls back to non-Sage code; change the test to assert an observable execution signal instead: update _is_sage_attention_enabled (and the similar checks at the other locations) to verify a production-side instrumentation point or mockable dispatch hook (e.g., incremented counter, telemetry flag, or a mocked TrtllmAttention.dispatch/choose_kernel call) that is set when SageAttention actually runs; locate the runtime dispatch in TrtllmAttention (or the backend dispatch function) and either read a runtime-executed counter/flag exposed by that component or patch the dispatch to set a test-only marker, and have the helper return True only if that runtime marker was observed.tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py (1)
286-297:⚠️ Potential issue | 🟠 Major | ⚡ Quick winThis autouse env override breaks the full-pipeline tests and misses the cached fixtures.
The module now mixes full-pipeline tests with transformer-only fixtures, but this function-scoped autouse fixture sets
TLLM_VG_SKIP_COMPfor every test. That can make the correctness/batch tests load incomplete pipelines, while the module/class-scoped transformer fixtures are created before the env override is applied. Please make skipping explicit on the transformer-only fixtures (or split them into a separate module) instead of using a global function-scoped autouse.Also applies to: 393-399, 413-446
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py` around lines 286 - 297, The autouse function-scoped fixture that sets the TLLM_VG_SKIP_COMP env var is globally affecting tests (including wan21_full_pipeline) and breaking full-pipeline caching; remove or disable the global autouse behavior and instead apply the skip logic only to transformer-only fixtures: delete autouse=True from the env-override fixture (or change its scope to module/class and limit its import), and add explicit skip/conditional checks (pytest.skip or pytest.mark.skipif tied to TLLM_VG_SKIP_COMP) on the transformer-only fixtures or move those fixtures into a separate module; ensure wan21_full_pipeline and other full-pipeline fixtures (e.g., wan21_full_pipeline) are not subject to the env override so they load the complete pipeline.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py (1)
289-292: ⚡ Quick winReplace assertion with explicit runtime validation.
This guard should raise a real exception instead of using
assert, so behavior is preserved in optimized runs.As per coding guidelines: use exceptions for error handling, not assertions.Proposed fix
- if self.quant_attention_config is not None: - assert k is not None and v is not None, ( - "SageAttention requires separate Q, K, V tensors" - ) + if self.quant_attention_config is not None: + if k is None or v is None: + raise ValueError("SageAttention requires separate Q, K, V tensors")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py` around lines 289 - 292, The current guard uses an assert to check that k and v are present when self.quant_attention_config is not None; replace the assert in the SageAttention-related code with an explicit runtime exception (e.g., raise ValueError or RuntimeError) that includes a clear message like "SageAttention requires separate Q, K, V tensors" so the check still triggers in optimized/production runs; update the block around self.quant_attention_config, k, and v in trtllm.py to raise the exception instead of using assert.tests/integration/defs/visual_gen/test_visual_gen_benchmark.py (1)
193-325: QA list updates appear unnecessary for this integration-test edit.This updates existing VisualGen benchmark test wiring/flags but does not introduce a new integration test definition that needs scheduling-list onboarding.
As per coding guidelines: if integration tests under
tests/integration/defs/are changed, explicitly call out whether QA list updates are needed.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/integration/defs/visual_gen/test_visual_gen_benchmark.py` around lines 193 - 325, The PR changed integration tests under tests/integration/defs/visual_gen (functions test_online_benchmark and test_offline_benchmark) but the reviewer noted QA list updates are unnecessary; update the PR by adding a brief note to the commit message or PR description stating that these edits only adjust existing benchmark wiring/flags and do not add new integration tests so no QA scheduling-list changes are required (alternatively add a single-line comment in the test file near test_online_benchmark/test_offline_benchmark indicating "No QA list update required — only wiring/flag changes").
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/visual_gen/visual_gen_wan_i2v.py`:
- Around line 331-339: When args.enable_sage_attention is set, ensure we fail
fast if the current backend is not TRTLLM: add a check before computing
k_block_size (use the same location where _wan_needs_fine_grained_sage is
called) that verifies the backend type (e.g., inspect args.backend or the
runtime backend variable) and raise/exit with a clear error or logger.error if
it's not "trtllm"; only proceed to set attention_cfg["quant_attention_config"]
and call logger.info when the backend is TRTLLM. Reference symbols:
args.enable_sage_attention, _wan_needs_fine_grained_sage, attention_cfg, and
logger.info.
In `@examples/visual_gen/visual_gen_wan_t2v.py`:
- Around line 345-354: When args.enable_sage_attention is true, ensure the
attention backend is TRTLLM before configuring SageAttention: check
args.attention_backend (or the variable holding the backend) and if it is not
"TRTLLM" set it to "TRTLLM" (or raise/exit if you prefer strict behavior), then
proceed to compute k_block_size with _wan_needs_fine_grained_sage and populate
attention_cfg and logger.info as currently done; this guarantees SageAttention
only runs with the TRTLLM backend and avoids misleading logs.
In `@tensorrt_llm/_torch/visual_gen/__init__.py`:
- Around line 11-23: The file imports AttentionConfig and
discover_pipeline_components but doesn't use or re-export them; either remove
those two imports to fix the F401 lint, or explicitly include them in the
module's public surface by adding their names to __all__ so they become
documented re-exports. Locate the import list (contains VisualGenArgs,
TeaCacheConfig, etc.) and either delete AttentionConfig and
discover_pipeline_components from that list, or update the module's __all__ to
include "AttentionConfig" and "discover_pipeline_components" alongside existing
public symbols like "VisualGenMapping" and "VisualGenArgs" to resolve the
mismatch.
In `@tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py`:
- Around line 43-47: The registry for FluxPipeline only lists
"black-forest-labs/FLUX.1-dev" in the register_pipeline decorator's hf_ids,
which can cause strict resolver failures for the schnell variant; update the
hf_ids list on the `@register_pipeline` for "FluxPipeline" to include
"black-forest-labs/FLUX.1-schnell" (and any other FLUX.1 family variants you
want supported) so the pipeline resolves for schnell builds—modify the hf_ids
parameter in the register_pipeline decorator that registers FluxPipeline
accordingly.
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py`:
- Around line 583-595: The decorator on LTX2TwoStagesPipeline registers HF IDs
that duplicate those already used by LTX2Pipeline, causing duplicates in
VisualGen.supported_models() and ambiguous pipeline_config lookups; remove the
conflicting entries from the register_pipeline call on the LTX2TwoStagesPipeline
(i.e., update the hf_ids argument for the decorator on the LTX2TwoStagesPipeline
class) so it does not include "Lightricks/LTX-2" or
"Lightricks/LTX-Video-13B-Distilled" (keep only unique IDs or leave hf_ids
empty), ensuring the two-stage pipeline is registered with distinct identifiers
while preserving the defaults and doc strings already present in the decorator.
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py`:
- Around line 707-708: The error message raised in the LTX-2 pipeline guard
references the old "SageAttention" name but the guard checks attn_cfg and its
"quant_attention_config" attribute; update the NotImplementedError text in the
block that checks attn_cfg and getattr(attn_cfg, "quant_attention_config", None)
to mention "quant_attention_config (quantized attention)" or similar accurate
term and reference the LTX-2 pipeline so the message aligns with the renamed
config and helps debugging.
In `@tensorrt_llm/_torch/visual_gen/pipeline_loader.py`:
- Around line 88-93: PipelineLoader accepts an explicit device via self.device
but the code that materializes meta tensors uses a hardcoded "cuda" device;
update all tensor allocations inside the PipelineLoader class that currently
pass device="cuda" (e.g., calls to torch.empty_strided, torch.empty,
torch.zeros_like, torch.tensor(..., device="cuda") or similar) to use
device=self.device (or device=str(self.device) where a string is required) so
meta tensors are created on the configured device; search for allocations in
PipelineLoader (including the meta-tensor materialization logic referenced near
the later allocation block) and replace "cuda" with self.device consistently.
In `@tensorrt_llm/commands/serve.py`:
- Around line 1049-1057: In _serve_visual_gen(), the local assignment to
visual_gen_args shadows the outer variable so the RHS reads an uninitialized
local and raises UnboundLocalError; rename the parsed result (e.g.,
parsed_visual_gen_args) and use that when calling launch_visual_gen_server,
i.e., call VisualGenArgs.from_yaml(visual_gen_args) into parsed_visual_gen_args
(if visual_gen_args is not None else None) and pass parsed_visual_gen_args to
launch_visual_gen_server; keep the metadata_server_cfg logic using
parse_metadata_server_config_file(metadata_server_config_file) unchanged.
In `@tensorrt_llm/visual_gen/__init__.py`:
- Around line 22-23: Update the module docstring/comments that claim QuantConfig
is re-exported from tensorrt_llm.llmapi to reflect the actual source import
(tensorrt_llm.models.modeling_utils); locate mentions in __init__.py (the lines
around the existing docstring and the other occurrence at the later doc block
near lines 62-63) and change the referenced module name to
tensorrt_llm.models.modeling_utils so the comment matches the actual import of
QuantConfig.
In `@tensorrt_llm/visual_gen/args.py`:
- Around line 518-531: The from_yaml classmethod in VisualGenArgs currently
assumes yaml.safe_load returns a mapping and calls
config_dict.update(overrides), which will raise an AttributeError for
scalars/lists; change it to validate the loader output: after config_dict =
yaml.safe_load(f) or {}, check that config_dict is a dict/mapping
(isinstance(config_dict, dict)), and if not raise a clear TypeError/ValueError
explaining the YAML root must be a mapping; only then merge overrides and return
cls(**config_dict) so overrides handling stays the same and errors are explicit.
In `@tests/integration/defs/visual_gen/test_visual_gen_benchmark.py`:
- Around line 109-113: The current truthy check drops empty dict inputs; change
the condition that guards config-file creation from a truthiness test to an
explicit None check (i.e., use "if visual_gen_args is not None:"), so that an
empty dict is treated as a valid present config; update the block that writes
self._config_file and appends "--visual_gen_args" (the variables
visual_gen_args, self._config_file and args) accordingly.
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.py`:
- Around line 326-330: The test is passing an unexpected keyword "model" to
_logic_wan_cfg_ulysses_pvae; change the call in run_test_in_distributed to use
the correct keyword that the worker expects (checkpoint_path) so replace
model=WAN21_1_3B_PATH with checkpoint_path=WAN21_1_3B_PATH when invoking
run_test_in_distributed for the _logic_wan_cfg_ulysses_pvae test.
In `@tests/unittest/_torch/visual_gen/test_attention_perf.py`:
- Line 940: Update the pytest.parametrization decorators that pair quant_cfg and
cfg_id to use strict zipping: find the `@pytest.mark.parametrize` decorator that
currently uses zip(_QUANT_CONFIGS, _QUANT_CONFIG_IDS) for the "quant_cfg,cfg_id"
parameters (and the two other identical parametrizations in the same test
module) and change each call to zip(_QUANT_CONFIGS, _QUANT_CONFIG_IDS,
strict=True) so mismatched lengths raise an error instead of silently
truncating.
In `@tests/unittest/_torch/visual_gen/test_flux_pipeline.py`:
- Around line 66-80: The autouse function-scoped fixture
_autouse_skip_components runs too late for class-level pipelines like
TestFluxBatchGeneration.flux1_pipeline and flux2_pipeline, so move the
environment control into a broader scope or into the pipeline fixtures
themselves: either change _autouse_skip_components to a class- or session-scoped
fixture (scope="class" or "session") so it runs before class-level fixtures, or
remove it and explicitly set/clear TLLM_VG_SKIP_COMP inside the pipeline
fixtures that create flux1_pipeline and flux2_pipeline; ensure the logic still
respects _FULL_PIPELINE_TEST_CLASSES and uses SKIP_COMPONENTS to build the env
value when opting into skipping components.
In `@tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py`:
- Around line 23-31: The test imports PipelineComponent from
tensorrt_llm.visual_gen.args but that module no longer exports it; update the
import to the module that actually defines/exports PipelineComponent (e.g.,
import PipelineComponent from tensorrt_llm._torch.visual_gen.args or the
specific module where PipelineComponent is declared) so the symbol resolves at
import time; modify the import line in
tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py alongside the existing
imports of DiffusionModelConfig, LTX2_FORCE_ONE_STAGE_ENV and PipelineLoader to
reference the correct module that provides PipelineComponent.
In `@tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py`:
- Around line 248-249: The test calls to _assert_pipeline_matches_hf use the
wrong keyword name; replace the keyword argument model with checkpoint_path in
the _assert_pipeline_matches_hf(...) invocations (e.g., the call currently
passing WAN21_1_3B_PATH as model) so the helper receives checkpoint_path
correctly—update both occurrences where _assert_pipeline_matches_hf is called in
this test file.
In `@tests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.py`:
- Line 289: The test calls _assert_pipeline_matches_hf with model=WAN22_I2V_PATH
but the helper signature expects checkpoint_path; update the call to use
checkpoint_path=WAN22_I2V_PATH (or alternatively rename the helper parameter to
model) so the keyword matches the function signature; locate the call to
_assert_pipeline_matches_hf and replace model=... with checkpoint_path=... to
avoid the TypeError at runtime.
In `@tests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.py`:
- Around line 265-267: The test call to _assert_pipeline_matches_hf uses the
wrong keyword: change the argument named model to checkpoint_path so the helper
receives the expected parameter (i.e., replace model=WAN22_A14B_PATH with
checkpoint_path=WAN22_A14B_PATH in the call to _assert_pipeline_matches_hf);
verify other calls to _assert_pipeline_matches_hf use checkpoint_path as well.
- Around line 282-287: Change the _autouse_skip_components fixture to module
scope and set the env var directly via os.environ so it is in place before
module-scoped pipeline fixtures (wan22_t2v_fp8, wan22_t2v_trtllm) initialize:
make the decorator `@pytest.fixture`(scope="module", autouse=True), compute the
value with ",".join(getattr(c, "value", c) for c in _SKIP_AUX), assign it to
os.environ["TLLM_VG_SKIP_COMP"] before yielding, and restore the previous env
value (or delete the key) after the fixture yields to avoid leaking state.
In `@tests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.py`:
- Line 288: The test calls to _assert_pipeline_matches_hf use the keyword model
(e.g., model=WAN22_TI2V_5B_PATH) but the helper still expects checkpoint_path;
update the test invocations to use checkpoint_path=<constant> (replace model=
with checkpoint_path=) for the occurrences that reference WAN22_TI2V_5B_PATH
(also fix the identical call at the other location) so the keyword matches the
helper function signature.
---
Outside diff comments:
In `@tensorrt_llm/serve/scripts/benchmark_visual_gen.py`:
- Around line 573-579: Update the stale help string for the parser.add_argument
call that defines the "--num-gpus" option: change the reference from
"--extra-visual-gen-options" to the new flag name "--visual-gen-args" so the
help text correctly matches the renamed argument (look for the
parser.add_argument(... "--num-gpus" ...) entry and edit its help parameter).
In `@tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md`:
- Around line 141-145: The README references the old key server_config.parallel;
update those occurrences to server_config.parallel_config (including the block
mentioning visual_gen_args_path merging and the note about
hardware.gpus_per_node matching GPU counts derived from server_config.parallel),
and ensure any guidance that derives GPU counts or examples use
server_config.parallel_config; also confirm the note about
client_configs[].generation_mode remains accurate after the key rename.
In `@tests/unittest/_torch/visual_gen/test_attention_perf.py`:
- Around line 201-205: The helper _is_sage_attention_enabled currently checks
only for the presence of quant_attention_config on the attention object which
can be set even when runtime dispatch falls back to non-Sage code; change the
test to assert an observable execution signal instead: update
_is_sage_attention_enabled (and the similar checks at the other locations) to
verify a production-side instrumentation point or mockable dispatch hook (e.g.,
incremented counter, telemetry flag, or a mocked
TrtllmAttention.dispatch/choose_kernel call) that is set when SageAttention
actually runs; locate the runtime dispatch in TrtllmAttention (or the backend
dispatch function) and either read a runtime-executed counter/flag exposed by
that component or patch the dispatch to set a test-only marker, and have the
helper return True only if that runtime marker was observed.
In `@tests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.py`:
- Around line 269-275: Call the helper with the renamed keyword by replacing the
model= argument with checkpoint_path= when invoking _assert_pipeline_matches_hf;
specifically change the call that currently passes model=WAN21_I2V_480P_PATH to
checkpoint_path=WAN21_I2V_480P_PATH (keeping the other args height, width,
num_frames, guidance_scale, model_label unchanged) so the helper receives the
expected parameter name.
In `@tests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.py`:
- Around line 286-297: The autouse function-scoped fixture that sets the
TLLM_VG_SKIP_COMP env var is globally affecting tests (including
wan21_full_pipeline) and breaking full-pipeline caching; remove or disable the
global autouse behavior and instead apply the skip logic only to
transformer-only fixtures: delete autouse=True from the env-override fixture (or
change its scope to module/class and limit its import), and add explicit
skip/conditional checks (pytest.skip or pytest.mark.skipif tied to
TLLM_VG_SKIP_COMP) on the transformer-only fixtures or move those fixtures into
a separate module; ensure wan21_full_pipeline and other full-pipeline fixtures
(e.g., wan21_full_pipeline) are not subject to the env override so they load the
complete pipeline.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py`:
- Around line 289-292: The current guard uses an assert to check that k and v
are present when self.quant_attention_config is not None; replace the assert in
the SageAttention-related code with an explicit runtime exception (e.g., raise
ValueError or RuntimeError) that includes a clear message like "SageAttention
requires separate Q, K, V tensors" so the check still triggers in
optimized/production runs; update the block around self.quant_attention_config,
k, and v in trtllm.py to raise the exception instead of using assert.
In `@tests/integration/defs/visual_gen/test_visual_gen_benchmark.py`:
- Around line 193-325: The PR changed integration tests under
tests/integration/defs/visual_gen (functions test_online_benchmark and
test_offline_benchmark) but the reviewer noted QA list updates are unnecessary;
update the PR by adding a brief note to the commit message or PR description
stating that these edits only adjust existing benchmark wiring/flags and do not
add new integration tests so no QA scheduling-list changes are required
(alternatively add a single-line comment in the test file near
test_online_benchmark/test_offline_benchmark indicating "No QA list update
required — only wiring/flag changes").
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f776491a-e627-41d0-b2a0-7bc01e223f72
📒 Files selected for processing (78)
docs/source/models/visual-generation.mdexamples/visual_gen/configs/wan2.2-t2v-fp4-1gpu.yamlexamples/visual_gen/configs/wan2.2-t2v-fp4-4gpu.yamlexamples/visual_gen/configs/wan2.2-t2v-fp8-8gpu.yamlexamples/visual_gen/serve/benchmark_visual_gen_mgmn_distributed.shexamples/visual_gen/serve/configs/flux1.ymlexamples/visual_gen/serve/configs/flux2.ymlexamples/visual_gen/serve/configs/ltx2.ymlexamples/visual_gen/serve/configs/wan21.ymlexamples/visual_gen/serve/configs/wan22.ymlexamples/visual_gen/visual_gen_flux.pyexamples/visual_gen/visual_gen_ltx2.pyexamples/visual_gen/visual_gen_wan_i2v.pyexamples/visual_gen/visual_gen_wan_t2v.pytensorrt_llm/_torch/visual_gen/__init__.pytensorrt_llm/_torch/visual_gen/attention_backend/trtllm.pytensorrt_llm/_torch/visual_gen/attention_backend/utils.pytensorrt_llm/_torch/visual_gen/cache/teacache.pytensorrt_llm/_torch/visual_gen/checkpoints/weight_loader.pytensorrt_llm/_torch/visual_gen/config.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/mapping.pytensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.pytensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.pytensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.pytensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.pytensorrt_llm/_torch/visual_gen/pipeline.pytensorrt_llm/_torch/visual_gen/pipeline_loader.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytensorrt_llm/bench/benchmark/visual_gen.pytensorrt_llm/commands/serve.pytensorrt_llm/commands/utils.pytensorrt_llm/serve/scripts/benchmark_visual_gen.pytensorrt_llm/visual_gen/__init__.pytensorrt_llm/visual_gen/args.pytensorrt_llm/visual_gen/visual_gen.pytests/integration/defs/examples/test_visual_gen.pytests/integration/defs/perf/README_test_visual_gen_perf_sanity.mdtests/integration/defs/perf/test_visual_gen_perf_sanity.pytests/integration/defs/perf/visual_gen_perf_utils.pytests/integration/defs/visual_gen/test_visual_gen_benchmark.pytests/scripts/perf-sanity/visual_gen/flux2_blackwell.yamltests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yamltests/scripts/perf-sanity/visual_gen/wan21_t2v_14b_blackwell.yamltests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yamltests/unittest/_torch/visual_gen/conftest.pytests/unittest/_torch/visual_gen/multi_gpu/test_flux_ulysses.pytests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_sage_attention.pytests/unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.pytests/unittest/_torch/visual_gen/multi_gpu/test_wan_attn2d.pytests/unittest/_torch/visual_gen/multi_gpu/test_wan_pipeline_parallel.pytests/unittest/_torch/visual_gen/test_attention_integration.pytests/unittest/_torch/visual_gen/test_attention_perf.pytests/unittest/_torch/visual_gen/test_cache_dit.pytests/unittest/_torch/visual_gen/test_flux_attention.pytests/unittest/_torch/visual_gen/test_flux_pipeline.pytests/unittest/_torch/visual_gen/test_flux_transformer.pytests/unittest/_torch/visual_gen/test_fused_qkv.pytests/unittest/_torch/visual_gen/test_ltx2_attention.pytests/unittest/_torch/visual_gen/test_ltx2_pipeline.pytests/unittest/_torch/visual_gen/test_ltx2_transformer.pytests/unittest/_torch/visual_gen/test_model_loader.pytests/unittest/_torch/visual_gen/test_teacache.pytests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.pytests/unittest/_torch/visual_gen/test_visual_gen_args.pytests/unittest/_torch/visual_gen/test_wan21_i2v_pipeline.pytests/unittest/_torch/visual_gen/test_wan21_i2v_teacache.pytests/unittest/_torch/visual_gen/test_wan21_t2v_pipeline.pytests/unittest/_torch/visual_gen/test_wan21_t2v_teacache.pytests/unittest/_torch/visual_gen/test_wan22_i2v_pipeline.pytests/unittest/_torch/visual_gen/test_wan22_t2v_pipeline.pytests/unittest/_torch/visual_gen/test_wan22_ti2v_5b_pipeline.pytests/unittest/_torch/visual_gen/test_wan_transformer.pytests/unittest/_torch/visual_gen/test_warmup.pytests/unittest/dynamo/test_imports.py
💤 Files with no reviewable changes (1)
- tests/unittest/_torch/visual_gen/test_fused_qkv.py
Member
Author
|
/bot run |
Collaborator
|
PR_Github #49220 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49220 [ run ] completed with state
|
yufeiwu-nv
approved these changes
May 20, 2026
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
Refactor the user-facing VisualGen configuration surface for clarity and consistency. The public ``VisualGenArgs`` schema is renamed/regrouped so each sub-config is named ``<area>_config`` and the field names match the CLI flags users already type. The internal ``DiffusionModelConfig`` and ``PipelineLoader`` stay backward-compatible in spirit but are decoupled from the user surface. Public schema changes (``tensorrt_llm.visual_gen.args.VisualGenArgs``): - ``checkpoint_path`` -> ``model``; revision/HF download flow unchanged. - ``parallel`` -> ``parallel_config`` with renamed fields: ``dit_cfg_size`` -> ``cfg_size``, ``dit_ulysses_size`` -> ``ulysses_size``, ``dit_ring_size`` -> ``ring_size``, ``dit_attn2d_*`` -> ``attn2d_size`` (a ``(rows, cols)`` tuple). ``parallel_vae_size`` retained. - ``attention`` -> ``attention_config``; ``backend`` kept; SageAttention's per-block quant fields collapsed into a single ``QuantAttentionConfig`` (``qk_dtype``, ``v_dtype``, per-tensor block sizes), mirroring ``SkipSoftmaxAttentionConfig`` naming. - ``cache`` -> ``cache_config`` (still ``Union[TeaCacheConfig, CacheDiTConfig]``). - ``cuda_graph`` / ``torch_compile`` -> ``cuda_graph_config`` / ``torch_compile_config``; their ``enable_cuda_graph`` / ``enable_torch_compile`` flags collapse to ``enable``. - ``pipeline`` -> ``pipeline_config`` (a strict ``Dict[str, Any]`` validated against per-pipeline registry defaults; LTX-2 asset paths ``text_encoder_path`` / ``spatial_upsampler_path`` / ``distilled_lora_path`` live here instead of on the top-level args). - ``skip_warmup`` moved from ``VisualGenArgs`` onto ``CompilationConfig``. - ``skip_components`` removed from the public surface; kept as an internal escape hatch on ``PipelineLoader.load(skip_components=...)`` (see below) for memory-constrained unit tests. Internals: - ``PipelineLoader.load(skip_warmup, skip_components)`` is now the single entry point that accepts skip_components, forwarded to ``BasePipeline.load_standard_components`` and gated per ``PipelineComponent`` in every pipeline subclass (Wan T2V/I2V, FLUX, FLUX.2, LTX-2 single + two-stage). LTX-2 native components still honor ``"audio_vae"`` / ``"vocoder"`` / ``"connectors"`` / ``"video_encoder"`` string tokens. - ``PipelineComponent`` enum moved to ``_torch.visual_gen.pipeline_loader``; re-exported from ``_torch.visual_gen`` for tests. - ``DiffusionModelConfig`` slimmed: drops the duplicated public fields that now live on ``VisualGenArgs``; ``from_pretrained`` accepts the pre-validated ``pipeline_config`` dict from ``PipelineLoader``. - Wan / FLUX / LTX-2 pipeline files updated to read renamed config fields and to translate the new ``QuantAttentionConfig`` shape to the existing kernel ABI keywords. Examples and tests: - ``examples/visual_gen/*`` scripts updated to the new CLI surface (``--model``, ``--cfg_size``, ``--attn2d_row_size`` etc.) and to the renamed ``parallel_config`` / ``attention_config`` / ``torch_compile_config`` / ``pipeline_config`` payloads. - All unit + integration tests under ``tests/unittest/_torch/visual_gen`` and ``tests/integration/defs/examples/test_visual_gen.py`` updated to the renamed schema. Tests that need to skip heavy components now pass ``skip_components=...`` to ``PipelineLoader.load()`` (the internal kwarg), preserving the pre-refactor memory budget so that ``test_fp8_vs_bf16_memory_comparison`` and the Wan2.2 batch / fp8_cache_dit tests stay within B200's 178 GB envelope. - LPIPS helpers in ``tests/integration/defs/examples/test_visual_gen.py`` ported to the renamed schema (drop ``device``/``dtype``/top-level ``skip_warmup``; move LTX-2 asset paths into ``pipeline_config``). - ``tests/integration/defs/perf/`` perf-sanity driver + YAML serve configs updated to the renamed match keys (``parallel_config``, ``attention_config``, ``cache_config``, ``cuda_graph_config``, ``torch_compile_config``, ``pipeline_config``). Signed-off-by: Zhenhua Wang <zhenhuaw@nvidia.com>
zhenhuaw-me
force-pushed
the
update-api-step5
branch
from
May 21, 2026 14:21
8925867 to
5ad402f
Compare
Member
Author
|
/bot run --add-multi-gpu-test |
Collaborator
|
PR_Github #49718 [ run ] triggered by Bot. Commit: |
o-stoner
approved these changes
May 21, 2026
Collaborator
|
PR_Github #49718 [ run ] completed with state
|
Member
Author
|
/bot run --add-multi-gpu-test |
Collaborator
|
PR_Github #49761 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49761 [ run ] completed with state |
zhenhuaw-me
enabled auto-merge (squash)
May 22, 2026 01:40
brb-nv
approved these changes
May 22, 2026
brb-nv
left a comment
brb-nv
left a comment
Collaborator
There was a problem hiding this comment.
Approving on behalf of dynamo-devs.
yibinl-nvidia
approved these changes
May 22, 2026
yibinl-nvidia
left a comment
yibinl-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
LGTM on the LTX-2, left a comment on the doc.
Member
Author
|
/bot reuse-pipeline --comment "last CI passed and the new commit is doc update and passing local UTs" |
Collaborator
|
PR_Github #49836 [ reuse-pipeline ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49836 [ reuse-pipeline ] completed with state |
chang-l
reviewed
May 22, 2026
…back LTX2Pipeline registry docstring now matches the one-stage/two-stage auto-discovery in resolve_variant. VisualGen docs are realigned to the refactored schema (cache_config.cache_backend, cache/ subpackage, the unified _setup_cache_acceleration hook). PipelineComponent is relocated to pipeline_registry so core modules no longer reach into pipeline_loader; pipeline_loader still re-exports it for external callers. Attention quant dtype literals are renamed e4m3 -> fp8 to match the rest of TRT-LLM (kv_cache_config.dtype, indexer_k_dtype, QuantAlgo.FP8); the kernel still uses torch.float8_e4m3fn and the docstring spells that out. Signed-off-by: Zhenhua Wang <zhenhuaw@nvidia.com>
Member
Author
|
/bot run --add-multi-gpu-test |
Collaborator
|
PR_Github #49865 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49865 [ run ] completed with state
|
Member
Author
|
/bot run --add-multi-gpu-test |
Collaborator
|
PR_Github #49920 [ run ] triggered by Bot. Commit: |
7 tasks
Collaborator
|
PR_Github #49920 [ run ] completed with state
|
Member
Author
|
/bot run --add-multi-gpu-test |
Collaborator
|
PR_Github #49980 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49980 [ run ] completed with state |
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
11 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Refactor the VisualGen public API surface (TRTLLM-11320 and follow-up cleanups) so the schema, sub-config naming, pipeline registry, and CLI flags are validated in one place.
VisualGenArgsand its sub-configs totensorrt_llm.visual_gen.args. Rename master switches to drop class-name prefixes (dit_cfg_size→cfg_size, …) and add_configsuffixes to sub-config attributes (attention→attention_config, …).skip_warmup/skip_componentswithCompilationConfig.skip_warmupand component-level controls; resolve quant lazily inPipelineLoader; reshapeSageAttentionConfiginto a recipe-basedQuantAttentionConfigvalidated centrally.VisualGen.supported_models()/VisualGen.pipeline_config(model)and validated at load time. LTX-2 paths now live insidepipeline_config.--visual_gen_argsontrtllm-serve/trtllm-bench, with--extra_visual_gen_optionsretained as a back-compat alias.Example YAML configs and example scripts are updated to the new names.
Test Coverage
VisualGen unit suites under
tests/unittest/_torch/visual_gen/(args, loader, pipeline, transformer, attention, per-model) and integration tests (tests/integration/defs/examples/test_visual_gen.py,tests/integration/defs/perf/test_visual_gen_perf_sanity.py) are updated to the new schema. Verified on GB200.PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Summary by CodeRabbit
Release Notes
Configuration Schema Updates
attention_config,parallel_config,cache_config,cuda_graph_config)cfg_size,ulysses_sizereplacingdit_cfg_size,dit_ulysses_size)CLI Updates
--extra_visual_gen_optionsto--visual_gen_argsfor model serving