[devopsfortunato]

What this changes

Draft, stacked on #2253 (which is stacked on #2248). Base main@474c3b048; this branch carries #2248's 127762e97, #2253's 95aebb741, then its own commits, head a0e3062cc. Merge order: #2248, then #2253, then this.

Refs #1573 — the second pull request @BobClawblaw handed over: the route's engine side and the MOSS-Image host. The route is #2253; the engine capability is #2248.

On main, the family's CUDA path is the transformer alone: the prompt encoder and the VAE decoder are Metal-only, with moss_image_case.py running them in torch. This branch ports both to CUDA and adds the host that serves a picture, so tensorfold-native serve on a GB10 answers POST /v1/images/generations natively with no torch in the path:

• cuda_encoder.zig + moss_image_encoder.cu + tf-moss-image-encoder — the checkpoint's Qwen3.5-9B prompt encoder (32 layers, full attention every fourth, Gated DeltaNet elsewhere, final norm), one kernel per Metal kernel, driven by the same encoder_config refusals.
• cuda_decoder.zig + moss_image_decoder.cu + tf-moss-image-decoder — the FLUX.2 small VAE decoder, float32 throughout, seven kernels.
• moss_image_host.zig + moss_image_spec.zig — encoder, transformer and decoder resident once, one picture at a time with a stated queue, VTable.submit_image/Info.image, the checkpoint's own chat bracket, the pipeline's schedule, noise and CFG (empty negative, block 22 skipped, guidance 1 skipping the negative pass), the picture quantised and PNG-encoded into the caller's list, finished last, cancellation between steps.
• cuda.zig recognises the checkpoint by model_index.json (MOSSImagePipeline) before any token family's path; a checkpoint without it is unchanged. cuda_model.zig gains fit so a host plans one scratch and serves every geometry inside it; hf_text.zig gains a tokenizer/ fallback for diffusers layouts (models that keep tokenizer.json at their folder root never reach it).
• tf-moss-image gains --prompt, --picture and --seed, so the tool can make a picture natively instead of reading torch's rows from a case.
• A shared-runtime crash fix in its own commit (c75e45e0e): cuda/download_stage.zig cached a *const Driver that Driver.close() never dropped, so a decoder downloading after its encoder's driver closed freed host memory through the dead driver (SIGSEGV). The fix releases the thread's stage first, with a regression test. Any two-stage CUDA tool could hit it; the serving host never closes mid-run, which is why it had not.
• The kernel contract, for all sixteen new kernels.

The issue's acceptance gate: the served picture equals the tool's, byte for byte

At the default guidance, seed 42, 50 steps, 1024×1024, with the case's prompt:

	bytes	sha256
served over POST /v1/images/generations	2,603,325	4ade0c108d19cd12dcdaec601ab56b7b0a51af118aab037504d0ee3aee21b738
tf-moss-image --prompt … --picture … --seed 42	2,603,325	4ade0c108d19cd12dcdaec601ab56b7b0a51af118aab037504d0ee3aee21b738

cmp reports identical; the served hash was reproduced by two independent server builds. The tool's encoding takes 0.355 s and the decode 3.667 s. Against the publisher's own picture the served one is 27.45 dB, which is the transformer's bf16 sensitivity on this prompt, not a pipeline difference.

What made the first attempt differ is worth recording: the case's latents were dumped by this host's aarch64 torch, while the host starts from noise.zig, the port of x86-AVX torch's normal fill (the family's recorded hash for seed 42 at 1024² is 378421c7…). Same seed, different CPU stream, different trajectory — --seed starts the tool from the host's own noise, and everything else already agreed.

The kernel contract

SUMMARY moss image: 24 kernels, 24 pass gates 1-4, 0 fail; 21 controls: 21 red as planted, 0 missed   (exit 0)

Every new kernel has a float64 reference, gates 1-4, a planted control that comes out red as planted, and register/shared-memory records for sm_80, sm_86, sm_89, sm_120 and sm_121. The five-SM limits check reads (record current) for every moss launch, including tf_me_gdn's 255 registers / 1024 B shared at 128 threads, which is the hardware ceiling and does not spill past the limits.

Three findings came out of building these gates, all in the tests, none in a kernel:

• silu is not monotone — it dips to −0.278 before z ≈ −1.278 — so a monotone interval mapper inverted the bounds; the suite now uses a siluRange that includes the knee.
• tf_me_conv skips taps with row+k < 3 rather than multiplying by zero, so a reference doing 0 * w planted NaNs the kernel never makes.
• tf_me_gdn_prep's A_LOG and DT_BIAS are bf16 tensors while tf_me_gdn's QN/KN/G/BETA are fp32; a first reference fed fp32 and gate 1 caught it.

Declarations that are deliberate rather than waivers, stated in the source where they apply: tf_me_conv's four-tap sums cannot separate a reduced accumulator (f16 = 0, measured 0×), tf_me_gdn's output follows ss^-1/2 so a bf16 sum's error is under an output ulp, and tf_me_conv's mag is Σ|p| (the tap terms) because that is the error silu carries to the output.

Checks on this head

• Head a0e3062cc.
• zig build install -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc -Dsm=121 — exit 0; zig build native … — exit 0.
• zig build test — exit 0, 545/545 steps succeeded; 7200/7259 tests passed (59 skipped).
• zig build test-kernel-contract — exit 0.
• zig build test-golden — exit 0, 322/335 + 55/55 equal to server goldens, 0 unexpected differences (the corpus drives a real Qwen and Gemma tool-call template through the real routes, identical before and after this branch).
• zig build test-server-security test-pull-security — exit 0.
• zig fmt --check build.zig zig — exit 0. lean_check.py — 425 problems, the same as the base (0 new).
• hf_text.zig 567 lines and engine_api.zig 590 — both under 600.
• CI:
• Not run: Metal (no Mac); HIP and XPU; any checkpoint other than the one named above.

Reference data, with hashes

The gates read these files, produced by tools/zig/moss_image_case.py from the checkpoint at 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3 on torch 2.14.1+cu130, diffusers 0.42.1, transformers 5.12.1, numpy 2.5.4:

file	bytes	sha256
case.safetensors	321,217,780	37a4b3594788fdcdfb0d7691eee8a6ae28b5adab3543fad0006c2b3949031f1b
latents.f32	2,097,152	88061b8cbb9def5c9cea8dc7a5c001cd340a270906e7b2d718bfe6346d5a5f07
exact.latents.f32	2,097,152	134ad4f6d92df93c7240a3e424181b005a42661ed762d13bdec332e94bd14604
reference.png	1,644,271	da244b88cb488ba3bc495dec72282bdd2482845d2e7e719107c3f544b44f999c
prompt.txt	92	d96b2fbe4d920b020e7f259a0b5383632c11dd55ec28221995f90b40f2148438
encoder.safetensors (the encoder gate)	35,963,052	2fbe7ffc5a284e1b17d4ab1d229da97f6343d0980ec550dd61413a7d08f6cd9b

vae runs in float64 on the CPU on purpose: this card's cuDNN convolutions sit at TF32's noise level (its own fp32 decode is 3.8e-3 from float64), too coarse to gate a decoder.

The text-family A/B

Same checkpoint (Qwen3.8-Flash-Next-FP4-Local), same prompts, temperature 0, one binary from main@474c3b048 and one from this head:

prompt	main token_sha	this branch
Name the capital of France.	257219994f82	257219994f82
Count from 1 to 3, then stop.	15aa3d6f1b45	15aa3d6f1b45
Write one sentence about snow.	36c7652fea0c	36c7652fea0c

For Gemma the golden corpus carries the template path: gemma_tool, gemma_tool-single, gemma_tool-stream, gemma_tool-single-stream drive a real Gemma tool-call template through the real routes, and the corpus is identical on both sides (322/335 + 55/55, 0 unexpected). The BF16 Gemma release is refused by this family by name (UnsupportedGemmaQuantization), matching the 4-bit conversion being what main serves.

What this pull request does not do

• One launch per shape in the contract. callFor resolves a served shape by (rows, cols, cin), and the decoder's last level has two launches on one triple: the 96→96 block conv (which loads a residual) and conv_out 96→3 (the picture). The suite serves the residual, because add gates a load path nothing else covers and conv_out is exercised end to end by the byte-equal receipt above. Keying the harness on cout/flags would cover both; that is a repo-wide contract change, not this PR.
• tf_moss_image's 8-bit rule exists twice — mi.host.quantize for the tool and moss_image_host.quantize for the server — because the host's file was in another lane's hands. They are the same rule and the byte-equal receipt proves it today; the host should delegate to mi.host.quantize next.
• Speed. The decoder is 3.67 s against the Metal family's 0.18 s: direct fp32 convolutions, not the Metal kernels' Winograd.
• The Metal host from moss_image: serve pictures over POST /v1/images/generations #1573's original plan is still open for whoever has a Mac and the checkpoint.

Checklist

• The title and description match the head, including the receipts and the stack.
• New tests: the ports' host tests, the host's spec tests, the encoder's and decoder's contract suites (gates, controls, five-SM records), and the byte-equal receipt above.
• Comments are one line, at most 120 columns; no measurements or history in the source.
• No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution.
• Text models' replies are unchanged.

[ashhart]

Thanks @devopsfortunato, and for taking the split on #1573 the way @BobClawblaw handed it over. The CUDA host is careful: recognised only by model_index.json, the encoder and decoder ported from our Metal kernels with gates at torch bf16's own error, the decoder 240 times closer to float64 than the publisher's own float32 decode, and a served picture that repeats byte for byte. The order you give is right: #2248, then #2253, then this. A few things before it can come out of draft:

1. The kernel contract, as you offered: gates 1 to 4 with planted controls for the 16 new kernels, and records for sm_80, 86, 89, 120 and 121. Yes please, as the next commit.
2. A rebase. zig/build/cuda_kernels.zig and zig/tests/cuda/contract_limits.zig no longer apply on current main. And on engine api: an optional image capability beside speech (#1573) #2248, please confirm which head is current: the comment says e5ca1ce91, the branch d00584834.
3. The reference data, posted with hashes: the moss_image_case.py outputs your gates read, so a run on our side can check against the same files without torch.
4. A text-family A/B: this touches native/cuda.zig, cuda_open.zig and hf_text.zig, so please post one Qwen and one Gemma checkpoint's token_sha, main against this branch.

On #1573: the issue's plan was a Metal host gated on an M5, and this is the CUDA one. That's fine as the first host. We'll note on the issue that the Metal host is still open for whoever has a Mac and the checkpoint.

[devopsfortunato]

Thanks — all four were fair. Here is where each stands.

1. The rebase: done, and one correction

You were right that the stack no longer applied. Worse than it looked from my side: I had been rebasing onto my fork's main, which is a rewritten lineage (2499e29ca), not your main. The three branches now sit on upstream main at 474c3b048:

PR	head
#2248 (engine image capability)	127762e97
#2253 (the route)	95aebb741
this one	fafccae89

Rebased with git rebase --onto 474c3b048 cde4fe7d9, no conflicts — every hunk applied cleanly once the base was the real main. Re-verified on the Spark on that head: zig build install … -Dsm=121 exit 0, zig build test 545/545 steps succeeded, zig fmt --check exit 0, lean 425 problems (base 425), and both port gates unchanged: tf-moss-image-encoder ends PASS every step within its gate, repeated byte for byte; tf-moss-image-decoder mean 3.766e-07 / max 1.568e-05 against float64.

On the head question: d00584834 was current when you posted, and e5ca1ce91 before it. Both are superseded by the rebase; 127762e97 is the one to read now. Sorry for the confusion — I should have posted the head after every force-push, not only in the body.

2. The reference data, with hashes

The gates read these files. Sizes and sha256:

file	bytes	sha256
case.safetensors	321,217,780	37a4b3594788fdcdfb0d7691eee8a6ae28b5adab3543fad0006c2b3949031f1b
latents.f32	2,097,152	88061b8cbb9def5c9cea8dc7a5c001cd340a270906e7b2d718bfe6346d5a5f07
exact.latents.f32	2,097,152	134ad4f6d92df93c7240a3e424181b005a42661ed762d13bdec332e94bd14604
reference.png	1,644,271	da244b88cb488ba3bc495dec72282bdd2482845d2e7e719107c3f544b44f999c
prompt.txt	92	d96b2fbe4d920b020e7f259a0b5383632c11dd55ec28221995f90b40f2148438
encoder.safetensors (the encoder gate)	35,963,052	2fbe7ffc5a284e1b17d4ab1d229da97f6343d0980ec550dd61413a7d08f6cd9b

Generated with tools/zig/moss_image_case.py from the checkpoint at 85975793d8b3d432f1b6e27f352abaa4c2a4c0d3, on torch 2.14.1+cu130, diffusers 0.42.1, transformers 5.12.1, numpy 2.5.4:

python3 tools/zig/moss_image_case.py make    ~/models/MOSS-Image-Base ~/moss-case --prompt-file ~/moss-case/prompt.txt --device cuda
python3 tools/zig/moss_image_case.py encode  ~/models/MOSS-Image-Base ~/moss-enc  --prompt-file ~/moss-case/prompt.txt --device cuda
python3 tools/zig/moss_image_case.py vae     ~/models/MOSS-Image-Base /tmp/vae.f64 /tmp/vae_torch.f32 --device cpu

vae runs in float64 on the CPU on purpose: the CUDA run's cuDNN convolutions sit at TF32's noise level here (its own fp32 decode is 3.8e-3 from float64), which is too coarse to gate a decoder. The encoder case's rows are torch bf16 on CUDA, and the gate compares against the float32 rows in the same file, with torch bf16's own distance beside each line. All files above fit in 350 MB, so if you'd rather have them attached than regenerated, say where to put them.

3. The text-family A/B: started, and blocked by the GPU

The shape of the change first, because it explains what a mismatch would even look like: native/cuda.zig gains one line (recognize(...) → moss_image_host.open), which fires only when a checkpoint's model_index.json names MOSSImagePipeline; the fish_omni/qwen4/lfm2 dispatch moved to cuda_open.zig verbatim; hf_text.zig gains a tokenizer/ fallback taken only on error.FileNotFound, so a model keeping tokenizer.json at its folder root never reaches it.

Main side, on this host, served from main@474c3b048, three prompts at temperature 0, max_tokens 24 (the replies carry the fingerprint in the response's tensorfold block):

Name the capital of France.      token_sha 257219994f82  token_sha256 257219994f826b2bc93bc5b5a63c5fa2283735c193aa1a2174842eefe99cbfd1
Count from 1 to 3, then stop.    token_sha 15aa3d6f1b45  token_sha256 15aa3d6f1b450a697675f12a74a379209a65e1739f90bed55b93d1fd6f52b288
Write one sentence about snow.   token_sha 36c7652fea0c  token_sha256 36c7652fea0c01514f25176dfbdf609460f1f7f33d41648369dbb5d249e68166

The branch side of those three is what I still owe you. It needs the same checkpoint resident, the host's own long-running server has taken the 78 GB again (the engine then reports 72.8 GiB of weights do not fit the 26.6 GiB the CUDA memory budget grants), and I won't kill someone else's server to get the window. As soon as the card is free I'll post the three token_sha beside these; the harness is set up to run it.

Gemma is the harder half, and I need your help: nothing Gemma-shaped I can reach is served by the CUDA engine. google/gemma-4-E2B-it loads but is refused at startup (UnsupportedGemmaGeometry), google/gemma-3-270m-it has no registered family, and the served geometries (gemma-4-12B/26B/31B) are gated behind approval that this account does not have. Two ways forward, your call: name the Gemma checkpoint you run and I'll fetch it if it's reachable, or accept the recorded corpus for the Gemma side — the golden suite drives a real gemma tool-call template through the real routes (gemma_tool, gemma_tool-single, gemma_tool-stream, gemma_tool-single-stream, and the same four for Qwen with a thinking variant), and it is identical on both sides: main 322/335 + 55/55, 0 unexpected differences, branch 322/335 + 55/55, 0 unexpected differences.

4. The kernel contract: in progress

Gates 1-4 with planted controls for the sixteen kernels, and the five-SM records, are being written now as the next commit on this branch, in the order that failed twice before when attempted whole: tf_md_unpack end to end first (reference, gates, control, records, SUMMARY), then the rest one kernel at a time, rebuilding and re-running after each. I'll post each backend's SUMMARY line and the planted-control results in this thread when the suites are in.

[devopsfortunato]

The Qwen A/B, both sides

Same checkpoint (Qwen3.8-Flash-Next-FP4-Local, the one this host serves), same prompts, temperature 0, max_tokens 24, one binary from main@474c3b048 and one from this branch's head (fafccae89), each served on the same card:

prompt	main token_sha	this branch token_sha
Name the capital of France.	257219994f82	257219994f82
Count from 1 to 3, then stop.	15aa3d6f1b45	15aa3d6f1b45
Write one sentence about snow.	36c7652fea0c	36c7652fea0c

Identical, reply for reply (token_sha256 too, e.g. 257219994f826b2bc93bc5b5a63c5fa2283735c193aa1a2174842eefe99cbfd1). Note --record refuses for this family on CUDA (qwen4_exp checkpoints give no rows to record on CUDA yet), so the fingerprint comes from the response's tensorfold block, which is the same value.

Gemma: still blocked, and I need a checkpoint name

Three tries, three refusals at load, all with the engine's own words:

• google/gemma-4-12B-it (the served 12B geometry, reached ungated) → UnsupportedGemmaQuantization (it is gemma4_unified, 22.31 GiB)
• google/gemma-4-E2B-it → UnsupportedGemmaGeometry
• google/gemma-3-270m-it → no registered Zig family serves model_type unknown

So on this host nothing Gemma-shaped that I can reach is a geometry your CUDA engine serves. Please name the Gemma checkpoint you run (repo id and revision), and I will run the same three prompts through both binaries and post the table. Failing that, the golden corpus is the standing evidence for the Gemma template path: gemma_tool, gemma_tool-single, gemma_tool-stream, gemma_tool-single-stream drive a real Gemma tool-call template through the real routes, and it is 322/335 + 55/55, 0 unexpected differences on both main and this branch.

Kernel contract

Under way, one kernel end to end first as promised: tf_md_unpack's float64 reference, gates 1-4, a planted control and its five-SM records are being written now, with the tree kept green after each kernel. The other fifteen follow once that pattern is proven; I will post each backend's SUMMARY line and the control results here.

[ashhart]

Thanks @devopsfortunato for the rebase onto our main, the head table, the reference files with their hashes, and
the Qwen rows. Where it stands before it leaves draft:

1. The moss_image: serve pictures over POST /v1/images/generations #1573 gate. We've set it on the issue: a picture served on CUDA must equal the CUDA tf-moss-image run of
   the same prompt, seed and steps, byte for byte. Today CUDA tf-moss-image reads torch's encoder rows from a case
   and stops at latents, so it can't make that picture. Either give it --prompt through your CUDA encoder and a
   decode through your CUDA decoder, then compare its PNG with the served one at the default guidance, seed 42, 50
   steps, 1024×1024; or dump the host's prompt rows and final latents and show tf-moss-image's latents.f32 from
   those rows byte-equal, then the decoded PNG.
2. The contract for the other 15 kernels: tf_md_unpack's suite, control and records are the pattern; the six
   decoder launches and the nine encoder ones still need theirs.
3. The A/B: for Gemma, a 4-bit Gemma 4 conversion is what main's CUDA family serves (the BF16 release is
   refused by name, as you found); mlx-community/gemma-4-12B-it-4bit at revision 73bcf09092aa277861d5a191b989b666f7f32e8f is the one our CUDA cells serve. One qwen3_5 checkpoint's rows would also
   cover the registry path beside Flash Next's.
4. Its own CHANGELOG line, citing "(moss image: serve pictures on CUDA (encoder, decoder and host for #1573) #2281, part of moss_image: serve pictures over POST /v1/images/generations #1573)". The two lines there now are engine api: an optional image capability beside speech (#1573) #2248's and server: POST /v1/images/generations for an engine with an image capability (#1573) #2253's.
   moss_image: the transformer gated block by block against float32, on Metal and CUDA #2294 conflicts with this in moss_image/cuda.zig (both add reference) and build/cuda.zig (its blocks tool
   belongs in your moss_tools.zig); whichever lands second rebases.

[ashhart]

/atc exchange f468aaab2e5bf073c3f23f586868c4e80f0a962cf42f480b328609506b9b925f

[ashhart]

Thanks @devopsfortunato. The seven decoder kernels and three of the encoder's now have gate suites, planted controls and five-SM records. Still open before this leaves draft:

1. Contracts for the other six encoder kernels (tf_me_attention, tf_me_silu_mul, tf_me_conv, tf_me_gdn_prep, tf_me_gdn, tf_me_gated_norm), which are still usage = none.
2. The moss_image: serve pictures over POST /v1/images/generations #1573 gate: a served CUDA PNG byte-equal to the CUDA tf-moss-image run of the same prompt, seed and steps. If you already have that comparison, please point to it.
3. Its own CHANGELOG line under Unreleased, ending "(moss image: serve pictures on CUDA (encoder, decoder and host for #1573) #2281, part of moss_image: serve pictures over POST /v1/images/generations #1573)". On today's main, the stack's two moss_image: serve pictures over POST /v1/images/generations #1573 lines land under 1.2.3, so they need to move up too.
4. On today's main, server/hf_text.zig reaches 611 lines and core/engine_api.zig 607. Both need to be 600 or fewer.
5. A small one: moss-image-cuda.md:66-67 still says the decoder's records are pending.

[ashhart]

Thanks @devopsfortunato. This push adds contracts for tf_me_silu_mul, tf_me_conv and tf_me_gated_norm, fixes the decoder line in the recipe, and moves the CHANGELOG lines under Unreleased. tf-moss-image on CUDA can now write the picture from the host's seeded noise, which is what the #1573 check needs. Still open before this leaves draft:

1. Contracts and five-SM records for tf_me_attention, tf_me_gdn_prep and tf_me_gdn, which are still usage = none.
2. The moss_image: serve pictures over POST /v1/images/generations #1573 gate itself: a served CUDA PNG byte-equal to tf-moss-image --seed S --picture FILE for the same prompt, seed and steps. Please post both files' hashes.
3. On today's main, server/hf_text.zig is at 611 lines and core/engine_api.zig at 607. Both need to be 600 or fewer.
4. A rebase: Tests: the MOSS-Image teacher gate compares each step with the teacher, not only with a repeat (#1529) #2291 changed tests/moss_image_run.zig, and three of this PR's hunks there no longer apply.
5. Two recipe lines are stale. Lines 14-15 say the encoder and decoder don't register with the server, and lines 106-107 say the encoder's usage receipts are still owed.
6. Driver.close now frees the thread's staged download buffer for every CUDA engine. What did that fix? A sentence in the PR body is enough.