[diffusion] optimization: INT8 Linear + pluggable DiT attention backends for MiniMax-H3 on consumer-level GPUs (#34581)

This commit is contained in:
WenhaoZhang
2026-08-18 21:44:39 +08:00
committed by GitHub
parent 0065fbfae1
commit 63d783bbe0
10 changed files with 616 additions and 45 deletions
+88 -2
View File
@@ -68,7 +68,9 @@ H200, and H100. Resident is the latency-oriented default; FSDP reduces DiT
weight residency at the cost of per-block parameter collectives. On H200 it
also selects the verified 2-node cross-node topology. **Online
Quantization** appears only on B200 and B300. AMD keeps its resident AITER
recipe, while RTX 5090 uses its dedicated layerwise-offload profile.
recipe, while RTX 5090 uses its dedicated layerwise-offload profile. A
single 24 GB card (RTX 4090) uses the same offload knobs plus online
`kitchen_int8`; that recipe is documented below rather than in the picker.
import { Deployment } from "/src/snippets/_deployment.jsx";
import { config } from "/src/snippets/configs/MiniMaxAI/minimax-h3.jsx";
@@ -168,6 +170,34 @@ transfer overhead. This exact recipe was validated on
2× RTX 5090 (32 GB each) and a 377 GiB host; use a 384 GiB-class machine. The
latency and memory comparison is collected in the benchmark section below.
For a single 24 GB consumer card (RTX 4090), stream the DiT and text encoder
and quantize DiT linear layers online with `kitchen_int8`. Keep `vae` out of
`--layerwise-offload-components`: putting the VAE decoder in layerwise
offload re-streams about 9 GiB on each of 167 decode tiles. Default
attention stays `fa` (exact). Approximate backends are opt-in; see
[Attention Backends](/docs/sglang-diffusion/attention_backends#sage-then-sol-hybrid).
Install `comfy-kitchen` first (`pip install comfy-kitchen`).
```bash 1×RTX 4090 24GB
sglang generate \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--quantization kitchen_int8 \
--attention-backend fa \
--performance-mode memory \
--layerwise-offload-components dit,text_encoder \
--dit-offload-prefetch-size 1 \
--dit-layerwise-resident-layers 0 \
--enable-torch-compile false \
--prompt "A cat walking on a sunny beach, gentle waves." \
--save-output
```
The same flags work on `sglang serve`. Drop `--quantization` for the BF16
baseline; everything else stays identical. GPU peak stays about 18 GB
either way because streaming offload is set by the offload buffers and VAE
decode, not the weight dtype.
The first launch downloads the model through the selected Hub. If the Hugging
Face repository requires authentication, export a Hugging Face token in the
server environment.
@@ -668,7 +698,7 @@ listed hardware and topology; it is not inherited by a similar GPU family.
| Tensor parallelism | Verified: B200 TP2 + Ulysses4; H100 TP2 + Ulysses2 and TP4 + Ulysses1 | `--tp-size` may be combined with Ulysses when the TP-local head count remains divisible by the Ulysses degree. On 4×H100, TP2 + Ulysses2 is the measured speed default. |
| FSDP inference | Verified: 4× B200 and 4× H100 + Ulysses4 | Preserves H3's mixed BF16/FP32 parameter policy. B200 completed the exact eager comparison; H100 completed consecutive real requests at about 57 GB peak memory per GPU. |
| Resident components | Verified: B200, H200, 4×H100 with TP, and 1/2/4/8× MI300X and MI355X | This is the recommended single-request latency path when the complete workload fits. |
| CPU and layerwise offload | Verified: 2× RTX 5090 TP2 | The measured lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. This status applies only to the listed topology. |
| CPU and layerwise offload | Verified: 2× RTX 5090 TP2; 1× RTX 4090 24 GB | The 5090 lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. The 4090 recipe streams DiT and the text encoder with zero resident DiT layers and **omits `vae`** from `--layerwise-offload-components`. |
| Breakable CUDA graph | Verified: B200 Ref2VA, opt-in | Matching eager output was observed for the captured signature, without a measured speedup. Re-capture for other shapes and reference sets. |
| `torch.compile` | Measured: H200, opt-in | Steady-state benefit was below measurement noise, while startup increased and numerical output changed. Do not use it for consistency ground truth. |
@@ -738,6 +768,40 @@ The picker exposes this option only on the B200 and B300 topologies used for
real H3 validation runs.
</Warning>
On a single 24 GB card, use `kitchen_int8` instead of FP8. It quantizes the
four GEMMs per DiT block online from the Hub BF16 weights (data-free, no
calibration) and dispatches them through `comfy_kitchen.int8_linear`.
Quantization happens after H3's grouped `qkv` reorder, so do not load an
externally pre-quantized INT8 checkpoint here.
```bash Command
sglang generate \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--quantization kitchen_int8 \
--attention-backend fa \
--performance-mode memory \
--layerwise-offload-components dit,text_encoder \
--dit-offload-prefetch-size 1 \
--dit-layerwise-resident-layers 0 \
--enable-torch-compile false \
--prompt "A cat walking on a sunny beach, gentle waves." \
--save-output
```
`fa` keeps exact attention. For a faster, approximate DiT path, use
`--attention-backend sol_attn` with
`--attention-backend-config dense_backend=sage_attn,dense_steps=10` and
`--component-attention-backends text_encoder=torch_sdpa,transformer=sol_attn`.
See [Quantization](/docs/sglang-diffusion/quantization#kitchen-int8-online-quantization)
and [Attention Backends](/docs/sglang-diffusion/attention_backends#sage-then-sol-hybrid).
<Warning>
`kitchen_int8` changes Linear numerics. `sol_attn` / `sage_attn` also change
the attention algorithm. Neither is a consistency ground-truth mode. The
BF16 path is unchanged when `comfy-kitchen` is not installed.
</Warning>
The Qwen3-VL text encoder can be replaced independently of the DiT. To reduce
its resident memory, point the text-encoder component at the serialized FP8
checkpoint used in validation:
@@ -799,6 +863,7 @@ the configurations with collected measurements:
| H100 | 4× TP2 + Ulysses2 resident | 4× TP4 + Ulysses1; 4× FSDP + Ulysses4 |
| MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs |
| RTX 5090 | 2× TP2 + layerwise offload | — |
| RTX 4090 24 GB | 1× layerwise offload + `kitchen_int8` | Approximate attention backends are opt-in |
### B300 precision and encoder placement
@@ -991,6 +1056,27 @@ peak per GPU.
| prefetch 2, resident 20 | 43.37 s | 78.06 s | 27.5 GiB | No measurable gain |
| Ulysses2, prefetch 2, resident 10 | Did not reach warmup | — | — | Rejected |
### RTX 4090 24 GB single-GPU run
One RTX 4090 D 24 GB completed the 1344×768, 107-frame, 20-NFE T2VA
workload (euler, `torch.compile` and step caching disabled) with DiT and
text-encoder layerwise offload. Same process: load → warmup (seed 0) →
timed (seed 42); only the timed pass is reported. GPU peak stayed about
18 GB.
| Config | Timed e2e | Denoise | vs BF16 | PSNR vs BF16 |
| --- | ---: | ---: | ---: | ---: |
| BF16 + FlashAttention | 405.6 s | 370.2 s | 1.00× | — |
| `kitchen_int8` + FA | 303.3 s | 273.7 s | 1.34× | 24.81 dB |
| `kitchen_int8` + `sol_attn` | 223.9 s | 203.9 s | 1.81× | 24.44 dB |
| `kitchen_int8` + `sage_attn` | 174.9 s | 154.2 s | 2.32× | 23.51 dB |
| `kitchen_int8` + Sage→Sol hybrid | 163.8 s | 143.1 s | 2.48× | 23.04 dB |
`kitchen_int8` + FA changes Linear numerics only. The `sol_attn` /
`sage_attn` / hybrid rows also change the attention algorithm, so speed
and pixel fidelity rank in opposite orders there. Default remains
`kitchen_int8` + `fa`.
### AMD Instinct task and scaling runs
The AMD recipes keep the released BF16/FP32 precision policy and use AITER
@@ -398,6 +398,12 @@ Some backends require additional configuration. You can pass these parameters vi
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Layer indices kept dense, e.g. `0,1` or `0-2`.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`0,1`</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`dense_backend`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`str`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Backend used for the dense prefix: `fa` (default) or `sage_attn`. `sage_attn` is approximate.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`fa`</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`kv_splits`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>`int | str`</td>
@@ -626,6 +632,25 @@ sglang generate \
Component keys match pipeline module names from `model_index.json`, such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`.
### Sage then Sol hybrid
`sol_attn` keeps the first `dense_steps` steps dense. Set
`dense_backend=sage_attn` to run that prefix on SageAttention and the tail on
Sol sparse attention. Keep the text encoder on `torch_sdpa`:
```bash
sglang generate \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--attention-backend sol_attn \
--attention-backend-config dense_backend=sage_attn,dense_steps=10 \
--component-attention-backends text_encoder=torch_sdpa,transformer=sol_attn \
--prompt "A cat walking on a sunny beach, gentle waves." \
--save-output
```
Both `sage_attn` and `sol_attn` are approximate. The default DiT backend remains `fa`.
### Using Sliding Tile Attention (STA)
```bash
@@ -88,6 +88,16 @@ description: "Configure SGLang diffusion behavior with environment variables."
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>true</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Use Run:AI model streamer for model loading</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_KITCHEN_INT8_MAX_ROWS</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>8192</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Max activation rows per <code>kitchen_int8</code> fused GEMM. Set <code>0</code> to disable row splitting.</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_KITCHEN_INT8_MIN_SPLIT_N</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>8192</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>Minimum output features before <code>kitchen_int8</code> splits large-M GEMMs. Narrow outputs stay as a single call.</td>
</tr>
</tbody>
</table>
@@ -79,6 +79,14 @@ backend.
<td>MXFP4: <code>aiter</code> on ROCm</td>
<td>MXFP4 requires ROCm and MI350+ (gfx95x). Weights quantized at load time, activations quantized to <code>fp8</code> / <code>mxfp4</code> dynamically.</td>
</tr>
<tr>
<td><code>kitchen_int8</code> (online quantization)</td>
<td>Unquantized BF16/FP16 checkpoint</td>
<td><code>--quantization kitchen_int8</code></td>
<td>MiniMax-H3 (validated on 1× RTX 4090 24 GB)</td>
<td><code>comfy-kitchen</code></td>
<td>Data-free INT8 ConvRot at load time via <code>comfy_kitchen.int8_linear</code>. Layers whose input dim is not divisible by the group size stay in BF16. Requires Turing+ (SM75).</td>
</tr>
<tr>
<td><code>fp8</code> (offline quantization)</td>
<td>Quantized transformer component folder, or safetensors with <code>quantization_config</code> metadata</td>
@@ -260,6 +268,50 @@ sglang generate \
```
**Note:** Requires `aiter` package with MXFP4 kernel support
### Kitchen INT8 Online Quantization
`kitchen_int8` quantizes DiT linear weights online from the stock BF16
checkpoint. Forward uses the fused `comfy_kitchen.int8_linear` op (rotation,
dynamic per-row activation quantization, INT8 GEMM, dequant, and bias).
Install the optional dependency first:
```bash
pip install comfy-kitchen
```
```bash
sglang generate \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--quantization kitchen_int8 \
--attention-backend fa \
--performance-mode memory \
--layerwise-offload-components dit,text_encoder \
--dit-offload-prefetch-size 1 \
--dit-layerwise-resident-layers 0 \
--enable-torch-compile false \
--prompt "A cat walking on a sunny beach, gentle waves." \
--save-output
```
Quantization runs after the model's weight loaders, so MiniMax-H3's grouped
`qkv` reorder is already applied. Layers whose input dim is not divisible by
the group size (256) stay in BF16 instead of failing the load; H3's AdaLN
projections take that path.
<Warning>
`kitchen_int8` is approximate and is not a consistency ground-truth mode.
The BF16 path is unchanged when `comfy-kitchen` is not installed. See the
[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#7-runtime-feature-recipes)
for the 24 GB offload recipe, including why `vae` must stay out of
`--layerwise-offload-components`.
</Warning>
Large-M GEMMs (`rows > 8192` and `out_features >= 8192`) are split by rows so
the fused kernel stays on the data-parallel CUTLASS config. Override the
thresholds with `SGLANG_KITCHEN_INT8_MAX_ROWS` and
`SGLANG_KITCHEN_INT8_MIN_SPLIT_N`.
### Skipping Layers
By default, online quantization quantizes every linear layer in