### What does this PR do?
Type of change: Bug fix
Adds a compact SDXL and SDXL-Turbo mixed-precision FP4 recipe:
- block-16 NVFP4 for non-QKV Linear/GEMM layers;
- FP8 for Conv2d layers;
- high-precision Q/K/V projection Linears to preserve TensorRT
horizontal fusion;
- optional FP8 MHA quantization.
For SDXL FP4 export, Conv2d quantizers export directly through the
shared FP8 custom-op path. The previous `generate_fp8_scales` plus
`convert_zp_fp8` INT8 zero-point workaround is removed. The graph then
uses the existing FP8 Q/DQ normalization and `NVFP4QuantExporter`
lowering, with opset 23 for FLOAT4 support. Flux FP8 export also saves
the graph returned by its RoPE weight conversion.
This PR also changes shared exporter behavior:
- `_fp8_quantize` refreshes ONNX shape/type inference after applying the
custom FP8 operator's uint8 output metadata, affecting all FP8 ONNX
exports through this symbolic.
- `_quantized_sdpa` derives `disable_fp8_mha` from the live Q/K/V
quantizer state instead of a restored private module flag.
Other model recipe configurations remain unchanged.
### Usage
```bash
python quantize.py \
--model sdxl-1.0 \
--model-dtype Half \
--trt-high-precision-dtype Half \
--format fp4 \
--block-size 16 \
--batch-size 2 \
--calib-size 128 \
--n-steps 20 \
--quantized-torch-ckpt-save-path ./sdxl-fp4 \
--onnx-dir ./onnx-sdxl-fp4
```
### Testing
- CPU-only focused and generic NVFP4 exporter tests: 44 passed in 4.35
seconds.
- Focused Flux returned-graph save test: 1 passed.
- Required Linux unit CI at `034fe23ec` passed with the `all` dependency
set, including `tests/unit/examples/test_diffusers_fp4.py`.
- Latest changed-file pre-commit checks: all passed.
- TensorRT 10.14 on a B200 GPU:
- 302 native block-scaled NVFP4 GEMM tactics;
- 38 native FP8 Conv tactics;
- no FP4 Q/K/V projections;
- all 11 FP16 Q/K/V projection-fusion groups preserved;
- three alternating batch-2 profiles measured 18.614 ms FP4 versus
20.028 ms FP16 median UNet latency, a 7.06% reduction.
- FP8 SDXL/SD3 ONNX-to-TensorRT end-to-end runs were not executed
because they require explicit approval. The existing end-to-end test
matrix now includes SD3 FP8 alongside SDXL FP8.
### Before your PR is "*Ready for review*"
Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).
Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).
- Is this change backward compatible?: ✅ — no public API or CLI flags
change; the shared changes preserve the intended FP8 export and
attention behavior.
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
✅ — the shared NVFP4 opset, FP8 shape-inference, and Diffusers
attention-policy changes are recorded under bug fixes.
- Did you get Claude approval on this PR?: N/A
### Additional Information
Tracking: [5565357]
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
- **New Features**
- Added SDXL support for mixed NVFP4/FP8 quantization, including
convolution and softmax handling.
- Added an SDXL quantization preset for streamlined post-training
quantization workflows.
- Expanded FP4 ONNX export support to Flux and SDXL, with improved
FP4/FP8 graph processing and export reliability.
- Added automatic quantization policy and format restoration from
checkpoints.
- **Documentation**
- Documented SDXL layer behavior, optional FP8 attention quantization,
and Blackwell/TensorRT requirements for FP4 and FP8 deployment.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
> 🤖 _Generated by Codex (AI agent)._
---------
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Co-authored-by: Codex <codex@openai.com>
### What does this PR do?
Type of change: chore
Bumps CI/dev tooling and test containers.
**Container bumps**
- NeMo test containers → 26.06
- TRT-LLM container → 1.3.0rc19
- transformers max version → 5.12
**Dev tooling bumps**
- ruff bump 0.12.11 → 0.15.18
- mypy 1.17.1 → 2.1.0: enable new defaults (`local_partial_types`,
`strict_bytes`); fix/narrow the errors newly surfaced by mypy 2.0 in 4
modules (rather than blanket-suppressing them); remove 2 stale `# type:
ignore` comments
- pre-commit 4.3.0 → 4.6.0
- sphinx 8.1 → 9.1 + sphinx-rtd-theme 3.0 → 3.1: add `suppress_warnings
= ["ref.python"]` to fix cross-reference ambiguity error new in sphinx
9.x
- trl fix for newly released 1.7 version
**Bug fixes surfaced by the bumps**
- sparsity (weight): make the weight mask DTensor-aware under FSDP. The
transformers→5.12 bump routes the HF Trainer FSDP optimizer-state save
through torch's DTensor-based `get_optimizer_state_dict`, which
triggered `aten.mul.Tensor got mixed torch.Tensor and DTensor` in the
dynamic `weight` getter. The mask is now distributed to the weight's
mesh/placements before masking, cached, and rebuilt only when the
sharding changes (invalidated on `set_mask`). Fixes the `llm_sparsity`
example test.
### Testing
- `pre-commit run --all-files` ✅ (including mypy 2.1.0)
- `nox -s docs` ✅
- `tests/unit/torch/sparsity` + `tests/unit/torch/nas` ✅
- `llm_sparsity` GPU example test (FSDP path) verified in CI
### Before your PR is "*Ready for review*"
- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: N/A
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A
- Did you get Claude approval on this PR?: ✅
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Summary
* **Documentation**
* Refreshed Docker pre-requisites across examples to recommend updated
container image tags (and streamlined some instructions).
* **Bug Fixes**
* Improved sparse weight mask handling for DTensor/FSDP by aligning and
caching distributed masks.
* Made TensorRT engine byte retrieval return immutable `bytes`.
* Reduced Sphinx cross-reference warnings and tuned Transformers
compatibility warning thresholds.
* **Tests**
* Increased default unit test timeout on Windows runners.
* **Chores**
* Updated CI workflow container tags and refreshed linting/typing/docs
version pins, plus related mypy configuration.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
## What does this PR do?
bug fix
**Overview:**
[NVBug 5659126] dummy inputs are kwargs maps, and rename the
generate_... function to more explicitly reflect the return value type
## Testing
`python diffusion_trt.py --model flux-dev --override-model-path
/models/FLUX.1-dev --torch --benchmark --skip-image
`
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
## What does this PR do?
**Type of change:**
New feature
**Overview:**
- Implemented the NVFP4QuantExporter
- Deprecated fp4qdq_to_2dq
- Updated tests
## Usage
```python
python torch_quant_to_onnx.py --quantize_mode=nvfp4 \
--onnx_save_path=vit_base_patch16_224.nvfp4.onnx \
--calibration_data_size 64 \
--batch_size 128
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
```
python evaluate.py --onnx_path=vit_base_patch16_224.nvfp4.onnx \
--model_name=vit_base_patch16_224 \
--results_path=./results.txt \
--batch_size 128
```
Results:
```
The top1 accuracy of the model is 84.39%
The top5 accuracy of the model is 97.312%
Inference latency of the model is 7.22412 ms
```
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: No
- Deprecated fp4qdq_to_2dq
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
## What does this PR do?
**Type of change:**
Bug fix
**Overview:**
1. The diffusion_trt.py needs the dynamic_shapes when running trtexec
for engine building. A previous change altered the format of
dynamic_shapes, fix it here.
2. the dynamic_shapes logic gets cleaned up. The existing logic is very
confusing
3. recover min-batch_size config for some pipelines. Previously some
pipelines set the min batch_size to be > 1, which was odd, so a previous
change sets them to be 1, but it turns out the oddity has a reason, the
trt engine building fails with the altered batch_size min/opt, thus
recover them.
## Testing
pytest tests/examples/diffusers
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
## What does this PR do?
**Type of change:**
new feature
**Overview:**
Support Wan2.2 t2v diffusers quantization
1. fix torch2.9 support
2. add Wan2.2 t2v diffusers pipeline quantization
Main difference of the Wan2.2 pipeline comparing to exisiting
pipelines is that there are 2 backbone models for denoising. For the
quantization therefore we need to quantize both of them.
However, it turns out our base library does not well support
quantization of multiple models in the same time. Therefore, the
change here just stick to quantize a single model each time, and then
run the quantization multiple times.
So, we need to allow users to pick which backbone to quantize,
therefore adding a new argment for it
3. add a workaround for the exporting ONNX issue when we upgrade
diffusers to >= 0.35.0. The issue lies is the exporting of the
torch.nn.RMSNorm. Some pipelines in the diffusers > 0.35.0 use the
torch version RMSNorm while before that they use the diffusers' own
version of RMSNorm. It turns out they are directly replacable so the
workaround is to simply replace the torch RMSNorm usages with
diffusers RMSNorm. But we need to fix it properly soon by porting our
ONNX export to be based on torch dynamo instead of torchscript. Issue
reported from external user:
https://github.com/NVIDIA/TensorRT-Model-Optimizer/issues/262
4. allow use of a prompts file, which is simply a text file with a list
of prompts, one prompt each line
5. allow each component of a pipeline to have different dtype accuracy.
added a new list stype command line arg --component-dtype for this.
example: --component-dtype vae:Float
6. print the summary of the quantized model so users can capture issues
from
log
## Usage
python quantize.py \
--model wan2.2-t2v-14b \
--format fp8 \
--batch-size 4 \
--calib-size 64 \
--n-steps 20 \
--backbone transformer \
--model-dtype BFloat16 \
--component-dtype vae:Float \
--trt-high-precision-dtype BFloat16 \
--quantized-torch-ckpt-save-path ./wan_transformer.pt \
--onnx-dir wan-transformer-onnx \
--prompts-file wan-prompts.txt
## Testing
Tested SDXL_BASE, LTX_VIDEO_DEV, WAN22_T2V
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
https://github.com/NVIDIA/TensorRT-Model-Optimizer/issues/262
---------
Signed-off-by: Shengliang Xu <shengliangx@nvidia.com>
## What does this PR do?
**Type of change:**
Minor code change
**Overview:**
- Delete backbone after Device Model creation
- Add assertion for torch compile
- Update dummy input generation function
## Testing
```
python diffusion_trt.py --model flux-dev --benchmark --skip-image
python diffusion_trt.py --model flux-dev --benchmark --skip-image --restore-from ./flux_dev_fp8_autodeploy_fake.pt
python diffusion_trt.py --model flux-dev --benchmark --skip-image --restore-from ./flux_dev_fp4_autodeploy_fake.pt
python diffusion_trt.py --model flux-dev --benchmark --skip-image --torch
python diffusion_trt.py --model flux-dev --benchmark --skip-image --restore-from ./flux_dev_fp8_autodeploy_fake.pt --torch
python diffusion_trt.py --model flux-dev --benchmark --skip-image --restore-from ./flux_dev_fp4_autodeploy_fake.pt --torch
python diffusion_trt.py --model flux-dev --benchmark --skip-image --torch --torch-compile
```
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: No
- **Did you add or update any necessary documentation?**: No
- **Did you update
[Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**:
No <!--- Only for new features, API changes, critical bug fixes or bw
breaking changes. -->
---------
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
