FLUX 3 Image Klein egy képészítésre és képszerkesztésre tervezett modell a Black Forest Labs (BFL) fejlesztésében. A modellt NVIDIA GB200 NVL72 rendszeren tanították, és a cikkben bemutatott munkamenet célja a modell optimalizálása NVIDIA GPU-kon a TensorRT LLM VisualGen futtatókörnyezet segítségével. A generatív diffúziós modellek jellemzően többszöri, egymás után végzett denoising lépést igényelnek milliárdos paraméterű hálókon, és a nagy felbontás figyelemköltségeket (attention) növel.
FLUX 3 Image Klein két kiadást kínál: a Base és a distilled checkpointokat, amelyek különböző késleltetés–minőség kompromisszumot nyújtanak.
Mi szerepel az optimalizálásban
A TensorRT LLM VisualGen több, egymásra épülő technikát használ a késleltetés csökkentésére:
- NVFP4 és FP8 kvantált kernelvariánsok (BFL által előkvantált checkpointok támogatása)
- FlashInfer/PrimsTS alapú GEMM-k és epilógusfúziók (CuTlass Primitives + Task Scheduling)
- kvantált (cuDNN FP8 és VC‑Attention FP8) és blokksparse (Sol‑Attention) attention-implementációk
- Cache‑DiT: lépésenkénti cache a transformer részintermediátok újrafelhasználására
- torch.compile alapú fordítás/kompiláció
- több GPU-s párhuzamosítás Ulysses és CFG megközelítéssel
- CUDA Graphs a kernel‑indítási overhead csökkentésére kisebb munkaterhelés mellett
A cikkben szereplő mérések elsősorban NVIDIA Blackwell adatközponti GPU-kon (például B200) és NVIDIA RTX PRO 6000 Blackwell Server Edition GPU-kon készültek.
Modellezés és pipeline jellemzők
A FLUX 3 Image Klein három fő összetevőből áll: Qwen3‑VL‑4B‑Instruct a szövegfeltételhez, egy diffusion transformer (DiT) a latens kép finomításához és egy autoencoder a latensek dekódolásához és I2I referencia kép enkódolásához. A Klein külön kezeli a szöveges, zajos célkép és opcionális referencia‑képi bemeneteket, majd egyesíti őket közös attention rétegen keresztül.
A cikkben alkalmazott mérési beállításoknál a Base 50 denoising lépéssel és 3.5‑ös classifier‑free guidance (CFG) skálával futott; a distilled Klein 16 lépéssel, CFG nélkül. A benchmarkok általában torch.compile és seed=0 használatával készültek.
Pipeline késleltetések (NVIDIA B200, 2048×2048)
A cikkben szereplő időmérések egy NVIDIA B200 GPU‑n, NVFP4, VANILLA attention és torch.compile mellett készültek. Az eredmények (átlag 10 kérés, 2 bemelegítő kérés után):
- T2I Base: Text+preparation 47 ms, denoising 14.24 s, decode+postprocess 416 ms, összesen 14.70 s (denoising arány: 96.8%)
- T2I Distilled: 24 ms, 2.28 s, 410 ms, összesen 2.71 s (84.0%)
- I2I Base: 468 ms, 39.40 s, 417 ms, összesen 40.29 s (97.8%)
- I2I Distilled: 407 ms, 6.65 s, 418 ms, összesen 7.48 s (89.0%)
Ezekből látszik, hogy a denoising (DiT) domesztikálja a teljes generálási időt, különösen a Base konfigurációnál.
Részletes optimalizációs lépések és hatásuk
-
GEMM-optimalizációk: a PrimsTS FP8/NVFP4 GEMM‑k (FlashInfer) a FLUX 3 tensorformákra hangolva futnak; a kiinduló pinned NVFP4 PrimsTS kernel 1.02–1.10× gyorsulást adott T2I-ben a B200-on az autotuned TensorRT LLM NVFP4 GEMM‑hez képest.
-
Kvantált attention: cuDNN FP8 és VC‑Attention FP8 csökkentik az attention költségét. A 2048×2048 NVFP4 T2I tesztek szerint cuDNN FP8 1.08× (Base) és 1.09× (distilled) gyorsulást hozott, LPIPS 0.072 és 0.161 értékekkel. VC‑Attention FP8 1.18× (Base) és 1.16× (distilled) gyorsulást adott, LPIPS 0.086 és 0.151 eredményekkel. A szerzők a VC‑Attention FP8‑at javasolják a 2048×2048 B200‑as környezethez, repair budget ρ=0.005 beállítással.
-
Sol‑Attention (blokksparse): Sol‑Attention τ paraméterrel szabályozza a ritkítást. NVFP4 Base‑re 2048×2048 B200‑on Sol‑Attention τ=1.0 mellett 1.24× gyorsulást mértek a denoising+dekódolás fázisra a VANILLA kompiltált sűrű baseline‑hoz képest (16 kezdeti dense lépéssel), vagy 1.30× ha 12 kezdeti dense lépést használnak. A nagyobb τ rontotta a kimenet minőségét relatíve kevés többletgyorsulással.
-
Cache‑DiT: a lépések közötti előfeldolgozott transformer eredmények újrafelhasználásával csökkenti az inferenciaidőt. A konfigurálható residual_diff_threshold segítségével lehet finomhangolni a sebesség/minőség trade‑offot. Például ρ=0.24 mellett NVFP4 Base T2I‑nél a 2048×2048‑as tesztek közel kétszeres (1.94×) gyorsulást mutattak; a distilled kisebb nyereséget ért el (1.20×). 1024×1024‑on Base: 1.97×, distilled: 1.26× gyorsulás volt megfigyelhető.
-
Multi‑GPU (Ulysses + CFG): Base esetén a pozitív/negatív ágakat külön GPU csoportokra rendelve és Ulysses sorparallelizmussal osztva a munkát, nyolc GPU‑ra skálázva 2048×2048‑on a medián T2I késleltetés 14.90 s‑ról 3.00 s‑re csökkent; distilled esetén 2.73 s‑ról 1.08 s‑re. Denoising sebességjavulás nagyobb: Base 6.01×, distilled 4.44×, de a teljes end‑to‑end javulást korlátozta, hogy a deszaturált részek nem mind denoisingben zajlanak.
-
CUDA Graphs: kisebb per‑GPU munkamennyiségnél a kernel‑indítási overhead csökkenthető. Egy 1024×1024, nyolc B200‑as kísérletben CUDA Graphs hozzáadásával Base‑n a latency 1.823 s‑ról 0.809 s‑ra (2.25×), distilled‑n 0.791 s‑ról 0.366 s‑re (2.16×) csökkent.
Összesített end‑to‑end eredmények
A cikkben bemutatott kumulatív optimalizációs pipeline (kompiláció, NVFP4 + PrimsTS GEMM, Sol‑Attention, Cache‑DiT) egyetlen NVIDIA B200 GPU‑n, 2048×2048‑on a következő javulásokat adta a nem kompiltált BF16 VANILLA összehasonlításhoz képest:
- T2I: akár 3.46× gyorsulás, a latency 22.36 s‑ról 6.46 s‑re csökkent
- I2I: akár 4.45× gyorsulás, a latency 54.64 s‑ról 12.28 s‑re csökkent A 16‑lépéses distilled Klein esetén legfeljebb 2.34× end‑to‑end gyorsulást mértek. Egy NVIDIA RTX PRO 6000 Blackwell Server Edition GPU‑n hasonló konfigurációk 1.73–3.25× gyorsulást adtak.
A minőségbeli eltéréseket LPIPS mutatja; a cikk számos összehasonlítást és példát közöl (figura‑sorozatok az optimalizációs lépések hatásáról).
Telepítési és futtatási példák
A szerzők mellékelnek gyakorlati lépéseket: előre elkészített TensorRT LLM konténer indítása, offline példa futtatása és OpenAI‑kompatibilis /v1/images/generations szerver elindítása. Példák:
- docker run parancs a megadott image‑szel és környezeti változóval (TLLM_FMHA_LIBS=+prims_ts)
- offline példa: python models/flux3.py --model black-forest-labs/flux-3-image-klein-base --visual_gen_args configs/flux3-images-1gpu.yaml --height 2048 --width 2048 --seed 0
- szerver indítása: trtllm-serve black-forest-labs/flux-3-image-klein-base --visual_gen_args configs/flux3-images-1gpu.yaml
A cikk tartalmazza továbbá a YAML konfigurációk mintáit NVFP4 + VC‑Attention FP8 és NVFP4 + Sol‑Attention + Cache‑DiT beállításokra, valamint utasításokat a variant (nvfp4 / fp8r / bf16) kiválasztására.
Következtetés
TensorRT LLM VisualGen részletes és konfigurálható eszközkészletet kínál a FLUX 3 Image Klein gyorsítására: NVFP4/FP8 kvantálás, fuzionált GEMM‑k, sparse attention, lépéscache és több GPU‑s párhuzamosítás egyaránt járulnak hozzá a késleltetés csökkentéséhez. A bemutatott kombinációk egyetlen B200 GPU‑n 3.46× T2I és 4.45× I2I gyorsulást értek el a nem kompiltált BF16 VANILLA alaphoz képest, miközben a kimenetek perceptuális hasonlóságát LPIPS metrikával mérik és közlik.
Hivatkozások és köszönetnyilvánítás
A cikk szerzői és közreműködői: Nikita Korobov, Harrison Zhang, Pamela Peng, Kanghwan Jang, Yuhang He és Charlotte Mallo. A részletes lépések, YAML minták és telepítési útmutató a TensorRT LLM VisualGen környezetben találhatók a cikk mintapéldái szerint.
