Skip to content

idol: the FLUX.2 VAE image encoder on CUDA (stacked on #1774, #1767; not yet run on a GPU) - #1859

Draft
jayleaton wants to merge 16 commits into
ashhart:mainfrom
jayleaton:idol-vae-cuda
Draft

jayleaton wants to merge 16 commits into
ashhart:mainfrom
jayleaton:idol-vae-cuda

Conversation

@jayleaton

Copy link
Copy Markdown
Contributor

Stacked on #1774 and #1767 (Idol's VAE encoder on Metal, and the denoiser on CUDA). This branch is #1774 with #1767 merged in. Its own change is 7f68bb85f. Review after them.

Compiles and passes host tests; never run on a GPU; GB10 qualification pending. The Sparks are busy, so nothing in this PR has run on an NVIDIA device.

What this changes

Idol's FLUX.2 VAE image encoder on CUDA, the CUDA side of #1774. For #1757. Details: the "VAE encoder" part of docs/recipes/idol-cuda.md.

  • zig/kernels/cuda/idol_vae.cu and idol_vae.cuh (glue flags): every convolution as an fp32 implicit GEMM of FMAs (same-size 3x3, stride-2 3x3 padded right and below, 1x1), conv_in, two-pass GroupNorm, GroupNorm with SiLU, bias epilogues, the attention's softmax, and the condition rows with the NaN-keeping clamp.
  • The mid block's q, k, v, out and P V run on cuBLASLt in true fp32 (CUBLAS_COMPUTE_32F, fp32 inputs, no TF32), through idol: Idol's denoiser on CUDA (stacked on #1766; not yet run on a GPU) #1767's cuda_gemm.zig.
  • vae_config.zig holds the geometry, config check, bf16 rounding and the gate, shared with Metal. tf-idol-vae is one driver (zig/tests/idol_vae_run.zig) with thin Metal and CUDA mains; the CUDA build step is idol-vae-cuda.
  • It admits sm_121 and sm_120 and refuses other GPUs by name. The posterior's noise stays an input.

Precision

fp32 throughout on the SDK's bf16 weights, as on Metal, and the same gate against float64. No Winograd and no tensor cores: the plain fp32 implicit GEMM keeps every sum an fp32 FMA chain, so it needs none of the Metal port's float64-made transforms. Whether it meets the float64 gate on a device is for GB10 qualification.

Kernel contract

tf-cuda-test idol-vae-contract: 16 forms through gates 1-4 and eleven controls built from idol_vae.cuh with one planted flaw each. It is in the full contract's family sets and contract_limits.zig. Gate-5 records (idol_vae_cuda_records.zig) are nvcc 13.0's for sm_80, 86, 89, 120 and 121. No SUMMARY line yet: it needs a GPU.

Receipt

  • Linux box (no GPU), nvcc 13.0:
    • idol_vae.cu and the controls compile for sm_80, 86, 89, 120 and 121.
    • zig build test-cuda-host passes; zig build (Linux) exit 0, 367 s; zig build fatbins -Dnvcc=... -Dsm=121,120: 220 of 220 steps.
    • tf-idol-vae for CUDA builds on the merged tree (173 of 173 steps).
  • Pre-existing host-suite failures: on the Linux box, zig build test-cuda-host failed the same 8 tests with and without this change: 6 in recall_store_test, plus recall_routes ("the record's state and size…") and late_system_test ("GET and DELETE /v1/recall…"). The tree with only idol: Idol's denoiser on CUDA (stacked on #1766; not yet run on a GPU) #1767 merged into the Metal branch fails the same 8 (4,117 of 4,134 passed, 8 failed). The recall store reads back nothing ("found 0") on a box whose disk was nearly full. None of them touches Idol. The final run on the merged CUDA tree (idol-e2e-cuda) passed with no failures (282 of 282 steps).
  • Metal, M5 Ultra: zig build test -Dcpu=apple_m1: 589 of 589 steps, 8,958 of 9,011 tests pass, 53 skipped. The pre-change and new tf-idol-vae on the VAE case: rows.f32 byte-identical, the JSON identical apart from timings, against float64 our worst row 5.144e-5 and the SDK's 2.729e-2: pass.
  • zig fmt --check build.zig zig: clean. lean_check.py: 424, the same as main.
  • Not run: every kernel's output on a device, the contract gates and controls on a device, cuBLASLt's fp32 plans at these shapes, case parity, timing and memory. tools/zig/idol_gb10_qualify.sh (in the CUDA end-to-end PR) runs this component's contract and case on a Spark with the rest; its section in docs/recipes/idol-cuda.md lists the case files to copy from the Mac.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (No token path.)
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (n/a.)
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt. (CUDA: registered; no SUMMARY line until a GPU runs it.)
  • Prompt processing is no slower than the last release. (Not touched.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Every pipeline of idol.metal and idol_attention.metal is registered: id_linear's eight forms (plain and GELU),
id_split's two, id_gemm32, id_norm (plain, affine, modulated), id_norm32 (affine, modulated), id_norm_split,
id_repo, id_qk, id_gate_add, id_embed, id_cond, id_bias32 and the attention's four forms. Each output is checked
at idol-sdk's fp16 points as the range its fp32 steps allow; the NAX projections against the tensor unit's
16-deep step bound, the attention against attention_ref's formula with P at the tensor unit's fp32. Inputs at the
large served shapes repeat at periods prime to every tile, so float64 stays cheap and every output is checked.

19 controls: rounding points dropped or added (gate 3), fp16 sums in id_norm32, id_repo, id_qk, id_embed and
id_gemm32, and the NAX fp32 path in id_gemm32's place (gate 4), and the attention's tail keys unmasked.
Gate 5 holds the 24 launches, each reserved by its source.

Known finding: id_qk gives a head holding a NaN +-inf in its other outputs where F.normalize gives NaN, as
Metal's max drops the NaN norm that clamp_min keeps.
Port Idol's denoiser from Metal to CUDA under the same precision contract:
fp16 block linears via cuBLASLt with fp32 sums and one rounding after the
bias (tf_idol_bias), LayerNorm32 to fp16 with fp16 modulation, true-fp32
RePo/condition/output projections (CUBLAS_COMPUTE_32F), and a head-64 fp16
attention with fp32 softmax and accumulation in its own image (idol.cu,
idol_attention.cuh), leaving prefill_attention.cu and its SASS untouched.

- family modules zig/src/families/idol/cuda*.zig; the model refuses any GPU
  but sm_120/sm_121 by name with a log line; timing via the sys shim
- tf-idol on CUDA (idol-cuda step) sharing the case driver with Metal
  (zig/tests/idol_run.zig)
- cuBLASLt Linear.initTyped for fp16 and fp32 inputs; bf16 path unchanged
- kernel contract: idol-contract (gates 1-4 over every idol.cu kernel and
  the attention, three controls dropping an fp16 rounding), joined to the
  family sets; nvcc 13.0 records for sm_80/86/89/120/121; image in
  contract_limits
- host tests for the family and the contract's range references
…f-idol-vae and its case

Builds the denoiser's feature2 [4101, 128] from the preprocessed picture as idol-sdk's ImageEncoder does: the
encoder, quant_conv, the posterior's sample with the SDK's noise as an input, the 2 x 2 patches, the batch norm's
statistics and five zero rows. fp32 throughout on the SDK's bf16 weights (FLUX.2-klein-4B, or FLUX.2-dev rounded):
Winograd 3x3s, implicit-GEMM downsamples and 1x1s, two-pass GroupNorm, single-head attention.

idol_case.py vae records the SDK's MPS rows, its noise and a float64 ruler; tf-idol-vae gates worst row and
1 - cosine within twice the SDK's distance from float64 (ours 5.1e-5 and 5.0e-11 against 2.7e-2 and 5.1e-5), and
repeat bytes. All eight kernels register with the contract: 13 forms pass gates 1-4, 12 controls red.
…d on GB10)

Port Idol's image encoder (ashhart#1774) to CUDA under the same float64 gate, fp32
throughout on the SDK's bf16 weights. Compiles and passes host tests; never
run on a GPU; GB10 qualification pending.

- idol_vae.cu / idol_vae.cuh: every convolution as an fp32 implicit GEMM of
  FMAs (same-size 3x3, stride-2 3x3 padded right and below, 1x1; no Winograd,
  no tensor cores, no TF32), conv_in, two-pass GroupNorm, GroupNorm + SiLU,
  bias epilogues, the attention's softmax and the condition rows (NaN-keeping
  clamp); glue flags. The mid block's q/k/v/out and P.V run on cuBLASLt in
  true fp32 (cuda_gemm.zig, CUBLAS_COMPUTE_32F), inside free picture tensors.
- cuda_vae*.zig: encoder, launchers, weights; admits sm_120/sm_121 alone and
  refuses others by name; posterior noise stays an input.
- vae_config.zig: geometry, config check, bf16 rounding and the gate shared
  with Metal; tf-idol-vae's driver shared too (idol_vae_run.zig), the Metal
  tool's output unchanged; tf-idol-vae on CUDA (idol-vae-cuda step).
- kernel contract: idol-vae-contract, 16 forms through gates 1-4 and eleven
  controls built from idol_vae.cuh with one planted flaw each, joined to the
  family sets; nvcc 13.0 records for sm_80/86/89/120/121; image in
  contract_limits; host tests for geometry, scratch, admission and references.
- contract_controls.cu: restore the closing brace of tf_ctl_pocket_in_fmax
  that the MOSS-Image CUDA commit dropped (as on main), so the controls build.
@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @jayleaton. Plain fp32 implicit GEMMs and true-fp32 cuBLASLt with no TF32 keep the CUDA side simple. This waits on #1774 and #1767. It has the same registry.zig line-limit point as #1767, and it needs the idol-vae-contract SUMMARY line plus the float64 gate on a device. We can run it on a GB10 if that helps.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants