FLUX.2 Vulkan Bridge

FLUX.2 Vulkan Bridge runs FLUX.2 Klein 4B through ONNX Runtime WebGPU on machines that have Vulkan/WebGPU available but do not have CUDA or ROCm. The immediate target is the AMD BC-250, an older mining card with useful graphics hardware but limited modern ML runtime support.

The project exports the model into a staged ONNX bundle, serves those files to a headless Chromium process, runs ONNX Runtime WebGPU in JavaScript, and wraps the browser engine with a small local API and CLI.

Export stages, bundle contents, and model-source details live in flux2-onnx-export/README.md. The packaged ONNX WebGPU model files are hosted on Hugging Face at https://huggingface.co/MarkShark2/flux2-klein-4b-onnx-webgpu-q4.

Runtime Shape

flux2-engine.js is the browser-side inference engine. It loads the configured model files, tokenizes and encodes the prompt, denoises latents with the transformer, decodes with the VAE, and returns a PNG. flux2_api.py starts the static/model server and browser runtime and exposes a local generation endpoint. flux2_cli.py is a small wrapper around that API.

The runtime intentionally keeps only one large ONNX session alive at a time. This matters on the BC-250 test machine, which has about 14 GiB RAM and 14 GiB swap.

Model Changes For WebGPU

The working model is deliberately close to the original FLUX.2 Klein graph, but a few changes are required for the ONNX Runtime WebGPU path.

Weight-Only q4

The text encoder and transformer replace learned constant-weight MatMul nodes with ONNX Runtime com.microsoft.MatMulNBits q4 nodes. The default block size is 16. The text encoder defaults to accuracy level 0, while the transformer defaults to accuracy level 4 for the faster WebGPU fp16-accumulate q4 GEMM path. The VAE decoder stays fp16.

The q4 stages handle both ONNX export shapes for linear weights: direct MatMul initializers and MatMul(..., Transpose(weight)) graphs. That second shape appears when constant folding is disabled to keep export memory under control.

ORT q4 Transformer With fp16 Activations

The transformer uses ONNX Runtime q4 packing for learned MatMul weights and keeps transformer activations in fp16 to reduce memory at larger resolutions. RoPE/position math and the fixed attention split-scale constant remain float32 where the WebGPU graph needs that precision.

Fixed Attention Split Scale

The hardest bug was an ONNX Runtime WebGPU scalar issue in the exported attention scale chain. CPU ORT and WebGPU matched through early blocks, q/k/v projections, RoPE composition, and the first attention block. The first major divergence appeared in block 2 attention scores.

The input q/k tensors were correct, but WebGPU computed a dynamic scalar Sqrt/Div/Sqrt chain incorrectly. A sqrt(128) value became about 0.545 instead of 11.313, which inflated the q/k split scale from about 0.2973 to 1.3543. Attention score matrices became roughly 20x too large, and the image turned into colorful noise.

The export now replaces those dynamic attention split-scale chains with a fixed float32 scalar for head dimension 128. This made CPU ORT and WebGPU transformer tensors match again at the tiny-contract level and produced coherent images in the full runtime.

Float32 RoPE And Position IDs

Position IDs and RoPE frequency math are kept in float32. This avoids drift from casting ID-derived position math to fp16 before attention.

Selective fp32 Residual Islands With fp16 Clamps

The transformer keeps fp16 activations everywhere except small fp32 "islands" wrapped around LayerNormalization, every residual Add, and every gate-Mul that feeds a residual Add. Each island casts in to fp32, runs the op in fp32, applies a Clip(+/-65504) for Add/Mul, then casts back to fp16. Islands are never chained in fp32; the residual highway between blocks stays in fp16. This fixes the gate*MLP_out overflow that produced infinities at step 2 of denoising while keeping the per-step memory cost essentially the same as the fully-fp16 build.

Staged VAE Decoder

The fp16 VAE decoder is exported as several smaller ONNX graphs instead of one large session. The exporter splits the decoder into a pre-attention graph, a small attention-chunk graph that is run repeatedly to cover the full mid-attention sequence, and a long sequence of post-attention residual/upsample stages. The runtime creates and releases each session in turn so the WebGPU EP never holds the entire decoder graph in memory at once.

Problems We Hit

Several plausible fixes did not solve the image-quality issue and are worth remembering for similar PyTorch to ONNX to WebGPU ports:

• q8 MatMulNBits is not supported by the tested ONNX Runtime WebGPU path, which supports 2-bit and 4-bit kernels there.
• Changing MatMulNBits accuracy level did not materially change WebGPU output for this model, so the transformer export now defaults to level 4 for throughput.
• Better q4 fitting, symmetric q4, and block-size changes did not fix the broken image by themselves.
• The VAE was not the main problem. Known-latent VAE WebGPU checks looked visually correct.
• Tokenizer and sampler defaults mattered, but they were not the root cause. Klein uses four denoising steps, guidance 1.0, and no guidance embed input.
• Full fp16 transformer diagnostics were too large for this machine and pushed memory/swap too hard to be a practical production path.
• Image-level testing was too coarse. Tensor-level CPU ORT versus browser WebGPU comparisons were what finally isolated the scalar attention bug.

For similar projects, expect small scalar or shape-derived ONNX subgraphs to be just as suspicious as large matrix kernels. A standalone MatMulNBits probe matched CPU and WebGPU, while the full graph was wrong because of a nearby scalar chain.

VAE Memory Blow-Up At 1024x1024

After the transformer image-quality work was finished, full-resolution generation would freeze the BC-250 the moment the VAE decoder ran. Memory dropped from about 12 GiB available to under 100 MiB in a few seconds, swap saturated, and Chromium had to be killed manually. 256 and 512 ran fine. The transformer was already releasing its session before the VAE started, so this was a VAE-only problem.

The path that finally pinned it down was a VAE-only diagnostic, not the full text-encoder/transformer/VAE pipeline. flux2-onnx-export/diag_vae_webgpu.py writes a deterministic latent, starts only the static model server plus Chromium WebGPU, and calls decodeLatent directly. Adding ONNX shape inference to that script before launching Chromium showed the actual cost: the FLUX.2 Klein VAE has a global self-attention block in decoder.attn_1. At 1024x1024 the decoder mid-attention sequence is 16384 tokens. The exported graph materialized the full attention score and softmax tensors:

/decoder/attn_1/MatMul_output_0   FLOAT16 [1,1,16384,16384]  512 MiB
/decoder/attn_1/Softmax_output_0  FLOAT16 [1,1,16384,16384]  512 MiB

Plus Q/K/V, output, and ORT/WebGPU temporaries, that single attention block plans for far more memory than the rest of the decoder combined.

What did not work:

• Chunking the attention math inside one ONNX graph. Replacing the big MatMul/Softmax/MatMul with a chain of per-query-chunk attention nodes plus a final Concat reduced the largest single tensor but ORT/WebGPU still planned the whole chunked graph at once. A 16-input Concat also hit the WebGPU per-stage storage-buffer limit. Adding a binary tree of two-input Concats fixed the storage-buffer issue but did not reduce the planned memory enough.
• Tiled VAE decode in the runtime. This was rejected as a band-aid; the VAE should not need tiling for a 14 GiB machine if the graph is exported sensibly.

What did work was splitting the VAE decoder ONNX into multiple smaller graphs and running them as separate sessions:

• flux2-klein-4b-vae-decoder-pre-attn-fp16.onnx runs everything up to and including the q/k/v projections for the mid-attention block. Its outputs are the residual tensor plus q, kT, and v.
• flux2-klein-4b-vae-decoder-attn-chunk-fp16.onnx is a small graph that takes one query chunk and the full kT/v and returns one chunk of attention output. The runtime creates this session once and runs it repeatedly to cover the 16384-token sequence.
• flux2-klein-4b-vae-decoder-post-attn-fp16.onnx, plus a series of flux2-klein-4b-vae-decoder-post-stageN-fp16.onnx graphs, run the post-attention residual/upsample tail of the decoder. The runtime creates and releases each post stage in turn.

Splitting the decoder into many smaller sessions was the actual fix. Each session creates, runs, and releases on its own, so the WebGPU EP only has to plan and hold one piece of the decoder at a time. The runtime path lives in runVaeDecoderSplitNchw in flux2-engine.js. The graph-side rewrites are in stage 05 (05_export_vae_decoder_fp16.py) and the bundle wiring is in stage 06 (06_package_bundle.py).

Debugging Similar Ports

The fastest path was not changing quantization profiles over and over. The useful path was bisection.

Start with a tiny deterministic contract input and compare the same model call across the reference PyTorch model, ONNX Runtime CPU, and the browser WebGPU runtime. Save raw tensors, not screenshots. If the final output differs, expose internal tensors as temporary graph outputs and bisect by subsystem: embeddings, input projections, position/RoPE math, norms, qkv projections, attention scores, softmax/value products, output projections, MLPs, residuals, and final layer.

When one boundary matches and the next does not, keep narrowing. In this project, the q4 MatMulNBits kernels looked suspicious, but standalone real-weight MatMulNBits probes matched CPU and WebGPU. The first real divergence was an attention score matrix, and the q/k tensors feeding it were already correct. That shifted attention to the scalar scale nodes and exposed the WebGPU Sqrt/Div/Sqrt issue.

For future PyTorch to ONNX to WebGPU/Vulkan ports, keep the production graph boring until evidence says otherwise. Use mixed precision, layer-by-layer dequantization, HQQ-style quantizers, context clipping, and other knobs only as diagnostics. Promote them into the main export only when a tensor comparison proves they are necessary for correctness or stability.

How We Should Have Debugged The VAE From The Start

The VAE blow-up is worth its own debugging note because the early instinct ("VAE weights are tiny, so it must be a transformer leak") sent us in the wrong direction for hours. For future ports, debug a suspected VAE problem in this order:

1. Run a VAE-only diagnostic before anything else. The very first step is flux2-onnx-export/diag_vae_webgpu.py-style isolation: only start the model server and the VAE session, write a deterministic latent, and call the decoder directly. If a VAE-only run reproduces the OOM, the transformer is irrelevant. If it does not reproduce, the bug is upstream.
2. Look at the exported VAE graph statically before re-running anything in WebGPU. Run ONNX shape inference, list the largest inferred tensors, and pay attention to anything proportional to seq*seq. The 512 MiB attention score tensor was visible from a static shape report alone, no Chromium needed.
3. Trust the seq^2 attention rule. The PyTorch VAE may have been written with memory-efficient attention or a fused SDPA call, but torch.onnx.export will lower it to plain MatMul -> Softmax -> MatMul and materialize the full attention matrix. For any image VAE with a self-attention block at full mid-resolution, the exported attention tensor is the first thing to size-check.
4. Do not assume that releasing one ONNX session frees its native/EP memory immediately. Use a fresh process for the VAE diagnostic. Browser/JS-heap counters do not see Dawn/Vulkan/WebGPU EP allocations, so they will look fine even when the system is being thrashed.
5. Add a host-memory watchdog to any large diagnostic. The diagnostic tool should kill its own browser child if MemAvailable drops under a small threshold, before the host needs to be rescued from another terminal. We learned this the hard way after multiple manual pkill chromium rescues.
6. Prefer splitting the graph over tiling at the runtime. Tiling is a runtime band-aid that papers over an expensive ONNX export. Splitting the decoder into pre-attention, repeated attention chunks, and a chain of post-attention stages is more code but it lets each session create/run/release on its own and keeps the WebGPU EP from planning the whole decoder at once.

Run Generation

The runtime loads the packaged model from the Hugging Face cache, downloading it from MarkShark2/flux2-klein-4b-onnx-webgpu-q4 on first use. Set FLUX2_RUNTIME_MODEL_DIR or pass --model-dir only when you want to override that with an explicit local bundle.

python flux2_cli.py \
	--prompt "A small ceramic robot on a wooden workbench, product photo, soft studio lighting" \
	--output outputs/flux2_cli_webgpu.png \
	--width 256 --height 256 --seed 123 --num-steps 4

The CLI auto-starts the local API if it is not already running.