Running the H3 diffusion transformer on INT8 kernels

kLLM includes INT8 kernels for the diffusion transformer (DiT) of MiniMax H3, for GPUs whose INT8 tensor cores run at twice the bf16 rate, such as the A100 and A40. On one A40 they reduced the steady-state time per denoising step by 25% (24.56 s to 18.38 s) on plain H3. This page lists the requirements, describes how the kernels are enabled and how they work, states how they interact with sealing, and reports the measured speed and numerical drift.

Important

The INT8 kernels are a kernel optimization, separate from sealing. They change the numerics of the DiT, with the measured drift reported below, and they are not part of the cost or the claims of a sealed model. The measured speed-up was obtained on plain H3, not on a sealed pipeline.

Requirements

  • The int8 extra, which installs triton>=3 and comfy-kitchen>=0.2.39:

    $ pip install -e '.[seal,serve,int8]'
    

    The linear-layer kernels need comfy-kitchen; the attention kernel needs Triton only.

  • An SGLang build with diffusion support. The kernels replace private functions of SGLang’s H3 implementation and were measured with SGLang 0.5.21; other versions are not checked at run time.

  • A GPU with INT8 tensor cores. The attention kernel is written for sm_80 and sm_86 (A100- and A40-class GPUs). No capability check is made at start.

Usage

With a sealed H3 pipeline

kllm serve --int8 [PARTS] applies to sealed MiniMax H3 pipelines served with --key. PARTS is attention, linear or all; a bare --int8 means all.

$ kllm serve models/MiniMax-H3/FL2VA-sealed --key --int8 --tp 1 --no-recipe -- \
    --dit-layerwise-offload true --layerwise-offload-components text_encoder vae

kllm serve prints the replaced parts, for example int8    attention, linear in the DiT blocks. For any other model it exits with --int8 applies to sealed MiniMax H3 pipelines (served with --key).

With plain H3

For a plain H3 pipeline, start sglang serve directly with the kernel directory on PYTHONPATH and KLLM_INT8 set. python -m kllm.kernels.h3 prints the directory:

$ PYTHONPATH=$(python -m kllm.kernels.h3) KLLM_INT8=all sglang serve \
    --model-path models/MiniMax-H3 --model-variant fl2va --num-gpus 1 \
    --dit-layerwise-offload true --layerwise-offload-components text_encoder vae

KLLM_INT8 accepts attention, linear, attention,linear, all or 1. An unknown value is reported as KLLM_INT8: unknown part(s) [...]; use attention, linear.

How it works

The kernel directory contains a sitecustomize.py that every Python process started with that PYTHONPATH imports, including SGLang’s worker processes. When KLLM_INT8 is set, it patches SGLang’s H3 DiT module right after the module is imported, and then chains to the next sitecustomize on the path. Only the 50 DiT blocks change: the text token refiner, the encoders and the VAE keep SGLang’s kernels, and so do Ulysses and Ring sequence-parallel runs and attention backends other than FlashAttention.

Attention (kllm.kernels.int8_attention, Triton)

After the approach of SageAttention. Q and K are quantized per token and head, K after removing its per-channel mean, which leaves softmax unchanged. P is quantized per row against the running block maximum that the online softmax computes anyway. V is quantized per channel after removing its mean, which is added back at the end. Both products run in INT8 with INT32 accumulation, and scaling and conversion are folded into one fused multiply-add per element.

Linear layers (kllm.kernels.w8a8)

The qkv_proj, out_proj, fc1 and fc2 layers of the DiT blocks run through comfy-kitchen’s INT8 GEMM. Their bf16 weights are quantized again on every call, per output channel after a 256-wide group-wise Hadamard rotation, at a cost of about 3 ms per H3 block. Quantizing at call time keeps the checkpoint, layer-wise offload and the online AdaLN cache unchanged. SGLang’s own --quantization kitchen_int8 quantizes at load time, which needs the whole model on the GPU and ran out of memory on a 48 GB GPU. Activations must be bf16.

Interaction with sealing

Part

Under perm

Under orth

Attention

Commutes with the signed permutations of heads, RoPE pairs and V: sealed and plain attention agree up to an occasional last-bit rounding of P·V. Verified by GPU tests on random tensors, which do not run in continuous integration.

Does not commute to the last bit: V/O take a dense per-head rotation, which changes V’s per-channel scales and means.

Linear layers

Not established. The Hadamard rotation does not commute with the signed permutations of the DiT’s residual stream or MLP, so sealed and plain weights quantize differently.

Not established.

No measurement of a sealed pipeline under --int8 has been published.

Measured results

Plain H3 (FL2VA), one NVIDIA A40, 768p, 4 s, layer-wise offload, SGLang 0.5.21, run with benchmarks/h3_int8.sh (which also enables --minimax-h3-adaln-online true and --warmup-mode off). Steady-state seconds per denoising step:

stock SGLang

--int8 linear

--int8 (all)

s / step

24.56

21.0

18.38 (−25%)

The linear run was on another GPU of the host, where stock SGLang took 24.9 s per step.

Per layer, on the Q, K and V of H3 (32704 tokens, 56 heads): FlashAttention 257 ms, SageAttention2 214 ms (1.42% error at the outlier-heavy last layer), the kLLM kernel 220 ms (1.01% error; FlashAttention itself 0.17%), all against an fp32 reference. Linear layers: 209 ms in bf16 against 122 ms in INT8. These per-layer figures have no script in the repository.

Fidelity against stock SGLang from the same noise: relative L2 error of the final latents (video / audio) and frame PSNR after 49 steps.

after 1 step

after 7 steps

after 48 steps

PSNR

SDPA instead of FlashAttention (noise floor)

1.7% / 2.5%

20.3% / 13.6%

46.7% / 21.2%

17.3 dB

--int8 attention

2.6% / 4.0%

20.1% / 25.7%

27.2% / 15.1%

23.4 dB

--int8 linear

3.2% / 5.6%

30.7% / 45.5%

37.5% / 17.9%

19.5 dB

--int8

4.4% / 8.4%

39.9% / 54.3%

39.0% / 17.3%

18.8 dB

INT8 moves the trajectory more than a kernel swap in the early steps, the audio stream most. After the full 49 steps every variant ended at least as close to stock as SDPA did, with the same scene, subject and lighting.

Reproducing

$ pip install -e '.[int8]'
$ benchmarks/h3_int8.sh ~/models/MiniMax-H3
$ STEPS=8 GPU=3 VARIANTS="plain int8" INT8=attention benchmarks/h3_int8.sh ~/models/MiniMax-H3

The argument is the H3 repository root, the directory that holds FL2VA/. The script serves plain FL2VA on one GPU, from the same noise, as stock SGLang, with the INT8 kernels, and with PyTorch SDPA in place of FlashAttention. It prints the total denoising time and the end-to-end time of each run, then latent drift, frame PSNR and audio SNR against stock, and writes OUT/grid.png. It reports total denoising time rather than steady-state seconds per step. See Benchmarking plain and sealed serving for its settings.

See also