Skip to content

radio CUDA: reuse one pinned pixel stage for every image - #1140

Closed
plotarmordev wants to merge 1 commit into
ashhart:mainfrom
plotarmordev:nemotron-radio-pinned-stage
Closed

plotarmordev wants to merge 1 commit into
ashhart:mainfrom
plotarmordev:nemotron-radio-pinned-stage

Conversation

@plotarmordev

Copy link
Copy Markdown
Contributor

Branch r3/vision-pinned-stage: one commit, 25b565cc8cb55e5d8764f510054e2f8ce1e04082, on main d1018b31.
The GPU numbers were measured with #1057 (the image tower runs once) under both builds; see Receipt.

What this changes

On main, each full-size image the tower encodes allocates a 64 MiB pinned buffer and frees the previous one, and
the frees are long. e9bf05f and cc50aef (both in 1.1.0, neither in 1.0.12) moved DeviceBuffer.uploadAsync onto
copy_source's temporary stages, and Tower.features uploads each image's f32 pixels that way. For a 2048 x 1664
image that is 13,312 patches × 768 values × 4 = 40,894,464 B, which asyncChunk rounds up to a 67,108,864 B (64 MiB)
stage, the whole 64 MiB spare budget; when the stage completes, keepSpare refuses it whenever any other spare is
held, and releaseTracked frees it (host/source-history.txt has the ancestry and the source lines). Smaller
images round to smaller stages; only this size was measured. In the trace, the base build makes 20
cuMemHostAlloc calls for the 20 images (629.4 ms) and 29 cuMemFreeHost calls totalling 13,089.5 ms, the longest
872.9 ms, about one queued tower; that the free waits for the queued tower is our reading of the durations, not
something the trace shows directly.

This commit gives the tower its own copy_source.Stage, reserved next to its scratch at init for the largest image
(max_patches * patchWidth() * 4 bytes, the features room limit; larger images and video are refused before
take, so the stage never grows). Each image takes the stage, which waits on the previous copy's event, copies
its pixels in, and copyTos the same scratch offset as before; the event is recorded before the cast and the tower
kernels. Upload bytes, destinations, kernels and arithmetic are unchanged. deinit releases the stage, and the
existing quarantine covers a copy that failed to complete.

What it does not do: the engine thread still waits on the previous copy's event. From stream order tower(k-1), copy(k), event(k), tower(k) we expect image 2 to wait about one copy and images 3 and 4 about one tower, so the host
stays about two images ahead rather than queuing a whole request (not timed per image). The stage costs 39 MiB of
pinned memory for the life of the process, charged to the copy-source budget, so a load that does not fit is refused
at start. Upload paths other than the tower's keep the temporary stages.

docs/runtime/cuda-copy-sources.md gets one paragraph saying the tower owns its stage.

Test: zig/src/families/radio/upload_test.zig, "RADIO upload reuses its pinned pixels without freeing between
images", drives the real Tower.features through a fake driver and stops before device arithmetic (a fake NotReady
launch). Four full-size images alternate between two payloads (0.5 and −0.5). It requires one pinned allocation for
all four, no free between images, an event wait before each reuse, the bytes in flight equal to the image that copy
read, and the release after the last event. With the old uploadAsync line in place it fails (expected 1, found 2). A variant that writes the next image into the stage before take aborts at the in-flight byte check.

Bug report

  • Native version: main d1018b31 (1.1.1), CUDA backend, with nemotron cuda: a prompt pass with images runs the RADIO tower once, not twice #1057. The upload path dates from 1.1.0
    (host/source-history.txt); 1.1.0 itself was not run for this report.
  • Model: nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4 at revision 82d93699acd3be9e4617c53360fe5b3f5dbeb1b3.
  • GPU/OS: one DGX Spark (GB10, sm_121), driver 580.173.02, CUDA 13.0, Linux aarch64.
  • Command: tensorfold-native serve MODEL --vision --vision-image-tokens 13312 --vision-max-images 4 --context 18432 --parallel 1 --prompt-cache-gib 0.
  • Requests: bench.py image-ttft (receipts' scripts/): one generated 2048 x 1664 PNG, then four copies of it
    (3,356 and 13,366-token prompts), a warmup and 3 measured requests each (20 images), greedy, thinking off,
    "draft": false, 32 tokens, cooled to 75 °C before every request.

Both servers under nsys profile --trace=cuda --sample=none --cpuctxsw=none, base = main + #1057, this branch =
base + this commit (scripts/host_api.py, scripts/timeline.py on the SQLite exports):

base this branch
cuMemHostAlloc from the first image on 20 calls (one per image), 629.4 ms 0
cuMemFreeHost from the first image on 29 calls, 13,089.5 ms, longest 872.9 ms 10 calls, 24.2 ms, longest 10.2 ms
cuMemHostAlloc while loading 77 78 (the tower's stage)
GPU idle inside each 4-image tower 0.145, 0.156, 0.150, 0.137 s 0.005, 0.005, 0.006, 0.005 s
GPU idle inside each measured 1-image request 0.024, 0.023, 0.024 s 0.007, 0.007, 0.007 s
kernels; host-to-device copies, bytes 34,387; 128,958, 76,011,157,900 the same
1 image token_sha256, 4 requests 541bca2d…9735ec 541bca2d…9735ec
4 images token_sha256, 4 requests ef22b44e…d444c4 ef22b44e…d444c4
1 image TTFT, median of 3 2.301 s 2.239 s (−2.7 %)
4 images TTFT, median of 3 15.058 s 14.030 s (−6.8 %)

Full hashes: 541bca2d52751d8bb9c44b2af53ce8ff2dfb93745d8cb6bdb740442dad9735ec,
ef22b44ea1d731737c2d2a757d671d03b1269563d6bb505944e58608ecd444c4 (the same as #1057's receipts).

How to reproduce

nsys profile --trace=cuda --sample=none --cpuctxsw=none --output=tower \
  tensorfold-native serve MODEL --port 8897 --vision --vision-image-tokens 13312 --vision-max-images 4 \
  --context 18432 --parallel 1 --prompt-cache-gib 0
API_URL=http://127.0.0.1:8897 python3 bench.py image-ttft   # receipts' scripts/, next to prompts.py
# stop the server, then:
nsys export --type sqlite --output tower.sqlite tower.nsys-rep
python3 host_api.py tower.sqlite    # pinned alloc/free before and after the first ViT kernel
python3 timeline.py tower.sqlite    # per-request tower/prompt busy time and idle

Receipt

Raw files: https://github.com/plotarmordev/TensorFold/tree/nemotron-receipts/receipts/nemotron/vision-pinned-stage (README.md maps every file; SHA256SUMS).

  • Environment, with the TensorFold commit, MLX or PyTorch version, machine and GPU, checkpoint and revision: this
    branch at 25b565cc on main d1018b31. Without nemotron cuda: a prompt pass with images runs the RADIO tower once, not twice #1057 every image runs through the tower twice, so both measured
    builds carry it: base = d1018b31 + nemotron cuda: a prompt pass with images runs the RADIO tower once, not twice #1057 (732b0e31), this branch = base + this commit (5a4d65a1). A third
    build adds the RADIO fused-glue pull request on top (f6d8ba5a, "both" below). env/commits.txt shows each
    measured commit's parent, tree and git patch-id --stable next to the submitted ones: 5a4d65a1 and f6d8ba5a
    carry 25b565cc's patch-id, every nemotron cuda: a prompt pass with images runs the RADIO tower once, not twice #1057 copy carries e2357b58's. Each built from a detached checkout of that
    commit (env/PROVENANCE-*.txt): zig build fatbins and zig build install with -Dtarget=aarch64-linux-gnu.2.28 -Dcpu=baseline -Doptimize=fast -Dstrip=true -Dsm=121 -Dversion=1.1.1 -Dnvcc=… (zig 0.17.0, nvcc 13.0.88). tensorfold-native sha256: base 5e70f78f…dcef69, this
    branch 7df4fbe7…d56ca8, both 34fba744…a29cda (full in env/ and each hold's .sha256). One DGX Spark (GB10),
    driver 580.173.02, CUDA 13.0, nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4@82d93699acd3be9e4617c53360fe5b3f5dbeb1b3,
    served as in the bug report on loopback with TF_HEAT_HIGH=85 TF_HEAT_LOW=75, a fresh server and HOME each hold,
    each hold started at a zone of 60 °C or less.

  • Exactness, from tools/bench_concurrent.py --alone --serial or token_sha against "draft": false: every
    image-ttft request (2048 x 1664 images, one and four) is "draft": false, and every request on every build (base,
    this branch, both; profiled and unprofiled) gives 541bca2d… for one image and ef22b44e… for four. bench.py decode on "both" runs 11 prompts drafted (count, prose, code, image, multilingual and six languages) plus count with
    "draft": false; every drafted reply equals "draft": false, and the hashes equal main's. Its image prompt
    carries one 320 x 208 picture, so that is the one drafted image request that ran, through this stage. The
    full-size image-ttft workload was not run with drafts; bench_concurrent.py, resumed and concurrent image
    requests were not run.

  • Decode speed before and after, from tools/bench_openai.py: not measured. Decode rounds and text prompts do not
    reach the tower's upload.

  • Prompt speed before and after: tools/bench_cache.py sends no images, so TTFT is the server's
    time_to_first_token from bench.py image-ttft. This commit alone was timed only under nsys, one hold each (base,
    then this branch; table above). Unprofiled, it was timed only together with the fused-glue pull request, in the
    order base, both, base, both:

    base both base both both vs base
    1 image, median of 3 2.293 2.177 2.288 2.172 −5.1 % (−0.116 s)
    4 images, median of 3 15.001 13.870 15.008 13.856 −7.6 % (−1.14 s)

    The ranges do not overlap (one image 2.286–2.295 s against 2.168–2.183 s). How the unprofiled gain splits between
    the two pull requests was not measured. Every four-image request on every build, warmups included, paused once at
    the heat gate inside its prompt pass (heat_wait_s 8.0 in each server log); no one-image request paused. The TTFTs
    above include that pause. The four-image saving is larger than the tower idle this removes (about 0.15 s a
    request); the rest is consistent with the engine thread no longer sitting in cuMemFreeHost while work is queued,
    so more work is queued when the heat gate pauses, but the trace does not split the two.

  • Tests run, with pass, fail and skip counts (CPU only, a DGX Spark, -j1, nice -n 19, fresh source copies), at this tip:
    zig build test-bolt-load-host -Dbolt-load-filter='RADIO upload' test-cuda-copy-source 10/10 steps, 49 of 49;
    zig build test 211/211 steps, 2960 of 2974 passed, 14 skipped, 0 failed; zig build test-kernel-contract test-golden fuzz-cortheon 19/19 steps, 76 of 76 tests (kernel contract 46, 9, 3 and 7; fuzz-cortheon 11 of 11),
    test-golden 26/26 and 306/319 equal to the frozen answers, 0 unexpected differences;
    zig fmt --check on the changed files clean; tools/zig/lean_check.py 444 problems, 444 on main, 0 new. Failing
    receipts: the new test with the old upload line (expected 1, found 2), and the early-overwrite variant (abort at
    the in-flight byte check).

  • Not run, and why:

    • Metal: the change is CUDA-only, and there is no Mac here.
    • Tensor parallelism: ranks past 0 build no vision; a TP=2 server with images was not run.
    • Video and Bolt with images: not run on a GPU (video goes through the same features and stage).
    • tf-cuda-test kernel-contract: no kernel changed.
    • Per-image cuEventSynchronize durations were not broken out.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Measured: full-size image replies byte-equal before and after with "draft": false, and drafted equal to "draft": false on 11 decode prompts, one with a 320 x 208 picture, with this commit in the build. Resumed and concurrent requests with images were not run.)
  • No precision traded for speed. If the bits change, the description says which sums change and why. (No sum changes: the same bytes are uploaded to the same place; hashes, kernel counts and copy bytes are equal.)
  • 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. (No kernel changed.)
  • Prompt processing is no slower than the last release. (Measured against main + nemotron cuda: a prompt pass with images runs the RADIO tower once, not twice #1057, not a release: image prompts are faster under nsys and, together with the fused-glue change, unprofiled. Text prompts do not upload pixels.)
  • 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. (A host test with a fake driver; it needs no GPU.)
  • 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.

The 1.1.0 temporary-upload path rounds full-size image pixels to a 64 MiB pinned stage. When smaller spares occupy the pool, the completed stage cannot be parked, so the next image frees it with cuMemFreeHost and allocates another while the GPU waits.

Reserve one stage with the tower scratch and wait on each copy event before refilling it. This removes per-image pinned allocation and disposal without changing upload bytes or device arithmetic. Exercise reuse and in-flight byte preservation with alternating image payloads.
@ashhart

ashhart commented Oct 25, 2026

Copy link
Copy Markdown
Owner

Landed on main in 93a215b..35e46c0, with your records under your authorship plus the maintainer fix-ups. Thank you. Measured with the whole group (#1139, #1140, #1143, #1147, #1164, #1169, #1171 and #1170's Flash Next slice) against main on GB10s:

  • Qwen3.8 Flash Next FP4-Local: tokens identical (0 of 20 rows differ). Prefill +23.7% to +28.9% from 2k to 64k (16k: 1356.6 → 1748.4 tok/s). Decode +1.0% to +2.2% at 1, 2, 4 and 8 streams.
  • Flash Next MLX 4-bit, alternating base/new/new/base: tokens identical, prefill +0.4% to +0.8%, decode within noise.

- Synthema, Ash's AI Assistant.

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

Labels

landed Landed on main by the maintainers; the commits keep the contributor's authorship

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants