Skip to content

idol: BiRefNet background removal on CUDA (stacked on #1787, #1767; not yet run on a GPU) - #1861

Draft
jayleaton wants to merge 31 commits into
ashhart:mainfrom
jayleaton:idol-birefnet-cuda
Draft

jayleaton wants to merge 31 commits into
ashhart:mainfrom
jayleaton:idol-birefnet-cuda

Conversation

@jayleaton

Copy link
Copy Markdown
Contributor

Stacked on #1787 and #1767 (Idol's background removal on Metal, and the denoiser on CUDA). This branch is #1787 with #1767 merged in. Its own change is 05e5ed31f. 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

BiRefNet's forward on CUDA, the CUDA side of #1787. For #1757. Details: the "Background removal (BiRefNet)" part of docs/recipes/idol-cuda.md.

  • zig/kernels/cuda/idol_birefnet.cu and idol_birefnet.cuh (all fp32, glue flags): tf_br_conv for every conv and linear, tf_br_norm, the Swin window attention tf_br_window, tf_br_resize, tf_br_tiles, tf_br_conv3, the deformable conv tf_br_deform, tf_br_pool, tf_br_gate and tf_br_head.
  • The host preprocessing (Pillow-exact LANCZOS, erosion, crop, composite, picture decode) is shared with Metal, unchanged.
  • tf-idol-birefnet is one driver (zig/tests/idol_birefnet_run.zig) with thin Metal and CUDA mains; the CUDA build step is idol-birefnet-cuda. It admits sm_121 and sm_120 and refuses other GPUs by name.

Precision

fp32 compute on the fp16 weights, as on Metal, gated against float64. tf_br_conv sums each output in ordered 256-term blocks, each a sequential fp32 FMA chain: on Metal one running sum over the squeeze's 51,840 terms put it at twice torch's distance. I did not use cuBLASLt here because its fp32 summation order is not documented. That may cost speed on the GB10; qualification will show how much. The deformable conv samples as torchvision's deform_conv2d (modulation 2·sigmoid, weighted sum not fused).

Kernel contract

tf-cuda-test idol-birefnet-contract: 24 kernel forms with 7 controls (conv in fp16 blocks, conv's ReLU through fmaxf, norm with an fp16 sum, the window attention with fp16 P, deform with fp16 columns, pool with an fp16 sum, resize in fp16). Gate-5 records (idol_birefnet_cuda_records.zig) are nvcc 13.0's for sm_80, 86, 89, 120 and 121.

JPEG on Linux

libjpeg-turbo 3.1.4.1 builds on Linux from the official tarball (checksum verified) into a private prefix, for TENSORFOLD_LIBTURBOJPEG. That box had no NASM, so its SIMD was off. I have not confirmed the JPEG loader picks up the Linux library.

Receipt

  • Linux box (no GPU), nvcc 13.0:
    • The image and the controls compile for sm_80, 86, 89, 120 and 121, exit 0.
    • test-idol-birefnet-cuda-host: 36 of 36. zig build (Linux): exit 0. zig build fatbins -Dnvcc=... -Dsm=121,120: exit 0.
    • zig build test-cuda-host: 4,154 of 4,171 passed, 9 skipped, 8 failed; the base fails the same 8.
  • 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 test-golden -Dcpu=apple_m1: 593 of 593 steps. The pre-change and new tf-idol-birefnet on all six cases (avocado, WebP, opaque in synthetic and real mode, RGBA, JPEG): logs and JSON identical apart from timings, all pass. The Windows cross-build of tf-idol-birefnet.exe and the host tests succeeds.
  • 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 controls on a device, the weight upload and the BN fold, case parity against the SDK and float64, the alpha and picture gates on CUDA, repeats, timing and memory, Windows at runtime. 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.
The octree decoder (4 blocks) and Gaussian decoder (16 blocks) run natively on
Metal from the released checkpoint. Host-side pieces: systematic resampling with
a pluggable RNG (native PRNG, or the SDK's recorded draws), compaction, leaves
and jitter, activations, the PLY writer and the camera JSON. tf-idol-decode
decodes a latent and gates it against an idol_case.py `decode` case.

New kernels live in idol_decoder.metal: the point embedding, K/V and head
packing, an fp32 attention for <= 8 query rows, and a copy of the idol
attention whose P.V uses an fp16 hi/lo split. This removes the -2^-12 per-pass
bias that fp32 P picks up on the tensor unit.

Gates are measured against fp32, allowing twice the worst of: the SDK's own
distance, the fp16 floor, and seeded nudged SDK reruns. Further gates:
per-block teacher-forced checks, counts identical given the SDK's
probabilities and draws, activated fields, the PLY's columns, the camera, and
byte-identical repeats. The comparators are covered by plant-a-mismatch unit
tests.

idol_case.py gains `decode`, which pins decoder.py, gaussian.py and camera.py by
sha256 and records levels, draws, counts, fields, the PLY and timings.
metal_idol_decoder*.zig registers dec_embed, dec_half, dec_kv, dec_heads,
dec_attention_few and dec_attention_{00,01,10,11} as "metal idol decoder
...", against float64 at idol-sdk's fp32 and fp16 points. Fifteen controls
each go red at their gate: an added or dropped rounding point, fp16 sums,
P's lo half dropped (caught on heads built so O is P's fp16 remainder
alone), keys past Lk unmasked, an input written, and the two NaN rules
below undone. Gate-5 records are in idol_records.zig; the common Lib is now
generic over a source and its launchers.

Gate 1 found two real rule breaks, fixed with finite outputs unchanged
(the 262k and 32k decode cases give byte-identical PLYs and gate lines):
dec_heads clamped its norm by max, dropping a NaN; and the tiled attention
gave NaN for an infinite value, because an fp16 lo of 0 or of the opposite
sign met the infinity. The lo half now meets the values clamped to fp16's
finite range (about 2% at 262k, 7% at 1M).

docs/recipes/idol.md gains a Decoders part with the precision, gates,
run-a-case commands, measurements and the contract SUMMARY.
…idol_case.py dino

The denoiser's feature1 [4101, 1280] from the SDK's 1024 x 1024 image, as idol-sdk's ImageEncoder makes it:
transformers' DINOv3ViTModel in bf16, its last_hidden_state, then F.layer_norm in fp32. Weights from the ungated
timm/vit_huge_plus_patch16_dinov3.lvd1689m (DINOv3 License), rounded to bf16 at load. Patch conv as one GEMM, class
token and four registers, 2D RoPE on the patches (transformers' fp32 periods, tables bit-equal to torch's), 32 blocks of
20 heads of 64 with LayerScale and SwiGLU, rounded where torch on MPS rounds (measured): one rounding a linear or norm,
two for SiLU times up and LayerScale with the residual, three for the rotary. The attention rounds P to bf16 itself.

idol_case.py dino pins idol-sdk's conditioning.py, transformers 5.22.1's DINOv3 files and the weights by sha256, and
writes the SDK's feature1 on MPS, a float32 CPU ruler, every block's rows and each block in float32. With --official it
proves timm's weights Meta's: 550 tensors byte for byte, feature1 bit for bit.

On the avocado: feature1 worst row against float32 0.1353 (torch 0.1686), 1 - cos 5.37e-5 (5.57e-5); every block on
torch's own input within the 2x gate; repeat byte-identical; 0.120 s an encode against the SDK's 0.261 s on MPS.
…ngelog line

Every pipeline of idol_dino.metal is registered: dn_linear_0 and dn_linear_1, dn_norm, dn_final, dn_heads,
dn_scale_add, dn_swiglu and dn_attention, each checked at transformers' bf16 points. 11 controls: a bf16 rounding
added or dropped (gate 3), fp16 sums in dn_norm and dn_final and P V rounded to bf16 each key block (gate 4), the
attention's tail keys unmasked. Gate 5 holds the 8 launches, each reserved by its source.

SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 11 controls: 11 red as planted, 0 missed; 0 limit problems
# Conflicts:
#	CHANGELOG.md
#	docs/kernel-contract.md
#	docs/recipes/idol.md
#	tools/zig/idol_case.py
#	zig/build/metal_targets.zig
#	zig/kernels/metal/sources.zig
#	zig/src/core/contract/idol_records.zig
#	zig/src/families/idol/idol.zig
#	zig/tests/contract/metal_limits.zig
#	zig/tests/contract/metal_main.zig
# Conflicts:
#	CHANGELOG.md
#	docs/kernel-contract.md
#	docs/recipes/idol.md
#	tools/zig/idol_case.py
#	zig/build/metal_targets.zig
#	zig/kernels/metal/sources.zig
#	zig/src/families/idol/idol.zig
#	zig/tests/contract/metal_limits.zig
#	zig/tests/contract/metal_main.zig
…erged case tool stays under 600 lines

The records file's header comment is one line again.
…oders, the denoiser and the decoders to a PLY

Each stage is timed and each model freed after it. The starting noise is torch's CPU-seeded randn, reimplemented;
the VAE's posterior noise and the starting noise can be given for parity with the SDK's MPS draws. The recipe's
End to end section has the stage timings, memory and PLY comparison against the SDK's own IdolPipeline run.
BiRefNet (Swin-L, ASPP-deformable decoder) in fp32 on Metal from the Idol
snapshot's birefnet.safetensors, the SDK's preprocess on the host byte for
byte (Pillow LANCZOS, erosion, crop, paste; torch's CPU bilinear input), the
tf-idol-birefnet gate tool, its case tool tools/zig/idol_birefnet_case.py
and the kernel contract for every idol_birefnet.metal kernel.
PICTURE may be a .png, .jpg, .jpeg or .webp photo, preprocessed natively
with BiRefNet; --mode synthetic|real and --crop-margin X; a photo writes
OUT_DIR/preprocessed.png and adds preprocess_s and birefnet_s to the JSON.
Recipe section and CHANGELOG line.
… on GB10)

Port BiRefNet's forward (Swin-L twice, the ASPP-deformable decoder) and its
resizes from Metal to CUDA in fp32, under the same precision contract. The
host preprocessing (preprocess.zig, photo.zig) is shared unchanged.

- idol_birefnet.cu/.cuh: fp32 kernels only, no TF32, no cuBLASLt. Every
  conv and linear is tf_br_conv, whose sums run in ordered 256-term blocks
  (a sequential fp32 FMA chain per block); cuBLASLt's fp32 summation order
  is undocumented. LayerNorm, window attention (fp32 scores, softmax and
  P V), resize, deform (torchvision deform_conv2d sampling), pool, gate and
  head are built with the glue flags.
- Bodies templated on a control, so the shipped image instantiates
  CTL_NONE only. contract_controls.cu includes
  idol_birefnet_controls.cuh for the seven planted faults.
- families/idol/cuda_birefnet*.zig and the module root
  zig/src/idol_birefnet_cuda.zig. Net.load refuses any GPU but
  sm_120/sm_121 by name with a log line.
- birefnet_config.zig holds the geometry shared with Metal.
  tf-idol-birefnet is split into a shared driver (idol_birefnet_run.zig)
  with thin Metal and CUDA mains; Metal output is byte-identical.
- Kernel contract: idol-birefnet-contract runs gates 1-4 over every kernel
  with 7 controls, joined to the family sets. nvcc 13.0 records cover
  sm_80/86/89/120/121, and the image is in contract_limits.
- Host tests for the family and the contract's references, in
  test-cuda-host.
- docs/recipes/idol-cuda.md section, including how to provide
  libjpeg-turbo 3.1.4.1 on Linux, and the changelog line.

Compiles and passes host tests; never run on a GPU; GB10 qualification
pending.
@ashhart

ashhart commented Nov 4, 2026

Copy link
Copy Markdown
Owner

Thanks @jayleaton. Keeping tf_br_conv in ordered 256-term blocks rather than cuBLASLt's unstated summation order is the right call. This waits on #1787 and #1767. It has the same line-limit point as #1767; here build/cuda.zig gains two lines, and it's at 600 on main. It needs the idol-birefnet-contract SUMMARY line and parity from a device, and qualification will show what skipping cuBLASLt costs in speed. 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