[None][feat] add Qwen-Image-Layered FP8 support - #16889
Draft
yumin066 wants to merge 2 commits into
Draft
Conversation
Signed-off-by: Min Yu <minyu@nvidia.com>
Add the FP8 block-scale and SageAttention FP8 preset for Qwen-Image-Layered. Route masked joint attention through the TRTLLM backend only when selected, while preserving the existing VANILLA path. Avoid two-step MXFP8 weight quantization by prequantizing Qwen-Image-Layered FP8 block-scale Linear weights directly with E8M0-compatible scales before the existing Qwen weight-loading path. This keeps the shared quantization APIs and Qwen-Image behavior unchanged while preventing additional QDQ error during the required post-load E8M0 repack. On B200 with the Qwen-Image-Layered FP8 block-scale and SageAttention FP8 preset, LPIPS against the official Layered golden improved from 0.322196 with two-step quantization to 0.035642 with direct one-step E8M0 quantization, passing the 0.05 threshold. Add unit and LPIPS integration coverage for the documented preset. Signed-off-by: Min Yu <minyu@nvidia.com>
Collaborator
Author
|
@coderabbitai full review |
Contributor
✅ Action performedFull review finished. |
Collaborator
Author
|
/bot run |
Contributor
WalkthroughChangesAdds Qwen-Image-Layered support across the PyTorch VisualGen runtime, including layered transformer inference, FP8 configuration, checkpoint routing, CLI examples, documentation, unit tests, parity tests, and LPIPS integration coverage. Qwen-Image-Layered support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant VisualGen
participant QwenImageLayeredPipeline
participant Qwen2VLProcessor
participant QwenImageLayeredTransformer2DModel
participant VAE
VisualGen->>QwenImageLayeredPipeline: generate image and prompt
QwenImageLayeredPipeline->>Qwen2VLProcessor: caption empty prompt input
QwenImageLayeredPipeline->>VAE: encode conditioning image
QwenImageLayeredPipeline->>QwenImageLayeredTransformer2DModel: denoise packed layered latents
QwenImageLayeredPipeline->>VAE: decode final layered latents
QwenImageLayeredPipeline-->>VisualGen: return tiled RGBA image grid
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![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: 6
🧹 Nitpick comments (5)
tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py (3)
743-761: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd a Google-style docstring for this public entrypoint.
forwardtakes 14 parameters, several with non-obvious semantics (resolutionas a bucket rather than an output size,layersexcluding the conditioning frame,latentsneeding to be pre-packed). As per coding guidelines: "Prefer docstrings for external interfaces, use Google-style docstrings, document public function arguments."🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py` around lines 743 - 761, Add a Google-style docstring to the public Qwen image layered pipeline forward method, documenting its purpose, return value, and all parameters, including that resolution selects a bucket, layers excludes the conditioning frame, and latents must be pre-packed. Keep the implementation and behavior unchanged.Source: Coding guidelines
727-741: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtra-param defaults are duplicated from
extra_param_specs.
layers=4,resolution=640,cfg_normalize=False, anduse_en_prompt=Falseare hardcoded here and again inextra_param_specs(Lines 219-239). Changing a schema default silently leavesinferon the old value. Sourcing them from the spec keeps the two in lockstep.♻️ Proposed refactor
+ specs = self.extra_param_specs + + def _extra(name): + return extra.get(name, specs[name].default) + return self.forward( image=req.params.image, prompt=prompts, negative_prompt=negative, height=req.params.height, width=req.params.width, true_cfg_scale=req.params.guidance_scale, - layers=extra.get("layers", 4), + layers=_extra("layers"), num_inference_steps=req.params.num_inference_steps, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, - resolution=extra.get("resolution", 640), - cfg_normalize=extra.get("cfg_normalize", False), - use_en_prompt=extra.get("use_en_prompt", False), + resolution=_extra("resolution"), + cfg_normalize=_extra("cfg_normalize"), + use_en_prompt=_extra("use_en_prompt"), )🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py` around lines 727 - 741, Update the request-to-forward mapping in the infer flow to source defaults for layers, resolution, cfg_normalize, and use_en_prompt from the corresponding entries in extra_param_specs instead of hardcoding them. Preserve values explicitly provided in extra and keep the schema and inference defaults synchronized.
511-550: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
_pack_layered_latentsand_unpack_layered_latentsare not symmetric APIs.Two mismatches make these easy to misuse: pack takes the literal frame count (callers pass
1at Line 620 andlayers + 1at Line 635) while unpack hardcodeslayers + 1internally; and pack takes latent-spaceheight/widthwhile unpack takes pixel-space dimensions it divides down itself. Making both take the same frame-count convention and the same coordinate space would remove a standing footgun.🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py` around lines 511 - 550, Make _pack_layered_latents and _unpack_layered_latents symmetric by using the same frame-count convention and dimension coordinate space in both APIs. Update their callers, including the paths currently passing 1 and layers + 1, so the pack/unpack round trip preserves the intended number of frames without internal layers + 1 assumptions or duplicated height/width conversion.tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py (2)
405-446: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
forwardis a near-verbatim copy of the parent implementation.Comparing with
QwenImageTransformer2DModel.forward(tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py:1001-1073), the only delta is the thirdadditional_t_condargument toself.time_text_embed. Since that class already accepts the argument, the parent could take an optionaladditional_t_condand this override could disappear — otherwise every future fix to the base mask/RoPE logic has to be applied twice.🤖 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/models/qwen_image_layered/transformer_qwen_image_layered.py` around lines 405 - 446, Refactor the duplicated forward path by adding optional additional_t_cond support to QwenImageTransformer2DModel.forward and passing it to time_text_embed, then remove the near-identical forward override from the layered transformer so it inherits the parent implementation. Preserve the existing mask, rotary embedding, block execution, and return behavior.
217-250: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valuePer-sample Python loop serializes attention for
batch_size > 1.Each row runs a separate
super().forward(), so a batch of N costs N sequential attention launches plus N host-side syncs fromtorch.nonzero. When all rows share the same valid text length (the common case, since_encode_promptpads to a sharedmax_len), you can take the fast path of compacting once and running a single batched call.🤖 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/models/qwen_image_layered/transformer_qwen_image_layered.py` around lines 217 - 250, Update the batched forward logic around the per-sample loop in the transformer’s forward method to add a fast path when all rows share the same valid text length: compact the text states and rotary embeddings once, invoke super().forward() once for the full batch, and restore outputs to the original text layout. Preserve the existing per-sample path for varying valid lengths and maintain timestep, mask, dtype, device, and bias behavior.
🤖 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 `@docs/source/models/supported-models.md`:
- Around line 185-186: Rename the Qwen-Image-Layered footnote reference and
definition from [^3] to a unique unused identifier, updating both the model row
and its footnote definition consistently. Remove the unused [^vg1] definition
while preserving the FLUX footnote content and all other references.
- Line 182: Update the Qwen-Image-Layered entries in
docs/source/models/supported-models.md lines 182-182 and
docs/source/models/visual-generation.md lines 57-70 to mark FP8 blockwise
support as available. Revise the associated footnote in supported-models.md and
remove FP8/Sage from the unsupported-feature list in visual-generation.md.
In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py`:
- Around line 663-671: Rename the reciprocal normalization tensor currently
stored as latents_std to inv_latents_std and update the denormalization
expression to preserve the existing math without ambiguity. In the VAE decode
flow, split the reshaped latents into bounded chunks before calling vae.decode,
then concatenate the decoded chunks in their original order so output shape and
layer ordering remain unchanged.
- Around line 803-822: Update forward’s dimension handling around the
calculated_width/calculated_height assignments to validate or snap supplied
height and width using resolution_multiple_of, ensuring the resulting dimensions
match the encoded latent extent and img_shapes. In the warmup forward call at
tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py:198-214,
pass height=height and width=width so the captured shape matches
warmup_cache_key; both sites require changes.
- Around line 604-639: Move the generator-list length validation before the
image encoding branch, specifically before _encode_vae_image can consume
generator entries, while preserving the existing batch-size error message.
Update the latents handling so caller-supplied values are packed through
_pack_layered_latents with the same batch, channel, height, width, and layers +
1 arguments as generated latents; document the parameter contract if needed.
In `@tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py`:
- Around line 201-220: Extend
test_layered_e8m0_quantization_only_for_repacking_backends to cover
disable_deep_gemm=True and the non-SM100F path with is_sm_100f=False and
get_sm_version() returning 120. Assert each case matches the expected
_uses_e8m0_post_load_repack behavior while preserving the existing repacking and
Cute DSL coverage.
---
Nitpick comments:
In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py`:
- Around line 743-761: Add a Google-style docstring to the public Qwen image
layered pipeline forward method, documenting its purpose, return value, and all
parameters, including that resolution selects a bucket, layers excludes the
conditioning frame, and latents must be pre-packed. Keep the implementation and
behavior unchanged.
- Around line 727-741: Update the request-to-forward mapping in the infer flow
to source defaults for layers, resolution, cfg_normalize, and use_en_prompt from
the corresponding entries in extra_param_specs instead of hardcoding them.
Preserve values explicitly provided in extra and keep the schema and inference
defaults synchronized.
- Around line 511-550: Make _pack_layered_latents and _unpack_layered_latents
symmetric by using the same frame-count convention and dimension coordinate
space in both APIs. Update their callers, including the paths currently passing
1 and layers + 1, so the pack/unpack round trip preserves the intended number of
frames without internal layers + 1 assumptions or duplicated height/width
conversion.
In
`@tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.py`:
- Around line 405-446: Refactor the duplicated forward path by adding optional
additional_t_cond support to QwenImageTransformer2DModel.forward and passing it
to time_text_embed, then remove the near-identical forward override from the
layered transformer so it inherits the parent implementation. Preserve the
existing mask, rotary embedding, block execution, and return behavior.
- Around line 217-250: Update the batched forward logic around the per-sample
loop in the transformer’s forward method to add a fast path when all rows share
the same valid text length: compact the text states and rotary embeddings once,
invoke super().forward() once for the full batch, and restore outputs to the
original text layout. Preserve the existing per-sample path for varying valid
lengths and maintain timestep, mask, dtype, device, and bias behavior.
🪄 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: bc5a6ecc-5f41-4fef-b0fa-b8d7aa92c6b8
⛔ Files ignored due to path filters (1)
tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zipis excluded by!**/*.zip
📒 Files selected for processing (17)
docs/source/models/supported-models.mddocs/source/models/visual-generation.mdexamples/visual_gen/README.mdexamples/visual_gen/configs/qwen-image-layered-1gpu.yamlexamples/visual_gen/configs/qwen-image-layered-fp8-1gpu.yamlexamples/visual_gen/models/qwen_image_layered.pytensorrt_llm/_torch/visual_gen/models/__init__.pytensorrt_llm/_torch/visual_gen/models/qwen_image_layered/__init__.pytensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.pytensorrt_llm/_torch/visual_gen/models/qwen_image_layered/transformer_qwen_image_layered.pytensorrt_llm/_torch/visual_gen/pipeline_registry.pytests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/qwen_image_layered_lpips_golden.jsontests/integration/defs/examples/visual_gen/test_visual_gen.pytests/integration/test_lists/test-db/l0_b200.ymltests/unittest/_torch/visual_gen/test_qwen_image_layered_parity.pytests/unittest/_torch/visual_gen/test_qwen_image_layered_pipeline_config.pytests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py
| | **Wan 2.2** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | ||
| | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | | ||
| | **Qwen-Image** [^2] | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | No | | ||
| | **Qwen-Image-Layered** [^3] | No | No | No | No | No | No | Yes | Yes | No | No | No | No | |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Document the newly added FP8/Sage capability accurately.
Both matrices state that Qwen-Image-Layered lacks FP8 blockwise support, and the notes also deny Sage support. This PR adds dynamic FP8 block-scale quantization and an FP8/Sage LPIPS regression test, so users will be told the released feature is unavailable.
docs/source/models/supported-models.md#L182-L182: mark FP8 blockwise as supported and revise the associated footnote to stop denying FP8/Sage support.docs/source/models/visual-generation.md#L57-L70: mark FP8 blockwise as supported and remove FP8/Sage from the unsupported-feature list.
📍 Affects 2 files
docs/source/models/supported-models.md#L182-L182(this comment)docs/source/models/visual-generation.md#L57-L70
🤖 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 `@docs/source/models/supported-models.md` at line 182, Update the
Qwen-Image-Layered entries in docs/source/models/supported-models.md lines
182-182 and docs/source/models/visual-generation.md lines 57-70 to mark FP8
blockwise support as available. Revise the associated footnote in
supported-models.md and remove FP8/Sage from the unsupported-feature list in
visual-generation.md.
Comment on lines
185
to
+186
| [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. | ||
| [^3]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition. FP8 blockwise, NVFP4, `trtllm-serve` image-edit routing, and attention-parallel backends are not enabled yet. |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Use a unique, referenced footnote ID.
[^3] duplicates an existing definition, while [^vg1] is unused. Rename the new Qwen-Image-Layered footnote consistently in its row and definition, and remove [^vg1]. This currently triggers markdownlint and can associate the row with the wrong note.
🧰 Tools
🪛 markdownlint-cli2 (0.23.0)
[warning] 185-185: Link and image reference definitions should be needed
Unused link or image reference definition: "^vg1"
(MD053, link-image-reference-definitions)
[warning] 186-186: Link and image reference definitions should be needed
Duplicate link or image reference definition: "^3"
(MD053, link-image-reference-definitions)
🤖 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 `@docs/source/models/supported-models.md` around lines 185 - 186, Rename the
Qwen-Image-Layered footnote reference and definition from [^3] to a unique
unused identifier, updating both the model row and its footnote definition
consistently. Remove the unused [^vg1] definition while preserving the FLUX
footnote content and all other references.
Source: Linters/SAST tools
Comment on lines
+604
to
+639
| image = image.to(device=device, dtype=dtype) | ||
| if image.shape[1] != self.latent_channels: | ||
| image_latents = self._encode_vae_image(image=image, generator=generator) | ||
| else: | ||
| image_latents = image | ||
| self._validate_single_conditioning_frame(image_latents) | ||
| image_latents = self._repeat_conditioning_batch(image_latents, batch_size, "image") | ||
|
|
||
| image_latent_height, image_latent_width = image_latents.shape[3:] | ||
| image_latents = image_latents.permute(0, 2, 1, 3, 4) | ||
| image_latents = self._pack_layered_latents( | ||
| image_latents, | ||
| batch_size, | ||
| num_channels_latents, | ||
| image_latent_height, | ||
| image_latent_width, | ||
| 1, | ||
| ) | ||
|
|
||
| if isinstance(generator, list) and len(generator) != batch_size: | ||
| raise ValueError( | ||
| f"Received {len(generator)} generators for effective batch size {batch_size}." | ||
| ) | ||
| if latents is None: | ||
| latents = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) | ||
| latents = self._pack_layered_latents( | ||
| latents, | ||
| batch_size, | ||
| num_channels_latents, | ||
| h, | ||
| w, | ||
| layers + 1, | ||
| ) | ||
| else: | ||
| latents = latents.to(device=device, dtype=dtype) | ||
| return latents, image_latents |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Generator-list validation runs after its first use, and caller-supplied latents skip packing.
Line 623 validates len(generator) == batch_size, but Line 606 already indexed generator[i] inside _encode_vae_image, so a mis-sized list raises a bare IndexError instead of the intended message. Move the check above the encode call.
Separately, the else at Lines 637-638 returns caller-supplied latents unpacked while the generated path packs them, so forward(latents=...) only works if the caller pre-packs. Worth documenting that contract on the parameter.
🐛 Proposed reordering
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"Received {len(generator)} generators for effective batch size {batch_size}."
+ )
image = image.to(device=device, dtype=dtype)
if image.shape[1] != self.latent_channels:
image_latents = self._encode_vae_image(image=image, generator=generator)
@@
- if isinstance(generator, list) and len(generator) != batch_size:
- raise ValueError(
- f"Received {len(generator)} generators for effective batch size {batch_size}."
- )
if latents is None:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| image = image.to(device=device, dtype=dtype) | |
| if image.shape[1] != self.latent_channels: | |
| image_latents = self._encode_vae_image(image=image, generator=generator) | |
| else: | |
| image_latents = image | |
| self._validate_single_conditioning_frame(image_latents) | |
| image_latents = self._repeat_conditioning_batch(image_latents, batch_size, "image") | |
| image_latent_height, image_latent_width = image_latents.shape[3:] | |
| image_latents = image_latents.permute(0, 2, 1, 3, 4) | |
| image_latents = self._pack_layered_latents( | |
| image_latents, | |
| batch_size, | |
| num_channels_latents, | |
| image_latent_height, | |
| image_latent_width, | |
| 1, | |
| ) | |
| if isinstance(generator, list) and len(generator) != batch_size: | |
| raise ValueError( | |
| f"Received {len(generator)} generators for effective batch size {batch_size}." | |
| ) | |
| if latents is None: | |
| latents = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) | |
| latents = self._pack_layered_latents( | |
| latents, | |
| batch_size, | |
| num_channels_latents, | |
| h, | |
| w, | |
| layers + 1, | |
| ) | |
| else: | |
| latents = latents.to(device=device, dtype=dtype) | |
| return latents, image_latents | |
| if isinstance(generator, list) and len(generator) != batch_size: | |
| raise ValueError( | |
| f"Received {len(generator)} generators for effective batch size {batch_size}." | |
| ) | |
| image = image.to(device=device, dtype=dtype) | |
| if image.shape[1] != self.latent_channels: | |
| image_latents = self._encode_vae_image(image=image, generator=generator) | |
| else: | |
| image_latents = image | |
| self._validate_single_conditioning_frame(image_latents) | |
| image_latents = self._repeat_conditioning_batch(image_latents, batch_size, "image") | |
| image_latent_height, image_latent_width = image_latents.shape[3:] | |
| image_latents = image_latents.permute(0, 2, 1, 3, 4) | |
| image_latents = self._pack_layered_latents( | |
| image_latents, | |
| batch_size, | |
| num_channels_latents, | |
| image_latent_height, | |
| image_latent_width, | |
| 1, | |
| ) | |
| if latents is None: | |
| latents = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype) | |
| latents = self._pack_layered_latents( | |
| latents, | |
| batch_size, | |
| num_channels_latents, | |
| h, | |
| w, | |
| layers + 1, | |
| ) | |
| else: | |
| latents = latents.to(device=device, dtype=dtype) | |
| return latents, image_latents |
🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py`
around lines 604 - 639, Move the generator-list length validation before the
image encoding branch, specifically before _encode_vae_image can consume
generator entries, while preserving the existing batch-size error message.
Update the latents handling so caller-supplied values are packed through
_pack_layered_latents with the same batch, channel, height, width, and layers +
1 arguments as generated latents; document the parameter contract if needed.
Comment on lines
+663
to
+671
| latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, z_dim, 1, 1, 1).to( | ||
| latents.device, latents.dtype | ||
| ) | ||
| latents = latents / latents_std + latents_mean | ||
|
|
||
| batch_size, channels, frames, latent_h, latent_w = latents.shape | ||
| latents = latents[:, :, 1:] | ||
| latents = latents.permute(0, 2, 1, 3, 4).reshape(-1, channels, 1, latent_h, latent_w) | ||
| image = self.vae.decode(latents, return_dict=False)[0].squeeze(2) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Misleading latents_std name, and the VAE decode batch is unbounded.
Line 663 stores 1 / std in a variable named latents_std and Line 666 then divides by it to multiply — correct, but a reader "simplifying" this into latents * latents_std + latents_mean would silently invert the denormalization. Storing the reciprocal under a name like inv_latents_std (or just multiplying by latents_std) removes the trap.
Line 671 also decodes batch_size * layers frames in a single vae.decode call; with layers up to 16 (the extra_param_specs ceiling) at 1024² this is a large activation spike. Chunking the decode would bound peak VRAM.
🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py`
around lines 663 - 671, Rename the reciprocal normalization tensor currently
stored as latents_std to inv_latents_std and update the denormalization
expression to preserve the existing math without ambiguity. In the VAE decode
flow, split the reshaped latents into bounded chunks before calling vae.decode,
then concatenate the decoded chunks in their original order so output shape and
layer ordering remain unchanged.
Comment on lines
+803
to
+822
| if height is None or width is None: | ||
| image_width, image_height = self._image_size(image) | ||
| calculated_width, calculated_height = _calculate_dimensions( | ||
| resolution * resolution, | ||
| image_width / image_height, | ||
| ) | ||
| else: | ||
| calculated_width = width | ||
| calculated_height = height | ||
| if not hasattr(self, "image_processor"): | ||
| from diffusers.image_processor import VaeImageProcessor | ||
|
|
||
| self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) | ||
| image = self.image_processor.resize(image, calculated_height, calculated_width) | ||
| prompt_image = image | ||
| image = self.image_processor.preprocess(image, calculated_height, calculated_width) | ||
| image = image.unsqueeze(2).to(dtype=self.dtype, device=device) | ||
|
|
||
| width = calculated_width // multiple_of * multiple_of | ||
| height = calculated_height // multiple_of * multiple_of |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Output-dimension contract in forward is under-specified in both directions. forward silently derives dimensions from the resolution bucket when height/width are omitted, and silently accepts them unsnapped when they are supplied — so neither caller gets the shape it asked for.
tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py#L803-L822: validate suppliedheight/widthagainstresolution_multiple_of(or snapcalculated_height/calculated_widththe same way as Lines 821-822) soimg_shapesmatches the real encoded latent extent.tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py#L198-L214: passheight=height, width=widthinto the warmupforwardcall so the captured shape matches thewarmup_cache_keythe executor checks.
📍 Affects 1 file
tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py#L803-L822(this comment)tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py#L198-L214
🤖 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/models/qwen_image_layered/pipeline_qwen_image_layered.py`
around lines 803 - 822, Update forward’s dimension handling around the
calculated_width/calculated_height assignments to validate or snap supplied
height and width using resolution_multiple_of, ensuring the resulting dimensions
match the encoded latent extent and img_shapes. In the warmup forward call at
tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py:198-214,
pass height=height and width=width so the captured shape matches
warmup_cache_key; both sites require changes.
Comment on lines
+201
to
+220
| def test_layered_e8m0_quantization_only_for_repacking_backends(): | ||
| module = SimpleNamespace( | ||
| use_cute_dsl_blockscaling_mm=False, | ||
| disable_deep_gemm=False, | ||
| ) | ||
| module_path = ( | ||
| "tensorrt_llm._torch.visual_gen.models.qwen_image_layered.transformer_qwen_image_layered" | ||
| ) | ||
| with ( | ||
| mock.patch(f"{module_path}.is_sm_100f", return_value=True), | ||
| mock.patch(f"{module_path}.get_sm_version", return_value=100), | ||
| ): | ||
| assert QwenImageLayeredTransformer2DModel._uses_e8m0_post_load_repack(module) | ||
|
|
||
| module.use_cute_dsl_blockscaling_mm = True | ||
| with ( | ||
| mock.patch(f"{module_path}.is_sm_100f", return_value=True), | ||
| mock.patch(f"{module_path}.get_sm_version", return_value=100), | ||
| ): | ||
| assert not QwenImageLayeredTransformer2DModel._uses_e8m0_post_load_repack(module) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Cover the remaining E8M0 repacking branches.
Line 201 only exercises the SM100F path. Add cases for disable_deep_gemm=True and is_sm_100f=False, get_sm_version()=120; otherwise regressions in the SM120 repacking path pass this test suite.
🤖 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_qwen_image_layered_registry.py` around
lines 201 - 220, Extend
test_layered_e8m0_quantization_only_for_repacking_backends to cover
disable_deep_gemm=True and the non-SM100F path with is_sm_100f=False and
get_sm_version() returning 120. Assert each case matches the expected
_uses_e8m0_post_load_repack behavior while preserving the existing repacking and
Cute DSL coverage.
Collaborator
|
PR_Github #61857 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61857 [ 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
2 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.
Dev Engineer Review
QA Engineer Review
Added test coverage includes:
The CI test list
tests/integration/test_lists/test-db/l0_b200.ymlincludes the layered example test and both LPIPS tests. Coverage is therefore registered for CI, including the B200 FP8 quality validation.Verdict: sufficient.
Description
Depends on #15096.
This PR adds an FP8 quality preset for Qwen-Image-Layered using dynamic FP8 block-scale Linear quantization and SageAttention FP8. The preset keeps non-backbone Linear layers and the first and last transformer blocks in BF16.
The implementation is scoped to the Qwen-Image-Layered model. It routes masked joint attention through the TRTLLM backend when selected while preserving the existing VANILLA path.
For MXFP8 Linear weights, the existing two-step path first quantized with FP32 block scales and then repacked to E8M0 scales. That introduced an additional QDQ error. This PR prequantizes Qwen-Image-Layered dynamic FP8 block-scale weights directly with E8M0-compatible scales before the existing weight-loading path, avoiding requantization without changing the shared quantization APIs or Qwen-Image behavior.
Test Coverage
0.322196with two-step quantization to0.035642with direct one-step E8M0 quantization, passing the< 0.05threshold.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.