WritingBasetenBasetenpublished Feb 18, 2026seen Jun 26

Captured source

source ↗
published Feb 18, 2026seen Jun 26captured Jun 27http 200method plain

4-Bit Quantization for Inference Optimization Announcing our Series F . Learn more

Model performance

Four Bits: 4-Bit Quantization for Inference Optimization in FLUX.2 [dev]

This writeup on FLUX.2 [dev] inference optimization reveals exactly how we achieved a 1.6x speed improvement with no perceptible quality loss.

Authors

Ali Taha

Last updated February 18, 2026

Share

One of these images was generated by a model compressed to a third of its size. The other generated by the same model in full precision. Which is which? ✕ Image A is the full-precision model, image B is the 4-bit model Trimming off two-thirds of a model’s weights implies, by definition, a loss of information and a degradation in overall quality. In this post, we explain exactly how we optimized the best image generation model on the market today with 4-bit quantization for inference optimization, achieving a 1.6x speed improvement with (almost) no quality impact. The FLUX.2 [dev] model First off, a deep dive into the architecture of the FLUX.2 [dev] model. This is a SOTA diffusion model, and follows most diffusion model architectures with some tweaks. We run our inference, denoising steps, and kernel optimizations in SGLang. Here’s an overview of the entire process. ✕ Full process of a single denoising step for FLUX.2 [dev]. We will examine this diagram in chunks throughout this writeup. How does FLUX.2 [dev] run inference? Image generation models start by passing the prompt through a text encoder for a single forward pass. For FLUX.2 [dev], this text encoder is a Mistral model. In this piece, we’ll skip an examination of the text encoder, and focus on the main inference loop for image generation. Optimizing this loop is the core of inference optimization for diffusion models. The Mistral model runs the text prompt ( e.g “a blue tiger with amber eyes”). It adds meaning to this prompt, each passing layer growing richer in meaning. FLUX.2 [dev] then extracts from three separate layers of the text encoder the hidden latents/activations and concatenates them together to create a 512x5120x3 embedding matrix. This captures the semantic meaning of what image we want to generate. ✕ Simplified embedding matrix capturing semantic meaning At its core, FLUX.2 [dev] takes in this embedding matrix above, and one more input- the noisy latent. A noisy latent is one additional matrix, instantiated with random numbers. This additional matrix is our starting point. It has no meaning and it can be thought of as an empty canvas onto which our model will draw its output. It then feeds both of these to two separate embedders (MLPs at heart): The X-embedder for the latent canvas

The context embedder for the embedding matrix

✕ Inputs to the embedders Outputting two separate hidden state matrices, the model then feeds the matrices into a transformer block that takes these states as input, representing the canvas (what has been drawn), and the state of the prompt (what it needs to draw). ✕ 8 Dual Transformer Blocks The output of this 8x transformer block is then fed to another transformer block, this one much larger, and takes as its input the concatenated result of the two output states. ✕ 48 Single Transformer Blocks The output of this second transformer block gets transformed via an MLP (linear projection) into a matrix that matches the dimensions of our latent canvas. More precisely, the final hidden matrix is just a delta that will be added to the original random canvas (the noisy latent instantiated at the start of inference ). ✕ Final output gets projected into the latent space Along with this delta comes a weight multiplier that emphasizes how seriously the model should take this nudge. Take the analogy of an artist. At the beginning, the canvas is empty, and the artist can do a lot of work (tracing, outlining, filling in the background, coloring the sky) in a few passes, with a large paintbrush and heavy strokes. Consequent passes refine the details (adding wrinkles to a human face), are done with care, and in a few light strokes. Similarly, at the start of inference (i.e., this is the first pass and our input latent canvas is complete noise), the model multiplies this nudge with a very high value. There is very little on the canvas, and it needs to do a lot of work On the other hand, by the 8th and final pass, the model is probably just refining very minor details and shouldn’t change much - the core image is already placed. The inference process then take this final answer, this updated latent canvas, and feeds it back as the input for the next pass. The number of times this process is repeated, ie the number of passes the model goes through, is referred to as denoising steps , and is usually 8. Inevitably, each step needs its own weight, referred to as a timestep . The earlier the denoising step, the higher the timestep. These time steps are used to create the shift, the scale, and the gate- modulation matrices used by the transformer to manipulate outputs (weight multipliers). ✕ Full FLUX.2 [dev] diagram For your convenience, here is a zoomed-in look at the two transformer blocks. ✕ Single Transformer block ✕ Dual Transformer block Profiling FLUX.2 [dev] inference Within this architecture, the most common operation during inference is a matrix multiplication. In fact, a profiling run of the model showed that it spent 67% of its time running GEMM (General Matrix Multiplication) kernels. ✕ Profiler output of a naive implementation shows matmul as the main bottleneck To be precise, the model spent 1.82 ms in each matmul kernel, and took, from prompt to image, 2.776s . This profiling step is essential for inference optimization: before applying 4-bit quantization, we needed to identify the true performance bottleneck. What if instead, we run these BF16 GEMM kernels in a lower precision, specifically FP4. Why would that make inference faster? On the memory side, memory transactions on the GPU don't care how many numbers they grab; they think in terms of bytes. Memory transactions move fixed-size chunks, so using smaller data types means we can fetch more values in the same transaction. Smaller data types increase effective bandwidth and cache residency (memory bound). On the compute side, Tensor Cores can perform more operations with lower precision formats. By moving from FP16 to FP4, we access cores with 4x higher FLOPS (compute bound). Whether our inference process is compute-bound or memory-bound (depends on...

Excerpt shown — open the source for the full document.

Notability

notability 3.0/10

Routine technical blog post by ML infra company.