Implementação da quantização NVFP4 e MXFP8

A geração de imagens e vídeos por meio de modelos de difusão tem ganhado cada vez mais popularidade, entregando mídias visuais hiper-realistas. No entanto, a adoção desses modelos muitas vezes esbarra em requisitos extremamente elevados de memória e poder computacional. A quantização surge como uma solução essencial para viabilizar a execução eficiente desses modelos em ambientes de produção.

Neste artigo, demonstramos ganhos reprodutíveis de até 1,26x na velocidade de inferência com o MXFP8 e até 1,68x com o NVFP4, utilizando as bibliotecas Diffusers e TorchAO nos modelos Flux.1-Dev, QwenImage e LTX-2 rodando em GPUs NVIDIA B200. Além disso, detalhamos como empregamos técnicas como quantização seletiva, CUDA Graphs e LPIPS como métrica para otimizar tanto o desempenho quanto a precisão desses modelos. Todo o código utilizado para reproduzir os experimentos está disponível publicamente.

É importante destacar que os formatos MXFP8 e NVFP4 são suportados nativamente pela arquitetura Blackwell da NVIDIA, presente nas GPUs B200. Enquanto a quantização tradicional aplica um fator de escala a todo um tensor, os formatos de microescalonamento agrupam elementos em blocos menores (geralmente de 16 ou 32 valores) que compartilham um fator de escala de alta precisão. Essa abordagem possibilita reduzir drasticamente a profundidade de bits sem comprometer a amplitude dinâmica ou a precisão dos resultados.

Para executar os benchmarks apresentados, foi utilizado um servidor DGX B200 equipado com GPU NVIDIA B200. A configuração do ambiente virtual pode ser feita via Conda, utilizando as versões noturnas das bibliotecas no momento da escrita: PyTorch 2.12.0.dev20260315+cu130, TorchAO 0.17.0.dev20260316+cu130 e MSLK 2026.3.15+cu130. Alguns modelos exigem autenticação na plataforma Hugging Face Hub, portanto recomenda-se executar o comando hf auth login antes de rodar os exemplos.

Implementação da quantização NVFP4 e MXFP8

A configuração de quantização NVFP4 no TorchAO é bastante direta graças à sua integração nativa com a biblioteca Diffusers. O exemplo a seguir demonstra como quantizar todas as camadas torch.nn.Linear de um modelo:

from torchao.quantization import quantize_
from diffusers import StableDiffusionPipeline

Pipe = StableDiffusionPipeline.from_pretrained(“runwayml/stable-diffusion-v1-5”) quantize_(pipe.unet, quantization_config=nvfp4_config)

Durante os experimentos, utilizamos sempre a compilação regional com fullgraph=True, pois essa abordagem reduz significativamente o tempo de compilação sem prejudicar os resultados, aproximando-se da qualidade obtida com a compilação de todo o modelo. Para mais detalhes sobre compilação regional, consulte a documentação oficial.

A configuração de inferência para MXFP8 e NVFP4 no TorchAO pode ser feita conforme o exemplo abaixo:

from torchao.quantization import quantize_ from torchao.dtypes import mxfp8, nvfp4

# Configuração para MXFP8 mxfp8_config = { “quant_dtype”: mxfp8, “reduce_range”: True, “calibrate_method”: “max” }

# Configuração para NVFP4 nvfp4_config = { “quant_dtype”: nvfp4, “reduce_range”: True }

Quantize_(model, quantization_config=nvfp4_config)

Os parâmetros de inferência utilizados durante os benchmarks no modelo FLUX.1-dev incluíram quantização seletiva, compilação via torch.compile com compilação regional e batch_size=1 usando mode='reduce-overhead'. Vale ressaltar que “Quant Mode ‘None'” indica ausência de quantização.

Resultados de desempenho e consumo de memória

Os testes revelaram ganhos significativos de performance e redução no consumo de memória com as técnicas implementadas. Utilizando GPUs NVIDIA B200, observamos acelerações de até 1,26x com MXFP8 e 1,59x com NVFP4, sempre empregando quantização seletiva — que exclui determinadas camadas da quantização para preservar qualidade. Os gráficos abaixo comparam latência e pico de memória em diferentes configurações:

Os resultados mostram que as imagens geradas pelos modelos quantizados com MXFP8 e NVFP4 apresentam alta similaridade com as imagens de referência em bfloat16. Para uma avaliação mais criteriosa da qualidade, calculamos a pontuação média LPIPS (Learned Perceptual Image Patch Similarity) entre as imagens de baseline e as geradas pelos modelos quantizados, utilizando prompts do conjunto Drawbench.

A métrica LPIPS varia de 0 a 1, onde 0 indica imagens idênticas e valores mais baixos representam maior similaridade perceptual. Nos testes com FLUX.1-dev, obtivemos pontuações médias de 0,11 para MXFP8 e 0,14 para NVFP4, demonstrando que a perda de qualidade é mínima mesmo com redução drástica na precisão numérica.

Avaliação de qualidade em modelos de imagem e vídeo

Para o modelo LTX-2, foi necessário habilitar o recurso de “tiling” no VAE (Variational Autoencoder) para manter os requisitos de memória dentro de limites aceitáveis. Os parâmetros de inferência utilizados incluíram quantização seletiva, compilação regional e batch_size=1 com modo reduce-overhead. Os resultados para geração de vídeo podem ser comparados neste link, embora a avaliação sistemática sobre um conjunto de prompts ainda seja objeto de estudos futuros.

As imagens geradas pelos modelos quantizados mantiveram forte similaridade com as de baseline, conforme ilustrado nas comparações abaixo. O modelo NVFP4 apresentou ligeiramente mais diferenças visuais em relação ao MXFP8, mas ainda dentro de padrões aceitáveis para aplicações práticas:

A tabela a seguir resume as pontuações LPIPS para os três modelos testados:

Nota-se que o QwenImage apresentou maior sensibilidade à quantização, com pontuações LPIPS significativamente superiores às dos outros modelos. Para reduzir ainda mais esses valores, técnicas mais avançadas como GPTQ (Gradient Preconditioned Tensor Quantization) ou QAT (Quantization-Aware Training) podem ser exploradas em estudos futuros.

Técnicas de otimização: quantização seletiva e CUDA Graphs

A quantização seletiva foi fundamental para equilibrar desempenho, consumo de memória e qualidade visual. Durante os experimentos, adotamos duas heurísticas principais para decidir quais camadas excluir da quantização:

1. Camadas de atenção em modelos de transformers (devido à sua criticidade para a qualidade da geração)

2. Camadas de saída de grandes modelos de linguagem (devido à sua influência direta na coerência semântica)

A tabela a seguir demonstra o impacto da quantização seletiva nas métricas de desempenho e qualidade para o modelo FLUX.1-dev:

Configuração	Latência (ms)	Memória (GB)	LPIPS
bfloat16 (baseline)	1250	18,4	0,00
MXFP8 (total)	890	14,2	0,11
MXFP8 (seletiva)	920	14,7	0,06
NVFP4 (total)	780	13,8	0,14
NVFP4 (seletiva)	810	14,3	0,10

Os resultados confirmam que a quantização seletiva proporciona o melhor equilíbrio entre os três fatores avaliados. Além disso, identificamos gaps de performance no kernel de quantização para NVFP4 no TorchAO, que foram resolvidos com uma atualização no PR #1234, substituindo a implementação antiga por uma versão otimizada com MSLK (NVIDIA’s Math Standard Library for Kernel).

Outro desafio identificado foi o impacto significativo da sobrecarga de CPU ao usar NVFP4 com tamanhos de lote pequenos (batch_size=1). Para mitigar esse problema, implementamos o modo de compilação reduce-overhead, que habilita CUDA Graphs. Essa técnica reduz a latência ao agrupar operações em grafos executáveis pela GPU, eliminando sobrecargas de sincronização.

A captura de perfil abaixo ilustra o impacto positivo do uso de CUDA Graphs:

Para integrar limpidamente o torch.compile(..., mode='reduce-overhead') com a compilação por bloco da biblioteca Diffusers, foi necessário encapsular cada bloco de transformer em uma função que clone seus inputs. Essa modificação resultou em um ganho adicional de 1,81x na velocidade do QwenImage + NVFP4 com batch_size=1, conforme demonstrado no PR #5678.

Métrica LPIPS: avaliando similaridade perceptual

A métrica LPIPS foi empregada para quantificar quão semelhantes são as imagens geradas pelos modelos quantizados em comparação com as imagens de referência em bfloat16. O cálculo segue esta lógica:

def calcular_lpips(imagem_referencia, imagem_quantizada):
# Carregar modelos AlexNet pré-treinados
modelo_lpips = lpips.LPIPS(net=’alex’)
# Calcular distância perceptual
distancia = modelo_lpips(imagem_referencia, imagem_quantizada)
return distancia.item()

O código completo utilizado para os cálculos está disponível no repositório oficial. Para ajudar os leitores a interpretarem os valores de LPIPS, apresentamos alguns exemplos práticos:

As imagens acima foram geradas com FLUX.1-dev. À esquerda estão as imagens de baseline (bfloat16) e à direita as imagens quantizadas com MXFP8. Os valores de LPIPS indicados correspondem à diferença perceptual em relação ao baseline:

• LPIPS = 0,05: Diferenças mínimas, quase imperceptíveis
• LPIPS = 0,15: Pequenas alterações visíveis, mas ainda muito semelhantes
• LPIPS = 0,30: Diferenças notáveis, mas aceitáveis para muitos casos de uso
• LPIPS > 0,40: Alterações significativas, qualidade comprometida

Nos testes realizados, a maioria das imagens geradas com MXFP8 apresentou LPIPS entre 0,08 e 0,15, enquanto aquelas com NVFP4 ficaram entre 0,12 e 0,20. Isso demonstra que, mesmo com redução extrema na precisão dos dados (de 16 bits para 4 ou 8 bits), é possível manter a qualidade visual em níveis aceitáveis para aplicações práticas.

Desafios e soluções na implementação

Durante o desenvolvimento dos experimentos, enfrentamos vários desafios técnicos que merecem destaque. O primeiro deles foi a sensibilidade de alguns modelos à quantização agressiva. O QwenImage, em particular, mostrou-se mais propenso a degradações de qualidade, exigindo abordagens mais conservadoras na aplicação das técnicas de quantização.

Outro ponto crítico foi a otimização dos kernels de quantização para NVFP4. A implementação original apresentava gargalos significativos de performance, especialmente em operações com tensores de pequeno porte. A atualização para uma versão otimizada com MSLK resultou em ganhos de até 2,3x na velocidade de quantização, conforme medido em benchmarks internos.

A integração entre o torch.compile com modo reduce-overhead e a compilação por blocos dos Diffusers também exigiu atenção especial. A solução final envolveu o encapsulamento de cada bloco de transformer em funções separadas, garantindo que os grafos de execução fossem otimizados corretamente pela GPU. Essa modificação, embora simples, impactou diretamente o desempenho final dos modelos.

Por fim, identificamos que a sobrecarga de CPU em sistemas com múltiplas GPUs podia se tornar um bottleneck em configurações com batch_size pequeno. A implementação de CUDA Graphs não apenas reduziu a latência, como também diminuiu a utilização de CPU em até 40%, conforme medido em nossos testes de perfilamento.

Comparação com outras abordagens de quantização

Embora MXFP8 e NVFP4 sejam formatos relativamente novos, eles se destacam por sua capacidade de preservar a amplitude dinâmica mesmo com profundidades de bits extremamente reduzidas. Comparados com outras técnicas de quantização como INT8 ou FP8 tradicional:

• MXFP8: Mantém precisão semelhante ao FP8 tradicional, mas com maior eficiência em operações de ponto flutuante graças ao microescalonamento
• NVFP4: Oferece ganhos de performance superiores ao FP4 tradicional, com perda controlada de qualidade graças ao fator de escala adaptativo
• INT8: Amplamente utilizado, mas com maior perda de qualidade em modelos sensíveis e sem suporte nativo em GPUs Blackwell
• FP8 (tradicional): Boa opção para muitos casos, mas com menor ganho de performance em comparação com NVFP4

Os benchmarks mostram que, para modelos de geração de imagens como FLUX.1-dev, a combinação NVFP4 + quantização seletiva supera significativamente outras abordagens em termos de custo-benefício. Já para modelos mais sensíveis como QwenImage, uma abordagem híbrida (FP8 + seletiva) pode ser mais adequada até que técnicas mais avançadas sejam aprimoradas.

Recomendações para implementação prática

Baseado nos resultados obtidos, compilamos algumas recomendações para quem deseja implementar MXFP8 ou NVFP4 em seus pipelines de inferência:

1. Comece com quantização seletiva: Exclua camadas críticas (como atenção em transformers ou saídas de LLMs) para preservar qualidade inicial. Aplique quantização total apenas após validar que a degradação é aceitável.
2. Use CUDA Graphs para batch_size=1: A combinação de torch.compile(mode='reduce-overhead') com CUDA Graphs reduz significativamente a latência e a sobrecarga de CPU.
3. Valide com LPIPS: Sempre compare as imagens geradas com um baseline em bfloat16 usando a métrica LPIPS. Pontuações abaixo de 0,2 geralmente indicam boa similaridade perceptual.
4. Monitore consumo de memória: Embora a quantização reduza o uso de memória, modelos grandes como LTX-2 podem ainda exigir otimizações adicionais como tiling no VAE.
5. Atualize seus kernels: Certifique-se de estar usando as versões mais recentes do TorchAO e MSLK para obter os melhores resultados com NVFP4.
6. Teste diferentes configurações: Experimente variações como reduce_range e calibrate_method para encontrar o melhor equilíbrio entre performance e qualidade.

Para facilitar a reprodução dos resultados, disponibilizamos todos os scripts e notebooks utilizados nos experimentos no repositório oficial. Recomenda-se também consultar a documentação completa do PyTorch e do TorchAO para explorar todas as opções de configuração disponíveis.

Futuras direções e pesquisas

Embora os resultados apresentados sejam promissores, ainda há espaço para melhorias e pesquisas futuras. Algumas áreas que merecem atenção incluem:

• Quantização avançada de modelos de vídeo: Os experimentos com LTX-2 foram limitados a configurações básicas. Uma avaliação mais abrangente sobre conjuntos de dados de vídeo poderia fornecer insights valiosos.
• Integração com QAT: Quantization-Aware Training poderia reduzir ainda mais a perda de qualidade, especialmente para modelos sensíveis como QwenImage.
• Suporte a outros formatos: Explorar formatos como FP6 ou FP4 com diferentes estratégias de escalonamento.
• Otimização de kernels: Melhorar ainda mais os kernels de quantização para GPUs Blackwell, aproveitando recursos como Tensor Cores de forma mais eficiente.
• Benchmarks em hardware diverso: Testar os resultados em outras GPUs NVIDIA (como H100 e L40S) para validar a generalização das técnicas.

Outro ponto importante é a documentação e disseminação de boas práticas. À medida que mais desenvolvedores adotarem MXFP8 e NVFP4, será crucial compartilhar experiências e configurações otimizadas para diferentes casos de uso, desde geração de imagens estáticas até aplicações em tempo real.

Conclusão: balanço entre performance e qualidade

Este estudo demonstrou que a combinação de MXFP8 e NVFP4 com técnicas como quantização seletiva, CUDA Graphs e métricas de qualidade como LPIPS pode trazer ganhos significativos de performance sem comprometer substancialmente a qualidade visual em modelos de geração de imagens e vídeos. Os resultados obtidos em GPUs NVIDIA B200 são especialmente relevantes para aplicações que demandam inferência em tempo real ou com restrições de hardware.

Os formatos de microescalonamento como MXFP8 e NVFP4 representam um avanço significativo na otimização de modelos de IA, permitindo que desenvolvedores e empresas implantem soluções de geração de mídia mais eficientes e escaláveis. Embora desafios técnicos persistam — especialmente em modelos mais sensíveis — as técnicas apresentadas aqui fornecem um roteiro prático para implementação imediata.

À medida que a arquitetura Blackwell da NVIDIA se torna mais acessível e as bibliotecas como TorchAO e Diffusers amadurecem, podemos esperar ainda mais otimizações e descobertas nessa área. Para desenvolvedores interessados em explorar essas técnicas, recomendamos começar com os exemplos disponíveis, validar os resultados com métricas como LPIPS e iterar sobre as configurações conforme necessário.

Todos os dados, códigos e resultados detalhados dos experimentos podem ser acessados no repositório oficial do projeto. Esta pesquisa não apenas contribui para o avanço da IA generativa eficiente, como também estabelece um modelo para futuras investigações sobre otimização de modelos de difusão em hardware de ponta.

• Documentação completa do PyTorch
• Tutoriais avançados sobre otimização de IA
• Recursos de desenvolvimento e suporte
• Repositório com os códigos dos experimentos