ComfyUI-fni8

INT8 dp4a acceleration for diffusion transformers (DiTs) in ComfyUI, targeting NVIDIA Volta (sm_70).

This is a ComfyUI custom-node pack. It runs the compute-heavy DiT (Flux, SD3.5, Qwen-Image, Z-Image, Wan, LTX, and others) through the fni8 int8 W8A8 dp4a GEMM and bidirectional int8 FlashAttention instead of the fp16 tensor cores. It patches ComfyUI's existing model ops rather than forking, so it composes with the sampler, LoRA, and VAE stack.

Contents

• Why this exists
• How it works
• Pre-quantized DiTs on Hugging Face
• Supported architectures
• Text encoders & VAE (fp vs. int8)
• Multi-GPU
• Install
• Publishing to the ComfyUI Registry
• Testing
• Status
• The fni8 family

Why this exists

On the target fleet (V100 / CMP 100-210, GV100 silicon) the fp16 tensor cores are firmware-crippled, but the integer __dp4a path is healthy at about 46 TOP/s. Existing quantized-DiT engines all require newer hardware: Nunchaku needs sm_75 (Turing), and the Ideogram/TensorRT int8 kernels need sm_80 (Ampere). int8-dp4a DiT inference on sm_70 is an unoccupied niche, and this fills it. int8 also halves weight HBM and read bandwidth, which matters on a bandwidth-starved fleet.

How it works

The approach follows the ComfyUI-GGUF ops-patch pattern, with one difference: GGUF dequantizes the weight back to floating point and matmuls on the (here, crippled) tensor cores, whereas this pack keeps the weight int8 and runs a true W8A8 dp4a GEMM.

• UnetLoaderFNI8 loads a .fni8-quantized DiT as a ComfyUI MODEL with model_options={"custom_operations": FNI8Ops}. The .fni8 container is shared with the fni8 LLM stack, so one weight format and one converter serve both.
• FNI8Ops is a comfy.ops.manual_cast subclass whose Linear.forward quantizes the activation per row and calls fni8.linear (dp4a). A per-layer SQNR gate falls back to the fp32 CUDA cores when needed, never to the gimped fp16 tensor cores.
• FNI8AttentionPatch swaps DiT self-attention (bidirectional, no causal mask) for fni8.attn_int8_fwd(causal=False) via set_model_attn1_replace, or an optimized_attention wrapper for Flux/SD3-style DiTs.

Pre-quantized DiTs on Hugging Face

Ready-to-load .fni8 DiT weights (int8 and int4) are published under huggingface.co/jajmangold. Each is a linked quantization, so it also shows under its source model's Quantizations tab (Flux, Qwen-Image, Z-Image, Wan, LTX, and more).

Supported architectures

comfyui_fni8/arch.py registers a broad set. Adding one is a data entry, since the pack patches ComfyUI's own DiT definitions.

• Image: Flux.1, Flux.2, SD3 / SD3.5, Qwen-Image / Qwen-Image-Edit, Z-Image (Turbo), PixArt-Sigma, Sana, Chroma, HiDream, Lumina-2, AuraFlow, Hunyuan-DiT (Ideogram is documented but its weights are closed).
• Video (arch registered; loading WIP): Wan 2.1 / 2.2, LTX-Video, HunyuanVideo, Mochi, CogVideoX, Cosmos. ⚠️ Wan/LTX .fni8s don't load end-to-end yet — ComfyUI's detect_unet_config doesn't recognize their diffusers layout (#30); "registered" means the arch entry exists, not that a checkpoint loads.

Z-Image Turbo is the first validation target (6B, 8-step, single joint-attention stream).

Text encoders & VAE (fp vs. int8)

The .fni8 archive only ever carries a quantized DiT — loader.py's quantize_dit_state_dict never touches a text encoder or VAE. The VAE (a convolutional autoencoder, not a transformer) is numerically load-bearing and small relative to the DiT's weight/HBM cost — there's nothing to gain from quantizing it, so the VAE always stays fp.

The text encoders (T5-XXL, Qwen, umT5, CLIP) ARE transformers — the same architectures the DiT int8 dp4a GEMM + attention already accelerate — so they can optionally run int8 dp4a too. FNI8ComponentLoader (in comfyui_fni8/nodes.py) defaults its clip_precision input to "fp" (the safe default); switching it to "int8_dp4a" patches the TE's Linear layers with FNI8Ops (the same int8 dp4a GEMM the DiT uses) via custom_operations, quantized lazily from the loaded fp weight on first forward — no separate quantized TE file needed. Norms and embeddings aren't nn.Linear, so they stay fp automatically; TE self-attention (the softmax(QK^T)V compute, not the QKV/out-proj linears) also stays fp for now — a low-risk starting point since it runs once per generation, not once per denoise step like the DiT's.

UnetLoaderFNI8 loads the int8 DiT; FNI8ComponentLoader loads the matching text encoder + VAE for a given arch, using arch.py's per-family text_encoder/clip_type/vae metadata so you don't have to guess which CLIPLoader/DualCLIPLoader type an arch needs. It rejects a .fni8 file passed into either slot rather than failing deep inside comfy.sd with a confusing parse error (the TE's int8 path never reads a .fni8 — it quantizes the original fp checkpoint on load). Get the TE/VAE weights themselves from each model's own release (its HF model card / ComfyUI's model docs) — this repo only publishes quantized DiTs, per Pre-quantized DiTs on Hugging Face above.

clip_type is a best-effort hint matching comfy.sd.CLIPType's member name (the same string CLIPLoader/DualCLIPLoader's type dropdown expects); it's left blank for archs without a confirmed native type (usually because support landed in ComfyUI after this table was last checked), and FNI8ComponentLoader raises a clear error rather than silently guessing when that happens — fall back to ComfyUI's own CLIPLoader node and pick the closest type by hand.

| Arch | Kind | Text encoder (fp) | clip_type | VAE (fp) | |---|---|---|---|---| | flux1 | image | CLIP-L + T5-XXL (dual) | flux | Flux 16-ch AE | | flux2 | image | Mistral-Small-3.1-24B (unconfirmed) | — | Flux.2 16-ch AE | | sd3 / sd35 | image | CLIP-L + CLIP-G + T5-XXL (triple) | sd3 | SD3 16-ch VAE | | qwen_image / qwen_image_edit | image | Qwen2.5-VL-7B | qwen_image | Qwen-Image 16-ch VAE | | zimage | image | Qwen3-4B | — (too new to confirm) | Z-Image VAE (dedicated) | | pixart | image | T5-XXL | pixart | SDXL VAE (4-ch) | | sana | image | Gemma-2 (2B/9B) | — | DC-AE (32x) | | chroma | image | T5-XXL | chroma | Flux 16-ch AE (shared) | | hidream | image | CLIP-L + CLIP-G + T5-XXL + Llama-3.1-8B (quad) | hidream | Flux-style 16-ch VAE | | lumina2 | image | Gemma-2-2B | lumina2 | Flux-style 16-ch VAE | | auraflow | image | Pile-T5-XL | — | SDXL-style 4-ch VAE | | hunyuan_dit | image | bilingual CLIP + T5 | hunyuan_dit | SDXL VAE (4-ch) | | wan21 / wan22 | video | umT5-XXL | wan | Wan causal 3D VAE | | ltx_video | video | T5-XXL (LTX-2: Gemma-3) | ltxv | LTX 3D causal VAE | | hunyuan_video | video | Llava-Llama-3 (MLLM) + CLIP-L | hunyuan_video | HunyuanVideo causal 3D VAE | | mochi | video | T5-XXL | mochi | Mochi causal 3D VAE | | cogvideox | video | T5-XXL | — | CogVideoX causal 3D VAE | | cosmos | video | T5-XXL | cosmos | Cosmos causal tokenizer VAE |

ideogram is omitted (closed weights — no known TE/VAE to document). This table lives as data in comfyui_fni8/arch.py (DiTArch.text_encoder/.clip_type/.vae); it's covered by test_arch_entries_document_text_encoder_and_vae in tests/test_core.py.

Multi-GPU

comfyui_fni8/multigpu.py encodes the fleet verdict, which matches the LLM side: per-layer communication is unusable, per-step and per-handoff communication survives. The supported strategies are CFG-parallel (cond/uncond on 2 GPUs, roughly free 2x), component-parallel (encoder / DiT / VAE on separate GPUs, zero denoise-loop traffic, as in ComfyUI-MultiGPU), and PipeFusion (displaced-patch pipeline, stage-boundary activations only). Avoid xDiT USP and tensor-parallel, whose per-layer collectives are dead on a 250 MB/s link. DisTorch block-offload is a fit-only fallback (bandwidth-bound weight streaming), not a distribution strategy.

Install

git clone https://github.com/jajmangold/ComfyUI-fni8 ComfyUI/custom_nodes/ComfyUI-fni8
# fni8 must be importable (the sm_70 kernels + .fni8 format):
pip install -e /path/to/fni8

Requires a Volta (sm_70) GPU and fni8 built against its pinned CUDA 12.9 / torch 2.10+cu129 stack. A runtime arch gate raises a clear error on other GPUs.

Publishing to the ComfyUI Registry

This pack carries registry.comfy.org metadata under [tool.comfy] in pyproject.toml (PublisherId, DisplayName), so it can be published as a node pack and installed through ComfyUI-Manager. Publishing is a manual step, never automatic on push/merge.

docs/publish-workflow.yml is a template for a workflow_dispatch-only GitHub Actions release workflow that runs comfy-cli. It lives outside .github/workflows/ because the bot account that maintains this repo's automation can't write workflow files itself (no workflow OAuth scope); copy it to .github/workflows/publish.yml yourself (drop its header comment) to enable a one-click CI publish. Until then, publish locally:

pip install comfy-cli

# bump the [project] version in pyproject.toml first — the registry rejects a
# re-publish of a version it has already seen.

# get an API key for the "fni8" publisher at https://registry.comfy.org
# (Publisher dashboard -> API Keys), then:
export REGISTRY_ACCESS_TOKEN=<your registry API key>
comfy node publish

To publish from CI instead of locally: add the same key as this repo's REGISTRY_ACCESS_TOKEN secret (Settings -> Secrets and variables -> Actions), then trigger the workflow by hand — gh workflow run publish.yml or the Actions tab. It is not on an automatic trigger, so a version bump alone will not publish; you still choose when a release goes out.

Testing

tests/test_core.py covers the pure numerics (int8 dp4a linear, bidirectional attention, the .fni8 round-trip, the arch/multi-GPU registries) against random tensors — no ComfyUI or real weights needed.

Random tensors can't catch the failure mode that actually matters here: an fp16/int8 overflow that silently produces a finite-looking but constant (black) image. Catching that needs a real checkpoint through a real ComfyUI process, gated behind the comfy_e2e pytest marker and run in Docker:

docker compose build          # ComfyUI (pinned) + fni8 (built for sm_70) + this pack
docker compose run --rm e2e   # per-arch matrix: loads each published .fni8 DiT via
                               # UnetLoaderFNI8, runs one denoise step, asserts the
                               # latent is finite + non-constant

Needs the weights archive mounted at /mnt/24tb/fni8-forge/weights (see docker-compose.yml) and a Volta/CMP (sm_70) GPU; each test skips cleanly without either. See docker/Dockerfile and tests/e2e/, and docs/e2e-coverage.md for the per-arch pass/fail coverage table (issue #12).

Status

Validated so far: the int8 dp4a linear and bidirectional attention numerics (cos >= 0.99 vs fp16), the .fni8 DiT round-trip, the arch and multi-GPU logic, and the ops-patch structure against the real ComfyUI API (FNI8Ops subclasses, patch seams, node contract). End-to-end (real ComfyUI process + real weights + Volta GPU): Z-Image, FLUX.1-dev, and Qwen-Image/Qwen-Image-Edit load and denoise correctly (see docs/e2e-coverage.md); Wan2.2 and LTX-Video's published .fni8s can't load yet — ComfyUI's own model detection doesn't recognize their diffusers-format checkpoint layout (same doc, xfail with the repro). Not affiliated with ComfyUI or the model authors.

The fni8 family

Part of a three-repo stack targeting the same sm_70 fleet:

• fni8: the dp4a kernels (int8 W8A8/W4A8 GEMM and int8 FlashAttention) and the .fni8 weight format this pack consumes.
• fni8-serve: W8A8 LLM inference server.
• ComfyUI-fni8 (this repo): diffusion DiT acceleration in ComfyUI.