Skip to content

FLUX.3 Klein fused family; one-time device move of fused group constants (0.13.0) - #19

Merged
iamwavecut merged 1 commit into
mainfrom
flux3-family-0.13.0
Oct 15, 2026
Merged

iamwavecut merged 1 commit into
mainfrom
flux3-family-0.13.0

Conversation

@iamwavecut

Copy link
Copy Markdown
Owner

Changes

  • orbitquant.fused.families.flux3 fuses Black Forest Labs' native flux3.dit.DiT (FLUX.3 Image Klein): per-mode and single-stream blocks, Q|K|V and SwiGLU GEMMs over one set of activation codes, Q/K RMSNorm and RoPE in one kernel, gated residuals in the GEMM epilogues. A mode whose tokens share one time embedding computes its AdaLN modulation on one row.
  • Fused groups built on the host kept their codebook surrogate codes on the CPU, and every decode copied them to the GPU: one blocking host-to-device transfer per GEMM. The constants now move to the weight device once. This affects every fused family that is fused after the model is on the GPU.
  • The activation quantizer no longer copies a plain bf16 input before the RPBH rotation; the prologue pass writes only the token norms.
  • Version 0.13.0.

Measurements

One RTX 5090, Torch 2.12.1 CUDA 13.0, FLUX.3 Klein W4A4 release, 1024x1024, warm median:

Runtime Text to image Edit, 1 reference GPU peak
0.12.0 + bundled family 2.53 s 5.32 s 16.9 GiB
this branch 2.06 s 4.72 s 12.4 GiB*

The images are bit-identical to the 0.12.0 runtime for both workloads. *The peak also includes loader changes in the model bundle (text encoder trimming, expandable allocator segments).

Checks

GPU test run and kernel checks on the RTX 5090 are in progress; results will be added here before merge.

… constants (0.13.0)

- orbitquant.fused.families.flux3 fuses Black Forest Labs' native flux3.dit.DiT: per-mode and
  single-stream blocks with Q|K|V and SwiGLU GEMMs over one set of codes, Q/K RMSNorm + RoPE in
  one kernel and gated residuals in the GEMM epilogues. A mode whose tokens share one time
  embedding computes its AdaLN modulation on one row.
- Fused groups built on the host kept their codebook surrogate codes on the CPU and copied them
  to the GPU in every decode: a blocking transfer per GEMM. The constants now move to the weight
  device once. On an RTX 5090 this cuts FLUX.3 Klein 1024x1024 generation from 2.44 s to 2.06 s
  with bit-identical images.
- The activation quantizer no longer copies a plain bf16 input before rotating it; it writes
  only the token norms.
@iamwavecut

Copy link
Copy Markdown
Owner Author

GPU checks, RTX 5090, Torch 2.12.1 CUDA 13.0:

  • Full test suite with CUDA/Triton: all fused, kernel and FLUX.3 tests pass. Five failures are outside this change: four tests read files that the synced test tree did not contain (.github workflows, sdist layout), and test_real_cuda_projection_reuse_is_bit_identical fails for main's src too when the native-kernels source tree sits next to it; this branch's src passes it in a bare tree.
  • scripts/run_cuda_kernel_checks.sh: all stages completed.
  • FLUX.3 Klein release runtime: images bit-identical to the 0.12.0 runtime for text to image and one-reference edit.

@iamwavecut
iamwavecut merged commit 62aec03 into main Oct 15, 2026
3 checks passed
@iamwavecut
iamwavecut deleted the flux3-family-0.13.0 branch October 15, 2026 19:53
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