Skip to content

[5565357] Fix SDXL NVFP4 export and performance - #2336

Draft
ajrasane wants to merge 1 commit into
mainfrom
arasane/fix_sdxl_nvfp4_export
Draft

[5565357] Fix SDXL NVFP4 export and performance#2336
ajrasane wants to merge 1 commit into
mainfrom
arasane/fix_sdxl_nvfp4_export

Conversation

@ajrasane

@ajrasane ajrasane commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: Bug fix

Adds an SDXL-specific 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.

The previous SDXL path could accept mismatched checkpoints and lose projection fusion. This change validates fresh and restored quantizer state, rejects incompatible checkpoints before export, normalizes supported FP8 Q/DQ, validates raw and final NVFP4 graph topology, and publishes checked ONNX artifacts atomically.

Non-SDXL quantization and the generic NVFP4 exporter remain unchanged.

Usage

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 unit tests: 100 passed.
  • Changed-file pre-commit checks: all passed.
  • TensorRT 10.14 on a Blackwell GPU:
    • 302 native block-scaled FP4 GEMM tactics;
    • 38 native FP8 Conv tactics;
    • zero FP4 Q/K/V projections;
    • Q/K/V projection fusion preserved;
    • three alternating batch-2 profiles measured 18.779 ms FP4 versus 20.149 ms FP16 median UNet latency, a 6.8% reduction.

Before your PR is "Ready for review"

  • Is this change backward compatible?: ✅ — supported interfaces and non-SDXL behavior are unchanged; mismatched SDXL FP4 checkpoints now fail early with instructions to recalibrate.
  • 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?: N/A — SDXL FP4 was not previously a supported deployment path.
  • Did you get Claude approval on this PR?: N/A

Additional Information

Tracking: [5565357]

🤖 Generated by Codex (AI agent).

Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
@copy-pr-bot

copy-pr-bot Bot commented Sep 4, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Comment @coderabbitai help to get the list of available commands.

@codecov

codecov Bot commented Sep 4, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.31%. Comparing base (f13a796) to head (ccece7f).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2336   +/-   ##
=======================================
  Coverage   79.31%   79.31%           
=======================================
  Files         527      527           
  Lines       61482    61482           
=======================================
  Hits        48765    48765           
  Misses      12717    12717           
Flag Coverage Δ
unit 55.89% <ø> (+0.01%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant