tensorRT: fp8 tier for the small DiTs (sm-music, sm-sfx) - #89
Merged
Conversation
fp8 E4M3 grafted onto the linear GEMMs of the known-good fp16mixed graph (attention fp16-fused, fp32 RoPE/RMSNorm islands left intact) — deliberately NOT the medium's baked-RoPE recipe, since the small DiTs use standard attention and never had bf16's long-angle RoPE problem. A clean weight-halving tier: engine 479 vs 936 MB, velocity-cos ~0.99 vs eager, clip% at/below fp16mixed; only marginally faster (~1.10-1.17x) because a small DiT's ~5 ms forward at batch 1 is overhead-bound, so fp8's GEMM-math savings barely show. Default stays fp16mixed; fp8 is selectable for the smaller engine/weight footprint. - canon (sa3_trt_core.py): register fp8 for sm-music/sm-sfx, drop the medium-only fp8 guard, update precision docs/CLI help. - build_from_onnx.py: sa3-sm-music-fp8 / sa3-sm-sfx-fp8 STRONGLY_TYPED targets. - build/make_dit_fp8_smalldit.py: the producer (calibrate + fp8-QDQ graft; fp16 scales, floored, so the fp16 trunk stays consistent). NB producing these via #47's ModelOpt path breaks the fp32 islands (collapses to velcos ~0.69). - docs: build/README.md + runtime README precision sections. Engines + ONNX uploaded to HF (tensorRT/sm_90/sa3-sm-{music,sfx}/dit_fp8.trt, onnx/sa3-sm-{music,sfx}/dit_fp8.onnx). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
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
1 participant
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.
Adds a selectable fp8 engine for
sm-musicandsm-sfx, alongside theirfp16mixeddefault.What it is. fp8 E4M3 grafted onto the linear GEMMs of each model's known-good
dit_fp16mixed.onnx— attention stays fp16-fused and the fp32 RMSNorm/RoPE islands are left exactly as the fp16mixed producer made them. This is not the medium DiT's baked-RoPE recipe: the small DiTs use standard (non-differential) attention and never had bf16's long-angle RoPE problem, so there's nothing to bake.Honest positioning — a clean weight-halving tier, not a speed tier.
The small DiTs' ~5 ms forward at batch 1 is overhead-bound, so fp8's GEMM-math savings barely show (~1.10–1.17×), unlike medium's ~1.3×. Default stays
fp16mixed; fp8 is for when the smaller engine / weight footprint helps. Not seed-reproducible vs fp16mixed.⚠ Recipe note. Do not produce these via
build_dit_fp8.py(#47's ModelOpt path): on the small graphs its island-flatten + reapply doesn't restore the fp32 islands and the engine collapses to velocity-cos ~0.69 with clipping (the fp8 GEMMs themselves are fine — measured 0.9995 per-GEMM / 0.984 whole-DiT in torch; it was purely the islands). Grafting onto the fp16mixed ONNX keeps the islands correct by construction.Changes. canon registration (
sa3_trt_core.py);build_from_onnx.pysa3-sm-music-fp8/sa3-sm-sfx-fp8STRONGLY_TYPED targets;build/make_dit_fp8_smalldit.pyproducer; build + runtime README precision sections.Artifacts. Engines + ONNX uploaded to HF:
tensorRT/sm_90/sa3-sm-{music,sfx}/dit_fp8.trt,onnx/sa3-sm-{music,sfx}/dit_fp8.onnx. sm_90 only; other arches rebuild via the build target. Smoke-tested through the production runtime (--dit sm-music --precision fp8, mega-graph render OK).🤖 Generated with Claude Code