NVFP4 en difusión: la misma cuantización, un motor completamente distinto

0. Misma técnica, arquitectura completamente distinta

En la página de Gemma-4, cuantizar a NVFP4 fue casi un one-liner: un flag (--qformat nvfp4_mse) que le pasás a una herramienta de NVIDIA, y vLLM se encarga del resto en runtime (--quantization modelopt_fp4). Acá no hay vLLM ni ModelOpt de por medio — es torch puro, invocando el kernel de matmul cuantizado de Blackwell directamente. Ver esta versión "a mano" es lo que te muestra qué es lo que vLLM te estaba escondiendo cuando usás el flag fácil.

Conviene tener presente desde el principio qué cambia y qué no respecto de las dos páginas anteriores. No cambia el formato: E2M1 en bloques de 16 con escala FP8, los mismos 16 valores representables, la misma cuenta de 4,5 bits efectivos por peso. Cambia todo lo que rodea al formato: quién carga los pesos, quién cuantiza la activación, quién elige el kernel, quién decide qué capas se salvan, y — lo más nuevo — un modelo que no genera un token por pasada sino que refina un tensor entero, muchas veces, antes de entregar algo.

1. Un diffusion transformer, en una frase

Un LLM decoder-only genera una vez: calcula la probabilidad del siguiente token y listo, un forward pass por token. Un modelo de difusión genera iterando: arranca de ruido puro y, en varias pasadas (típicamente 20-50, a veces mucho menos — sección 4), predice cuánto ruido hay que restarle a la imagen actual para acercarla un poco más al resultado final. Flux2 usa bloques transformer para esa predicción — de ahí "diffusion transformer" — pero no hay nada autoregresivo ni causal en juego: no genera token por token, refina la imagen entera en cada paso.

Para lo que importa en esta página, el detalle que sí es igual a un LLM: por dentro, la mayoría del cómputo son las mismas operaciones — proyecciones lineales grandes (attn.to_q/k/v, capas de feed-forward) — y son exactamente esas matrices las que se cuantizan a NVFP4.

De la U-Net al DiT

Si hiciste la Parte 2 de fast.ai, el modelo de difusión que tenés en la cabeza es el de Stable Diffusion: una U-Net, una red convolucional en forma de U que baja la resolución del latente en la mitad izquierda, la vuelve a subir en la derecha, y conecta ambas mitades con skip connections — con bloques de atención intercalados en las resoluciones más bajas. Es una arquitectura heredada de segmentación de imágenes, y funciona muy bien, pero tiene un problema práctico: los sesgos convolucionales y la estructura de resoluciones hacen que no escale tan limpiamente como un transformer cuando le das más cómputo.

El Diffusion Transformer (DiT, Peebles & Xie, 2022) reemplaza la U-Net por lo mismo que ya conocés del Capítulo 1: el latente se corta en parches (como ViT hace con una imagen), cada parche es un "token", y una pila de bloques transformer — atención + feed-forward — procesa la secuencia de parches. El condicionamiento (el paso de difusión t y, en los modelos texto→imagen, el texto) entra modulando cada bloque. El hallazgo del paper es que la calidad escala de forma predecible con los GFLOPs del transformer, igual que en LLMs — y eso es lo que llevó a que Flux, Wan y casi todos los modelos generativos recientes sean DiTs y no U-Nets.

Flux agrega un detalle de diseño que explica los nombres que vas a ver en el §2: sus primeros bloques son de doble flujo (double_blocks) — el texto y la imagen viajan por dos pilas separadas que solo se mezclan en la atención — y después vienen bloques de flujo único donde texto e imagen ya son una sola secuencia. Cuando el checkpoint dice img_attn.qkv, está nombrando la proyección Q/K/V del lado imagen de uno de esos bloques dobles.

Lo que este cambio de arquitectura le hace a la inferencia es lo que importa para el resto de la página: si el DiT es "un transformer que corre 20-50 veces sobre la misma secuencia", entonces el costo se multiplica por los pasos, la memoria se domina por los pesos del transformer (no por un KV cache que crece), y cada palanca de velocidad tiene que actuar o sobre el costo de cada pasada (cuantización, kernels de atención) o sobre la cantidad de pasadas (destilación, §4).

2. NVFP4 sin vLLM: torch._scaled_mm a mano

blackwell_utils.py define BlackwellLinear, un reemplazo a medida para nn.Linear que hace exactamente lo que ModelOpt+vLLM hacen por vos en la página anterior, pero visible línea por línea:

# Pesos: (out, in // 2) uint8 — dos valores FP4 empaquetados por byte
self.register_buffer("qweight", torch.empty((out_features, in_features // 2), dtype=torch.uint8))
# Escala: un float8_e4m3fn cada 16 valores — mismo esquema de bloque que Gemma-4
self.register_buffer("weight_scale", torch.empty((out_features, in_features // 16), dtype=torch.float8_e4m3fn))

Hacé la cuenta con esas dos formas, porque es la misma que la de la página de Gemma-4 vista desde el otro lado: para una matriz de out × in, los pesos ocupan out × in / 2 bytes (4 bits cada uno) y las escalas out × in / 16 bytes (una por bloque de 16, a lo largo de la dimensión de entrada). Total: out × in × (0,5 + 0,0625) bytes = 4,5 bits por peso, contra 16 en BF16. Las escalas son el 11 % del tamaño — no son gratis, pero son lo que hace que 4 bits alcancen.

En el forward, la activación se cuantiza dinámicamente a FP4 (a diferencia de los pesos, que ya vienen cuantizados desde el checkpoint), y el matmul cuantizado corre directo en hardware:

res = torch._scaled_mm(
    qact.view(torch.float4_e2m1fn_x2),
    self.qweight.view(torch.float4_e2m1fn_x2).t(),
    scale_a=act_scale,
    scale_b=self.weight_scale,
    out_dtype=torch.bfloat16,
)

float4_e2m1fn_x2 es exactamente el E2M1 de la página de Gemma-4 (1 signo, 2 exponente, 1 mantisa), con el "x2" del nombre del dtype indicando que cada elemento de PyTorch empaqueta dos valores FP4 en un byte — la razón por la que el tensor de pesos declara la mitad de columnas (in_features // 2). torch._scaled_mm es la función que sabe desempaquetar eso y multiplicar en el tensor core de Blackwell sin pasar por una conversión intermedia a BF16.

Leé la llamada como la ecuación que es: out = (A · s_a) × (B · s_b)ᵀ, donde A es la activación en FP4, B los pesos en FP4, y las escalas se aplican por bloque dentro del kernel. Dos cosas que la llamada deja explícitas y que vLLM esconde: (1) las escalas viajan como tensores aparte, no dentro de los pesos — el kernel las multiplica sobre la marcha; (2) out_dtype es BF16: el resultado de acumular productos de 4 bits se devuelve en 16 bits, porque el siguiente paso (normalización, residual, la próxima capa) no sabe ni quiere saber de FP4. Cada capa cuantizada es una isla de 4 bits con costas en BF16.

Por qué la activación se cuantiza en runtime

Es la distinción W4A16 / W4A4 de la página de Gemma-4, vista en código. Los pesos son fijos, así que se cuantizan una vez, offline, con toda la calma del mundo. La activación es distinta en cada forward — y en un modelo de difusión, distinta en cada paso de denoising, porque el latente cambia de "ruido puro" a "casi imagen" — así que sus escalas solo se pueden calcular cuando la activación existe: qact y act_scale se producen en cada llamada, bloque de 16 por bloque de 16, con el mismo absmax o criterio que use la implementación. Esto es lo que compra el tensor core FP4 × FP4; el precio es que un outlier de activación en un solo paso de difusión puede degradar ese paso, y el ejercicio 1 del §8 te muestra cuánto.

El mapeo de nombres parece un detalle administrativo hasta que pensás qué pasa si está mal: no hay error. Una matriz de to_q cargada con el tercio de to_k tiene exactamente la misma forma y el mismo dtype — el modelo corre, y genera imágenes que no se parecen a nada. Es el mismo tipo de "basura silenciosa" de la página de vLLM, con una diferencia: acá el eval no es un GSM8K con respuesta exacta, es mirar la imagen. Por eso conviene tener un golden (misma semilla, mismo prompt, salida en BF16 de referencia) y comparar la salida NVFP4 contra ese, como parte del pipeline y no como inspección a ojo.

3. ¿Cómo fine-tuneás algo que ya está en 4 bits?

Con un modelo en BF16, aplicar un LoRA es casi trivial: sumás la matriz de bajo rango entrenada (A·B) directamente a los pesos originales, o la dejás separada y la sumás en runtime. Con pesos empaquetados en NVFP4 eso deja de ser posible — no hay forma de "sumarle" una corrección continua a un valor que ya está discretizado a uno de 16 niveles por bloque.

Pensalo con números: en un bloque cuya escala es 0,01, los pesos cuantizados solo pueden valer 0, ±0,005, ±0,01, ±0,015, ±0,02, ±0,03, ±0,04 o ±0,06. Un LoRA típico corrige cada peso en cantidades del orden de 0,001 — un décimo de la distancia entre niveles. Si "mergeás" y volvés a cuantizar, casi todas esas correcciones se redondean a cero, y las pocas que cruzan un umbral saltan un nivel entero: el LoRA queda entre borrado y amplificado, al azar. Por eso la única alternativa correcta es no tocar los pesos cuantizados.

benchmark_lora_nunchaku.py lo deja explícito en el código: el método merge() de la capa LoRA-sobre-NVFP4 directamente levanta NotImplementedError("Merging into NVFP4 relies on fusion at runtime."). La única salida es fusionar el cómputo del LoRA dentro del kernel en cada forward — nunca modificar los pesos cuantizados en sí. nunchaku (la librería de la que viene svdq_quantize_w4a4_act_fuse_lora_cuda) hace justamente eso: cuantiza la activación y aplica la corrección de LoRA en la misma pasada CUDA.

En una ecuación: en vez de y = (W + BA)·x con un W ya redondeado, el kernel calcula y = W_4bit·x + B·(A·x), donde el término de la derecha corre en 16 bits — es chico, porque A tiene rango r (decenas) contra los miles de la matriz completa — y se suma al resultado del matmul cuantizado antes de devolver. "Fusionar" significa que ese cálculo pasa dentro del mismo kernel que cuantiza la activación y hace el matmul FP4, sin que la activación tenga que salir a memoria y volver; si no, el LoRA costaría un lanzamiento de kernel extra y una lectura extra de la activación por capa.

nunchaku es el motor del paper SVDQuant, y vale la pena saber qué resuelve además del LoRA, porque conecta con el §2: cuantizar activaciones a 4 bits en un DiT es difícil por sus outliers. SVDQuant primero "migra" esos outliers de la activación a los pesos (escalando canal por canal), y después le saca a los pesos una componente de bajo rango con SVD que se queda en 16 bits — la misma forma que un LoRA — dejando un residuo mucho más fácil de cuantizar a 4 bits. Que el LoRA del usuario y esa rama de bajo rango tengan la misma forma es lo que hace natural fusionarlos en el mismo kernel.

4. Video: lo mismo, con una dimensión más

Tu configuración de producción para Wan 2.1 (image-to-video) en LightX2V agrega, sobre la misma base de NVFP4, un conjunto de técnicas que no tienen equivalente en un LLM de texto:

Por qué "una dimensión más" cambia tanto

Un DiT de video es el mismo transformer del §1 con parches en tres dimensiones — alto, ancho y tiempo — sobre un latente que un VAE 3D ya comprimió también en el eje temporal. El problema es la cantidad de tokens: donde una imagen tiene, digamos, unos miles de parches, un clip de decenas de frames tiene decenas de miles, y la atención completa escala con el cuadrado de esa cantidad. Por eso en video el kernel de atención deja de ser un detalle y pasa a ser el costo dominante — y por eso un kernel que cuantiza la atención misma (SageAttention: Q y K en pocos bits, con corrección para que el softmax no se rompa) rinde acá mucho más de lo que rendiría en un LLM decodificando de a un token.

Destilación de pasos, en una frase cada una

Un modelo de difusión se entrena para dar un paso chico desde cualquier nivel de ruido; generar exige encadenar muchos. Destilar pasos es entrenar un segundo modelo (o afinar el mismo) para que un paso grande suyo reproduzca lo que el original hacía en varios. Las tres familias que vas a encontrar:

• Consistency models / LCM: el modelo aprende a mapear cualquier punto de la trayectoria de denoising directamente a su origen limpio, con la restricción de que dos puntos de la misma trayectoria den la misma respuesta ("consistencia"). LCM aplica la idea a modelos latentes como Stable Diffusion, y llega a 2-4 pasos.
• Destilación adversarial: además de imitar al maestro, un discriminador penaliza las salidas que no parecen imágenes reales — recupera nitidez que la destilación pura tiende a perder, al costo de un entrenamiento más delicado (es la familia de los modelos "turbo").
• Destilación de guidance: el estudiante aprende a producir directamente la salida que el maestro daba con classifier-free guidance, así no hace falta la segunda pasada sin condicionamiento. Es la razón por la que un checkpoint destilado se sirve con guidance_scale = 1.0: la guía ya está adentro de los pesos.

Lo que las tres tienen en común, y que importa para esta página: ninguna cambia la arquitectura ni el formato de los pesos — un DiT destilado a 4 pasos se cuantiza a NVFP4 exactamente igual que uno de 50. Se apilan: 4 pasos en vez de 50 es ~12× menos pasadas; sin CFG, cada pasada cuesta la mitad; NVFP4 hace cada pasada más barata todavía. Las tres palancas son independientes y multiplicativas.

Por qué el text encoder se queda en BF16

Tres razones, y las tres generalizan. Primero, se ejecuta una sola vez por prompt, no una vez por paso: el embedding de texto se calcula al principio y se reutiliza en los 4 (o 50) pasos del DiT. Cuantizarlo ahorraría memoria pero casi nada de tiempo — y de hecho, una vez calculado el embedding, el encoder entero se puede sacar de la GPU. Segundo, su error se propaga a todo lo demás: un embedding levemente mal es una condición levemente mal en cada uno de los pasos, sin posibilidad de corregirse. Tercero, es la capa que traduce entre modalidades — texto a espacio vectorial — que es exactamente la clase de capa que Gemma-4 también excluye (embed_tokens, lm_head). El mismo criterio para CLIP en image-to-video: es la traducción de la imagen de partida a condición, se corre una vez, y se queda en BF16.

Offload por bloque

Cuando ni la cuantización alcanza, queda mover pesos entre la RAM del sistema y la GPU durante el forward. La granularidad es la decisión: mover el modelo entero de a uno (útil para alternar text encoder → DiT → VAE, que nunca corren a la vez) o mover bloque por bloque del transformer — mientras el bloque n calcula, el n+1 viaja por PCIe hacia la GPU y el n−1 vuelve a RAM. Con bloques grandes y regulares, como los de un DiT, la transferencia se solapa casi por completo con el cómputo; con capas chicas y heterogéneas (una U-Net), el costo fijo de cada transferencia se come la ventaja. Lo que hay que saber para usarlo bien: no toca la precisión (los pesos van y vienen intactos), el límite es el ancho de banda de PCIe, y necesita RAM de sistema suficiente para alojar lo que no está en la GPU. Es una palanca de memoria, ortogonal a la de precisión (NVFP4) y a la de pasos (destilación).

5. Lo que generaliza

Sacando los nombres propios, lo que esta página deja para cualquier proyecto que cuantice o sirva un modelo generativo — de texto, imagen o video:

6. Resumen

1. NVFP4 es un formato de pesos — no le importa si el modelo es un decoder autoregresivo o un diffusion transformer; el mismo E2M1 en bloques de 16 aplica a los dos.
2. Un DiT reemplaza la U-Net por una pila de bloques transformer sobre parches del latente; casi todo su cómputo son matrices lineales grandes, repetidas en cada paso de denoising — por eso cuantizarlas rinde tanto.
3. Integrarlo sin un framework que lo resuelva por vos expone lo que ModelOpt+vLLM esconden: empaquetado de bits, matmul cuantizado explícito con escalas por bloque como tensores aparte, activación cuantizada en runtime, y el problema de mapear convenciones de nombres entre checkpoints de distinto origen.
4. Fine-tunear con LoRA sobre pesos en 4 bits no se puede hacer por fusión de pesos — se fusiona el cómputo en el kernel, en cada forward; SVDQuant/nunchaku ya usan una rama de bajo rango en 16 bits para absorber outliers, y el LoRA viaja por la misma rama.
5. Video agrega memoria (offload por bloque), velocidad de sampling (destilación, sin CFG), atención cuantizada (SageAttention) y un truco de post-proceso (interpolación de frames) que no tienen contraparte en texto. Son palancas independientes: precisión, cantidad de pasadas y memoria.
6. El text encoder y CLIP se quedan en BF16 por la misma razón que embeddings y lm_head en Gemma-4: corren una vez, condicionan todo, y traducen entre modalidades.

7. Ponete a prueba

Respondé antes de abrir cada respuesta. Si fallás una, volvé a la sección indicada.

8. Para practicar

import torch
torch.manual_seed(0)

POS = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
NIVELES = torch.cat([-POS.flip(0)[:-1], POS])     # los 16 valores de E2M1 (el -0 no aporta)

def cuantizar(x, escala):
    """x: (..., 16) bloques; escala: (..., 1). Devuelve la reconstrucción en float."""
    y = x / escala
    idx = (y.unsqueeze(-1) - NIVELES).abs().argmin(-1)   # nivel más cercano
    return NIVELES[idx] * escala

def escala_absmax(b):
    return b.abs().amax(-1, keepdim=True).clamp_min(1e-8) / 6.0   # 6 = máximo de E2M1

def escala_mse(b, factores=torch.linspace(0.3, 1.0, 36)):
    base = escala_absmax(b)
    mejor_err = mejor_esc = None
    for f in factores:                              # barre escalas más chicas que absmax
        esc = base * f
        err = ((cuantizar(b, esc) - b) ** 2).mean(-1, keepdim=True)
        if mejor_err is None:
            mejor_err, mejor_esc = err, esc
        else:
            gana = err < mejor_err
            mejor_err = torch.where(gana, err, mejor_err)
            mejor_esc = torch.where(gana, esc, mejor_esc)
    return mejor_esc

def mse(w, elegir_escala):
    return ((cuantizar(w, elegir_escala(w)) - w) ** 2).mean().item()

print(f"{'outlier ×':>10} {'MSE absmax':>12} {'MSE búsqueda':>13} {'mejora':>7}")
for factor in (1, 3, 6, 12, 30):
    w = torch.randn(4096, 16) * 0.02            # 4096 bloques de 16 "pesos"
    w[:, 0] *= factor                          # un outlier por bloque
    a, m = mse(w, escala_absmax), mse(w, escala_mse)
    print(f"{factor:>10} {a:12.3e} {m:13.3e} {a/m:6.2f}×")

def params_dit(bloques, d, expansion_ff=4):
    attn = 4 * d * d                    # q, k, v, o
    ff = 2 * d * (expansion_ff * d)      # subida y bajada
    return bloques * (attn + ff)         # sin normas ni modulación (despreciables)

cfg = {"bloques": 40, "d": 5120, "expansion_ff": 4}   # geometría de ejemplo, ~12,6B
p = params_dit(**cfg)
text_encoder_bf16_gb = 4.7 * 2       # p. ej. un encoder de ~4.7B parámetros × 2 bytes ≈ 9.4 GB
for nombre, bits in [("BF16", 16), ("FP8", 8), ("NVFP4 (4.5 efectivos)", 4.5)]:
    gb = p * bits / 8 / 1e9
    print(f"{nombre:<22} DiT {gb:5.1f} GB  + text encoder BF16 {text_encoder_bf16_gb:.1f} GB = {gb + text_encoder_bf16_gb:5.1f} GB")
print(f"parámetros del DiT: {p/1e9:.1f} B")

9. Fuentes y para seguir

• Peebles & Xie — Scalable Diffusion Models with Transformers (arXiv:2212.09748) — el paper del DiT: parches del latente como tokens, condicionamiento por modulación, y la curva de escala con GFLOPs que explica por qué Flux y Wan no son U-Nets.
• Black Forest Labs — repositorio de Flux — el código de referencia de la familia Flux: bloques de doble flujo y flujo único, y la convención de nombres (double_blocks.N.img_attn.qkv) que tu mapeo traduce.
• Li et al. — SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models (arXiv:2411.05007) — el paper detrás de nunchaku: migración de outliers, rama de bajo rango en 16 bits, y por qué el LoRA se fusiona en el kernel en vez de mergearse.
• Zhang et al. — SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration (arXiv:2410.02367) — la primera versión de la familia de kernels que cuantiza la atención misma; leelo para entender qué se cuantiza (Q, K) y cómo se corrige antes del softmax.
• Luo et al. — Latent Consistency Models: Synthesizing High-Resolution Images with Few-Step Inference (arXiv:2310.04378) — destilación por consistencia sobre modelos latentes: la explicación más directa de cómo se llega de 50 pasos a 4.
• Huang et al. — RIFE: Real-Time Intermediate Flow Estimation for Video Frame Interpolation (arXiv:2011.06294) — el interpolador de frames del post-proceso: estima flujo óptico intermedio con una red chica, sin tocar el modelo de difusión.
• Wan-Video — Wan 2.1 (repositorio) — el modelo de video del §4: arquitectura, VAE 3D, text encoder y CLIP para image-to-video, y la configuración de referencia contra la que comparar la tuya.