Skip to content

Fix pre-warm accumulation; correct resident_generate()'s @return - #64

Merged
TroyHernandez merged 6 commits into
mainfrom
fix/prewarm-accumulation
Aug 24, 2026
Merged

Fix pre-warm accumulation; correct resident_generate()'s @return#64
TroyHernandez merged 6 commits into
mainfrom
fix/prewarm-accumulation

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

Two unrelated defects in code merged today, in one PR at Troy's request.
They are separate commits and separate NEWS entries — the behavioural fix
is 653eedb, the docs correction is fddd739, so bisecting the allocator
change does not mean reading past roxygen.

Found by vientito (gpu.ctl) reviewing #62 after it merged.


1. Pre-warm accumulation (653eedb)

.resident_prewarm(), added in 0.2.2.4, asked for the full onload need
on every activation
. After the first render that doubles the CUDA
caching allocator's pool, and under release = FALSE it keeps doubling
until activation is refused.

Three activate/render/deactivate cycles on SDXL, one process:

before after
release = FALSE 5.299 → 10.322 → REFUSED 5.396 → 5.498 → 5.498
release = TRUE 5.299 → 10.123 → 10.125 5.396 → 5.396 → 5.396

The cycle 1→2 step was 5.023 GiB against a pre-warm block of
need * 1.05 = 5.021 GiB. That is the mechanism: a render fragments the
cache into its activation blocks, so the next single large request cannot
be served from what remains and takes a fresh cudaMalloc beside the old
one. release = TRUE masked it, because .resident_release_vram()'s
empty_cache caps the pool at two blocks.

Why this matters to the caller it was written for. A residency broker
passes release = FALSE deliberately — an exclusive device grant means
blocks never returned cannot be taken by anything else between requests.
So an SDXL entry would serve two requests and refuse every one after. On a
12 GB card sooner.

It broke the budget independently of the refusal. Both arms double
between cycle 1 and cycle 2, so any peak measured on a single activation
understates steady state about 2x — declaring 5.398 for something that
needs 10.125 on its second request, with the first looking perfect.

The fix reads what the allocator already holds and grows only what is
missing: skip when the pool covers the transfer, otherwise request the
shortfall. On a cold pool held is 0, so it is byte-identical to the
previous behaviour and the 74x cold-start win is untouched.

How it was found matters

Activate/deactivate cycles with no render in between hold dead flat at
5.098 GiB in both arms — the block is reused cleanly across bare
transitions. It takes a generation in between to fragment the pool. So one
cycle is not a test, and neither is a cycle without a render. The obvious
probe would have shown nothing wrong.

held is injectable for the same reason .resident_check_fits() takes
free_gb: without CUDA the measurement is 0, which would always warm, so
a test that wants the skip has to state what the pool holds rather than
depend on the machine having a card. Six new assertions cover skip,
shortfall, cold pool, and two nonsense readings.


2. resident_generate()'s @return was wrong (fddd739)

The @return added in 0.2.2.4 said flux1/flux2/zimage return bare
image arrays, ltx returns a video array, and sdxl is the odd one out
with list(image, metadata). All four claims are wrong, and in the
damaging direction: it tells a consumer to special-case sdxl and treat
flux2 as an array, which fails on first contact.

Every generator returns a list:

txt2img_flux    invisible(list(image = img_array, metadata = metadata))
txt2img_flux2   invisible(list(image = img_array, metadata = metadata))
txt2img_zimage  invisible(list(image = img_array, metadata = metadata))
txt2img_sdxl    return(list(image = img_array, metadata = metadata))
txt2img_sd21    return(list(image = img_array, metadata = metadata))
txt2vid_ltx2    invisible(result)   # video, audio, sample_rate, latents, ...

So sdxl agrees with the other image families rather than diverging,
and the corrected contract is simpler than the one it replaces: unwrap
$image across all five image families, $video for ltx. The only real
inconsistency left is return() versus invisible(), which affects
console auto-printing and nothing else.

The decision not to normalise the shapes stands; only the stated reason
was inverted.

Observed at runtime for sdxl, sd21 and ltx during this work,
confirmed for flux2 by the gpu.ctl side, and read from source for
flux1 and zimage.


Suite 1203 assertions, 0 failures. vientito is independently re-running
both arms against this branch before merge — they hold the before-numbers
and the harness, and the author measuring his own fix is the arrangement
their own measurement doc argues against.

.resident_prewarm() asked for the full onload need on EVERY activation.
After the first render that doubles the caching allocator's pool, and
under release = FALSE it keeps doubling until activation is refused.

Measured on SDXL, three activate/render/deactivate cycles in one process:

  release = FALSE   5.299 -> 10.322 GiB -> REFUSED
  release = TRUE    5.299 -> 10.123 -> 10.125 (plateaus)

The step from cycle 1 to cycle 2 was 5.023 GiB against a pre-warm block of
need * 1.05 = 5.021 GiB, which is the whole mechanism: a render fragments
the cache into its activation blocks, so the next single large request
cannot be served from what is left and takes a fresh cudaMalloc beside the
old one. release = TRUE masked it, because .resident_release_vram()'s
empty_cache caps the pool at two blocks; release = FALSE never empties, so
the third activation found the card full of retained pre-warm blocks.

This matters most to the caller it was written for. A residency broker
passes release = FALSE deliberately -- an exclusive device grant means
blocks never returned cannot be taken by anything else between requests --
so an SDXL entry would serve two requests and refuse every one after.

The doubling also broke the budget independently of the crash: both arms
double from cycle 1 to cycle 2, so a peak measured on a single cycle
understates steady state about 2x. Declaring 5.398 for something that
needs 10.125 on its second request is admit-then-OOM with the first
request looking perfect.

The fix reads what the allocator already holds and only grows what is
missing: skip entirely when the pool covers the transfer, otherwise
request the shortfall. On a cold pool `held` is 0, so it is byte-identical
to the previous behaviour and the 74x cold-start win is untouched.

After, same probe:

  release = FALSE   5.396 -> 5.498 -> 5.498
  release = TRUE    5.396 -> 5.396 -> 5.396

Flat from cycle 2 in both arms rather than merely surviving cycle 3. The
residual 0.102 GiB in the FALSE arm is the render's own fragmentation and
does not compound.

Worth recording how it was found, because an obvious test would have
missed it: activate/deactivate cycles with NO render in between hold dead
flat at 5.098 GiB in both arms. The pre-warm block is reused cleanly
across bare transitions. It takes a generation in between to fragment the
pool, so one cycle is not a test, and neither is a cycle without a render.

`held` is injectable for the same reason .resident_check_fits() takes
free_gb: without CUDA the measurement is 0, which would always warm, so a
test that wants the skip has to state what the pool holds rather than
depend on the machine having a card.
The @return added in 0.2.2.4 said flux1, flux2 and zimage return bare
image arrays, ltx returns a video array, and sdxl is the odd one out with
list(image, metadata). All four claims are wrong, and the shape of the
error is the damaging kind: it tells a consumer to special-case sdxl and
treat flux2 as an array, which fails on first contact.

Every generator returns a list:

  txt2img_flux    invisible(list(image = img_array, metadata = metadata))
  txt2img_flux2   invisible(list(image = img_array, metadata = metadata))
  txt2img_zimage  invisible(list(image = img_array, metadata = metadata))
  txt2img_sdxl    return(list(image = img_array, metadata = metadata))
  txt2img_sd21    return(list(image = img_array, metadata = metadata))
  txt2vid_ltx2    invisible(result)   # video, audio, sample_rate, latents,
                                      # audio_latents, latent_shape

So sdxl agrees with the other image families rather than diverging from
them, and the corrected contract is simpler than the one it replaces:
unwrap $image across all five image families, $video for ltx. The only
real inconsistency left is return() against invisible(), which changes
auto-printing at the console and nothing else.

The decision not to normalise the shapes stands; only the stated reason
for it was inverted.

Observed at runtime for sdxl, sd21 and ltx during this work, confirmed for
flux2 by the gpu.ctl side, and read from source for flux1 and zimage.
…ixes

Three defects found reviewing the pre-warm fix, all real.

READ THE ALLOCATOR ON THE HANDLE'S OWN DEVICE. cuda_memory_stats() takes
`device` and defaults to cuda_current_device(), so calling it bare reports
whichever device happens to be current rather than the one the handle
bound. resident_load() binds an explicit "cuda:N" precisely so transitions
cannot drift, and a cuda:1 handle deciding from cuda:0's pool would either
skip a pre-warm it needs or repeat one it does not. Added .cuda_index()
to derive the ordinal, with a test for an unqualified device and for a
malformed one that must not become NA and poison the lookup.

MEASURE FREE CACHE, NOT RESERVED. reserved includes live allocations,
which belong to something else and cannot serve this transfer, so it
overstates what is reusable. Now reserved minus allocated. This also
tightened the result: the residual growth across cycles fell from 0.102 GiB
to 0.002 GiB, because the decision is no longer made on bytes that were
never available.

PUT THE MARGIN ON THE FINAL POOL. (bytes - held) * 1.05 asks for 5% of the
GAP, which undershoots the intended target whenever held > 0. The figure
wanted is bytes * 1.05 - held.

Growing the pool is best-effort in the partial case and the docs now say
so rather than implying a guarantee: a request smaller than a free block
already in the cache is served from that block and grows nothing. That is
bounded and harmless -- the onload falls back to the per-tensor path for
the remainder, which is the old behaviour -- and the alternative, asking
for the whole figure to force a new segment, is the accumulation bug this
function exists to avoid. The cold pool, which is the case worth
optimising and the one a broker's first request hits, is unaffected.

THE @return WAS STILL WRONG. It said sdxl alone uses return() while every
other family uses invisible(); txt2img_sd21 uses return() too, which the
commit that introduced the text had listed correctly two paragraphs
earlier. And ltx's `video` and `audio` are produced only when
decode_video/decode_audio are TRUE -- a caller that turns either off gets
a list without the field, not a NULL one, which is why txt2vid_ltx2
indexes it with [[ ]] internally.

VERIFIED, closing the three gaps the reviewer named rather than leaving
them for verification:

  five cycles, release = FALSE, render each cycle
    5.396 -> 5.398 -> 5.398 -> 5.398 -> 5.398 GiB

  cold-start win intact, measured not argued
    cycle 1 activate 2.48 s (2.51 s before the accumulation fix)
    cycles 2-5 activate 0.16 s

  a phase-offloading family is untouched
    flux2, three cycles: state=active, components_on_gpu=0, reserved
    0.002 GiB throughout -- the bulk branch is never taken, so the
    pre-warm is never reached

Suite 1209 assertions, 0 failures.
@TroyHernandez

Copy link
Copy Markdown
Contributor Author

Review addressed — all three findings were real, and the third had a part I had not spotted.

1. Allocator stats read from the wrong device — fixed

cuda_memory_stats() does take device and defaults to cuda_current_device(), so calling it bare reports whichever device is current rather than the one the handle bound. resident_load() binds an explicit "cuda:N" precisely so transitions cannot drift, so a cuda:1 handle deciding from cuda:0's pool would skip a pre-warm it needs or repeat one it does not. Added .cuda_index(), with tests for an unqualified device and for a malformed ordinal that must not become NA and poison the lookup.

The second half of that finding was also right: reserved includes live allocations, which cannot serve this transfer. Now reserved - allocated. That tightened the result as well — residual growth across cycles fell from 0.102 GiB to 0.002 GiB, because the decision is no longer made on bytes that were never available.

2. Shortfall request may not grow the pool — fixed the arithmetic, documented the limit

The margin correction is right and now reads bytes * 1.05 - held; (bytes - held) * 1.05 asks for 5% of the gap and undershoots the target whenever held > 0.

On the deeper point — a request smaller than a free block already in cache is served from that block and grows nothing — that is correct and I have documented it as best-effort rather than implying a guarantee. It is bounded and harmless: the onload falls back to the per-tensor path for the remainder, which is the pre-fix behaviour. The alternative, asking for the whole figure to force a new segment, is exactly the accumulation bug this function exists to avoid. The cold pool, which is the case worth optimising and the one a broker's first request hits, is unaffected.

3. @return still inaccurate — fixed, and worse than reported

txt2img_sd21() uses return() too — which the commit introducing that text had listed correctly two paragraphs earlier, so the doc contradicted its own commit message. And ltx's video/audio are produced only when decode_video/decode_audio are TRUE; a caller that disables either gets a list without the field rather than a NULL one, which is why txt2vid_ltx2() indexes with [[ ]] internally. Both now stated.

Verification

Also closed the three gaps flagged for the verifier, rather than leaving them:

Five cycles, release = FALSE, render each cycle

5.396 → 5.398 → 5.398 → 5.398 → 5.398 GiB

Flat at 3 is not flat at N, so this runs to 5.

Cold-start win intact — measured, not argued from the code

cycle 1   activate 2.48 s     (2.51 s before the accumulation fix)
cycles 2-5 activate 0.16 s

A phase-offloading family is untouched

flux2, three cycles: state=active, components_on_gpu=0, reserved 0.002 GiB

components_on_gpu = 0 while active proves the bulk branch was never taken, so the pre-warm was never reached.

Suite 1209 assertions, 0 failures.

…NEWS

Three follow-ups from the re-review, none functional bugs but the first
was a real inconsistency and the second and third would have shipped
wrong.

ONE TARGET FOR BOTH THE SKIP AND THE SIZE. The skip tested `held >= bytes`
while the growth aimed at `bytes * 1.05`, which put a step in the middle:
3.999 GiB held asked for 0.201 GiB and 4.000 GiB held asked for nothing.
`target` is now computed once and governs both, so the two cannot
disagree. Added a test that holding a hair under and a hair over the raw
need differ by a hair rather than by the whole margin -- the case the old
threshold got wrong.

ROXYGEN ORDERING SWALLOWED THE EXPLANATION. The partial-cache paragraph
was written after `@param device`, so roxygen read it as a continuation of
that tag rather than as description, and the block also carried a
duplicate `@param bytes`/`@param device` pair from before the edit. The
generated page therefore did not contain the best-effort rationale at all
-- the reason the function does not force pool growth was documented in
the source and absent from the docs. Prose now precedes the tags and the
duplicates are gone; verified against the generated Rd rather than
assumed.

man/dot-cuda_index.Rd WAS UNTRACKED. `git add -u` stages tracked files
only, so a newly generated Rd was silently left out of the previous
commit. Staged by explicit path. This is the "never git add -A" rule
biting from the other side: -u is the safe default precisely because it
ignores untracked files, which means new files need naming.

NEWS CARRIED THE SUPERSEDED MEASUREMENTS. It still reported 5.396 / 5.498
/ 5.498 from before the free-cache correction, rather than the current
5.396 / 5.398 across five cycles. Also now records the cold-start figure
and the phase-offload check, since those are the numbers a reader would
otherwise have to take on trust.

Re-verified after the threshold change: five cycles at 5.396 / 5.398 /
5.398 / 5.398 / 5.398 GiB, cold start 2.52 s, warm activations 0.17 s.
Suite 1211 assertions, 0 failures.
@TroyHernandez

Copy link
Copy Markdown
Contributor Author

Re-review addressed. All three were right; the second would have shipped a doc page missing the very rationale it was added to explain.

1. Margin threshold inconsistent — fixed

The skip tested held >= bytes while growth aimed at bytes * 1.05, which put a step in the middle: 3.999 GiB held asked for 0.201 GiB, 4.000 GiB held asked for nothing. target is now computed once and governs both the skip and the size, so they cannot disagree.

Added a test for the discontinuity itself rather than just the new value — holding a hair under and a hair over the raw need must differ by a hair, not by the whole margin.

2. Generated documentation incomplete — fixed, both halves

man/dot-cuda_index.Rd was untracked. git add -u stages tracked files only, so a newly generated Rd was silently left out. Staged by explicit path. This is the "never git add -A" rule biting from the other side: -u is the safe default precisely because it ignores untracked files, which means new files have to be named.

The duplicate @param tags were worse than cosmetic. The partial-cache paragraph was written after @param device, so roxygen read it as a continuation of that tag rather than as description — meaning the generated page did not contain the best-effort rationale at all. The reason the function deliberately does not force pool growth was in the source and absent from the docs, which is the half of that finding that mattered. Prose now precedes the tags, duplicates removed, and I verified against the generated Rd rather than assuming the regeneration fixed it.

3. NEWS carried superseded measurements — fixed

It still reported 5.396 / 5.498 / 5.498 from before the free-cache correction. Now 5.396 / 5.398 across five cycles, plus the cold-start figure and the phase-offload check, since those are otherwise numbers a reader has to take on trust.

Re-verified after the threshold change

five cycles, release=FALSE   5.396 → 5.398 → 5.398 → 5.398 → 5.398 GiB
cold start                   2.52 s
warm activations             0.17 s

Suite 1211 assertions, 0 failures. Working tree clean, nothing untracked.

@TroyHernandez
TroyHernandez merged commit ee9c134 into main Aug 24, 2026
2 checks passed
@TroyHernandez
TroyHernandez deleted the fix/prewarm-accumulation branch August 24, 2026 20:47
TroyHernandez added a commit that referenced this pull request Aug 25, 2026
)

* rformat: normalise a continuation line carried over from #64

* Point safetensors messaging at CRAN 0.3.0 instead of a GitHub build

safetensors 0.3.0 reached CRAN on 2026-08-21 with all four fixes diffuseR
had been routing users around: float8 (mlverse/safetensors#13), bfloat16
write (#11), the >2 GB offset overflow (#14), and empty tensor names
(#10). Every message telling users to install a development build from
GitHub was, as of that date, advising the harder path to a fix already
sitting in `install.packages("safetensors")`.

Five messages carried it: recommend()'s tier note and the graceful
fp8/bf16 fallback (both via the shared helper), flux_quantize()'s two
errors, and the >2 GB read breadcrumb. Plus README's tier table and the
st_caps/reshard docs.

THE PROBES ARE UNCHANGED, and that is the point. They were written as
runtime capability probes rather than a version floor precisely so this
day would need no code change, and they earned it: a 0.3.0 user got the
higher tiers the moment they updated, with nothing in the package to
adjust. They also still cover what a version test cannot -- the fixes
existed for three weeks in builds reporting 0.2.1, so the version number
never distinguished them. No floor added to Suggests for the same reason:
nf4 works on older safetensors, so a stale install costs a tier, not the
model.

TESTS NOW PIN THE REMEDY, NOT JUST THE PHRASING. The existing assertions
checked "best fit for your card", the absence of an em dash, and the file
size in the breadcrumb -- everything except what the message tells a user
to DO. That is how the advice went stale silently and would have stayed
stale: nothing failed. Added guards that the messages name
install.packages and do NOT mention a development version or GitHub, and
checked the guards reject the old wording rather than merely passing on
the new.

.st_fork_note is renamed .st_update_note, since it no longer suggests a
fork. Four call sites, all internal.

recommend()'s returned `fork_suggested` field KEEPS its name. It is part
of a documented return contract that memory_flux.R propagates and five
tests read, so renaming it would break callers for a cosmetic gain. The
docs now say what it means and that the name is historical.

reshard_safetensors() is no longer required to make a large artifact
readable. It stays useful for publishing: its shards load on every
safetensors including older ones, which is what makes a hosted artifact
safe to redistribute.

NEWS records this as a new 0.2.2.7 entry rather than editing 0.2.2.3's
rationale, which described the situation accurately when it was written.

Suite 1220 assertions, 0 failures.

* Finish the safetensors 0.3.0 sweep: docs, and guards with teeth

The first pass fixed the messages and left the documentation, which is
where most of the stale advice actually lived. The shard_bytes help for
flux_quantize(), ltx23_quantize_nf4(), ltx23_quantize_fp8() and
gemma3_quantize_nf4() still described the 1.9e9 default as what "stock
CRAN safetensors" can read, and told users to install a fork for
anything larger. README, the performance-levers vignette, and the
unet_safetensors / download_prebuilt / convert_sd21_pt_to_diffusers
pages carried variants of the same. All now describe the floor as
"older than 0.3.0" and name install.packages("safetensors") as the
remedy.

st_caps.R contradicted itself: the opening said safetensors 0.2.1 lacks
the fixes, while the explanation below correctly noted that fixed
development builds also reported 0.2.1. The opening now says the CRAN
0.2.1 release, which is what was meant, and the probe-don't-pin
rationale survives intact.

The two direct fp8 errors had their wording corrected last commit but
nothing asserted it: both tests only matched "float8", so the old
GitHub-development-build advice would have passed. Both gates now assert
the exact CRAN remedy and reject "GitHub" and "development version".
Checked that these guards are not vacuous by running all four
pre-change messages through them: every one fails all three assertions.

The load-path fixture needs the class. flux_load_transformer() runs
stopifnot(inherits(ckpt, "ltx23_checkpoint")) before the fp8 gate, so a
bare list(format = "fp8") errors on the stopifnot and never reaches the
message under test.

Suite 1229 assertions, 0 failures. man/ regenerates byte-identical from
source. Historical NEWS entries left alone.

* Diagnose the capability that actually failed in the bf16 upgrade note

recommend() gates tiers on .st_can_read(), but the note it produced for
a blocked bf16 tier cited mlverse/safetensors#11, which is the bfloat16
WRITE fix. bfloat16 read worked on CRAN 0.2.1, so a reader that lacks it
is not waiting on #11: the message was well-formed, carried the right
remedy, and pointed at an issue unrelated to the user's failure.

.st_update_note() now takes mode, and both call sites pass the
capability they gated on: recommend() passes "read",
.st_graceful_precision() forwards its own mode. The read path for bf16
drops the issue reference entirely, because no release added bfloat16
read and there is no fix to point at. float8 is untouched in both modes,
since 0.2.1 had neither read nor write for it.

The suite was holding this in place rather than catching it: its
read-mode assertion matched "safetensors#11", so the wrong reference was
pinned by the test. That assertion now matches the read wording. The
remedy-only guards from the previous commit could not have caught this
class of defect, which is the point of the added assertions: checked
that both new ones fail against the old output while the remedy guard
still passes it.

Also stopped making the version the requirement in prose. The shard_bytes
docs for flux_quantize(), ltx23_quantize_nf4(), ltx23_quantize_fp8(),
plus reshard_safetensors() and the README fp8 bullet, said "requires
safetensors 0.3.0 or newer". That contradicts the probe-don't-pin
rationale in the same PR, which turns on capable builds having existed
while reporting 0.2.1. They now name the fix and note where it landed.

Suite 1242 assertions, 0 failures, re-run after the rformat pass against
a fresh install. man/ regenerated from source.

* Finish the version-to-capability vocabulary sweep

Four sites still described the missing capability as a version. The
worst contradicted itself in adjacent sentences: flux_quantize()'s
resident-dtype comment said "safetensors before 0.3.0 cannot write it.
The probe decides, not the version: the fix existed for three weeks in
builds still reporting 0.2.1." Both halves cannot be true.

The others: the >2 GB read breadcrumb told users "safetensors before
0.3.0 overflows", unet_safetensors said "on safetensors older than
0.3.0", and .st_read_or_breadcrumb's comment used "safetensors 0.3.0+"
as shorthand for a reader with the fix. All now name the capability and
say where it landed.

NEWS was narrating the iteration rather than the result: one paragraph
said the docs now require 0.3.0 or newer, the next said that wording was
replaced because it was wrong. Collapsed to the end state.

Left alone deliberately: convert_sd_pt, quantize_gemma3 and
download_prebuilt say sub-2 GB artifacts load on readers older than
0.3.0, which is a compatibility claim that holds for every build, capable
or not. README says "releases before 0.3.0 overflow", which is accurate
as scoped, since no release before 0.3.0 carried the fix.

Suite 1242 assertions, 0 failures. The overflow breadcrumb still carries
2^31, the size, the shard name and the CRAN remedy.
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.

1 participant