ROCm on Windows: keep VRAM resident and bound attention and VAE memory - #304
Pfannkuchensack wants to merge 18 commits into
Conversation
With pytorch_cuda_alloc_conf unset, a ROCm build on Windows now runs with expandable_segments:True, so fragmented VRAM no longer pushes allocations into shared system memory. An explicit allocator setting or allocator env var still wins; docs and generated config artifacts updated.
On ROCm under Windows, free VRAM is now capped by the WDDM budget (gdi32 D3DKMTQueryVideoMemoryInfo), past which Windows pages allocations into system memory instead of failing them. torch's figure there ignores other processes and runs up to the physical total; other platforms are unchanged.
After each session, a worker on a Windows ROCm device reads the process's shared GPU memory (PDH) and warns once per episode when more than 512 MiB of it sits in system memory, where every generation that touches it slows down. Adds a low-VRAM docs section on AMD GPUs under Windows.
The ROCm SDPA guard now splits any math-kernel call over 1 GiB of scores into head groups (bitwise identical) or query rows (one oversized head), and working-memory estimates price one chunk instead of the whole matrix. Z-Image 1024px peaks at 0.5 GiB instead of 8 GiB per attention on an RX 9060 XT; renamed to install_rocm_sdpa_guard.
The Krea-2 working-memory estimate now adds the score matrix where the build materializes it (no fused kernel), priced from the loaded model's heads and the attended sequence; one 1 GiB chunk on ROCm, nothing on CUDA. Before, Krea-2 fp8 on an RX 9060 XT loaded the whole transformer and paged each 12.9 GiB attention into system RAM.
On ROCm the FLUX.1 VAE (also Z-Image's) now budgets 3600 decode / 2750 encode bytes per pixel·byte instead of 2200/1100, measured 3451/2688 on an RX 9060 XT: the same convolution stack and numbers as the FLUX.2 VAE, so both share one table. The estimate follows the VAE's compute device, so a cpu_only VAE keeps the cuDNN column.
…code The FLUX.1, Z-Image, Qwen-Image (Krea-2) and Wan decodes now tile when the untiled estimate exceeds 90% of the VAE's GPU memory (Anima keeps its measured 70%), instead of relying on an OOM that Windows paging never raises. New auto_tiled_decode setting (default on) turns it off; force_tiled_decode and the tiled field still win.
- Cap free VRAM at the WDDM budget minus live allocations: HIP keeps freed memory that CurrentUsage still counts - Empty the allocator cache after each offload under expandable segments so the cache stops once enough is free - Hardware tests for the headroom after a free and for the PDH lookup
- Skip the allocator default once torch is imported; stop the test fixture leaking it - Warn only about memory paged across two sessions, with actionable advice - Leave K/V broadcasts the chunks cannot slice whole; cover masked row chunks and Krea-2 CFG/regional pricing
- Price and pre-tile FLUX.1 decodes with a diffusers-layout VAE - Log up-front tiling at debug; document what auto_tiled_decode off means - Emit POSIX path defaults in the generated docs settings
github-actions
Bot
added
python
python-tests
backend
services
frontend
invocations
docs
labels
|
I can verify this on Linux ROCm, but I don't have a Windows boot on my AMD rig. |
|
Passes all the ROCm tests, including |
These tests fail on purpose. They pin behaviour the current up-front tiling rule does not have, and should go green with the fix. - A Qwen-Image/Krea-2 decode whose measured peak fits the card is tiled anyway, because the gate is fed the padded reservation figure (a flat 5500 B/pixel-byte) instead of the expected peak (3273 measured at 1536px). Tiling is not pixel-identical, so this silently changes output for images that would have decoded in a single pass. - should_pretile_vae_decode compares against the card's nameplate total, not the Windows video-memory budget this PR added TorchDevice.cuda_mem_get_info for, so a decode Windows will page into system memory is left untiled.
Adversarial reviewReviewed against I've pushed 33e16af with three intentionally-failing tests for the two findings I think are blockers. CI will be red by design — they pin behaviour the current rule does not have and should go green with the fix. Happy to drop that commit if you'd rather not carry red tests. 1. Reservation headroom is reused as a tiling trigger, so decodes that fit are silently tiled
The estimator constants are deliberately over-provisioned ("max observed + ~8% headroom"). Over-reserving used to cost only cache eviction. Comparing that padded figure against a 90% line converts the conservatism into a non-pixel-identical output change, with no error and nothing above DEBUG in the log. Qwen-Image is the sharp case, because its ROCm decode constant is a flat
On a 24 GiB ROCm card a 1536² Qwen-Image decode that previously ran untiled and succeeded now tiles. Krea-2 decodes through this node too. On CUDA FLUX.1 the crossovers are benign (shipped 2200 vs measured 2185), so this is specifically Qwen/Krea-2 on MIOpen. The reservation is right as it is — it's the tiling decision that shouldn't be carrying reservation headroom. Test: One knock-on: 2. The gate measures total VRAM on the platform this PR taught the cache not to trust
Trigger: 16 GiB RX 9060 XT, a browser holding ~3 GiB so the budget is ~12.5 GiB. FLUX.1/Z-Image at 1408² estimates 13.29 GiB — under the 14.4 GiB threshold, over the budget. No pre-tile, no OOM (Windows pages instead), so the retry at Also reachable on Linux/CUDA: Test: Other material findings (no tests pushed)3. The expandable-segments offload fix is inert on multi-GPU. 4. The "bitwise identical" claim for head-group chunking is false. 5. 6. 7. SD1/SDXL, SD3 and CogView4 left on the cuDNN constant. Smaller
Verified as claimed
I could not complete a full-suite run — this box was saturated and |
…droom - Compare against the Windows video-memory budget where there is one, not the card's total - Qwen-Image ROCm constants now follow the bounded math attention this branch ships: measured 2650-2772 decode and 1541-1552 encode across 512-2048px, against 5500/6300 fitted while a call built its whole score matrix - Calibrate with the attention guard installed, so the script measures the path the app runs
…ld up - Credit freed bytes while a peer device defers empty_cache, so multi-GPU stops over-unloading - Never let the budget or paged-bytes lookups raise, and close adapter handles if enumeration throws - Detect a ROCm build from torch's version.py too, so a locally built wheel gets the allocator default - Correct the bitwise-identical claim for head-group chunking
…kend - auto_tiled_decode no longer disables Anima's own rule, which is a speed optimization, not an OOM fallback - SD1/SDXL, SD3 and CogView4 take the MIOpen constant on ROCm, like the FLUX.1 autoencoder they share a stack with - Encode nodes place their tensors on the VAE's device, matching the decode nodes
…eiling - Keep the padded Qwen-Image reservations; the tiling decision reads the measured curve instead, so a decode whose real peak fits the card is no longer tiled by a reservation's headroom - Compare against the video-memory budget plus what this process can release: Windows halves the budget once the process passes ~12 of 16 GiB, and the bare figure tiled a 7.0 GiB decode that fits
|
Thanks — this was a genuinely useful review, and the two tests made the first finding much easier to act on. All three are green now, and everything else you raised is either fixed or answered below. One of my first attempts at finding 1 was wrong in a way worth recording, so I've left that in rather than quietly dropping it. 1. Reservation headroom as a tiling trigger — fixed the way you framed itImplemented as you described it: the reservation keeps its headroom, and the tiling decision no longer carries it. I went the wrong way first and want to record why, since it nearly shipped: I re-measured on an RX 9060 XT (gfx1200, torch 2.12+rocm7.14, fp16) and got a flat 2650-2772 decode / 1541-1552 encode across 512²..2048², attributed that to the score-matrix chunking this branch adds, and made the constants conditional on the guard. That attribution does not survive your own table. At 512² the gap between the two cards is 1.24 GB while the entire score matrix there is ~0.3 GB, and at 1536² an unbounded matrix would be ~23 GB against a measured 15.4 GB total — so the W7900 run was never dominated by an unbounded score matrix in the first place. The spread is the card, not the guard, and shipping the lower pair would have under-reserved by 1.9x (decode) to 3.8x (encode) on the hardware the shipped figures were measured on, with Qwen Image Edit's encode having no tiled retry to fall back on. Reverted; the constants are unchanged. The gfx1200 numbers stay in the comment, because that spread is the actual argument for your finding: a reservation that has to cover a 1.9x card difference is not a figure to decide non-pixel-identical output with. One thing I did keep: 2. The gate measured the nameplate total — fixedThe gate consults the budget now. My first version compared against the bare budget, and an end-to-end run on the card caught what that does mid-session: Windows holds the budget at 15.09 GiB until the process passes about 12 of 16 GiB and then halves it to 7.62 — below what we already hold. With models resident, the gate read 7.6 GiB, drew a 6.9 GiB line, and tiled a 7.0 GiB Z-Image decode that fits. Output stopped being pixel-identical (PSNR 42 dB) and the decode went from 0.89 s to 1.5 s: your finding 1, reintroduced at the other end. So the ceiling is the budget plus what this process itself holds, capped at the card: the cache evicts models to honour the reservation, so that memory is available to the decode. Measured:
Your test passes unchanged (it patches The 3. Offload fix inert on multi-GPU — fixed
4. "Bitwise identical" — corrected, thank youYou are right, and the CPU test was validating the claim on the one platform it isn't about. The comment and the docstring now say what actually holds: exact in exact arithmetic, bitwise where the kernel's arithmetic does not depend on the batch it runs over (CPU), and within a bf16 ulp on a card whose batched GEMM re-tiles with the group count — with your 0.0005 measurement named. The test keeps its exact assertion and says why it is exact there. The PR description is updated too. 5.
|
# Conflicts: # tests/app/invocations/test_z_image_tiled_decode.py
main added a case reading _FLUX2_VAE_SCALING_CONSTANTS, which this branch merged into _FLUX_VAE_SCALING_CONSTANTS.
The new device argument made a device-less estimate resolve the session device, so the same call answered MIOpen on a GPU box and cuDNN on CI. A named device still answers for that device. Pin the convolution backend separately from the attention guard in the ROCm estimator tests, which only diverged on a ROCm build.
This branch has not been deployed
2 participants
Summary
Generations on a ROCm build of PyTorch under Windows ran 4-5x slower than the hardware allows, and Krea-2 fp8 did not run at all. Windows never fails an allocation that does not fit: it moves it into shared system memory and keeps it there, so nothing in the log says why a run crawls. Measured on an RX 9060 XT (gfx1200, torch 2.12+rocm7.14.1): Z-Image Turbo nvfp4 at 1024px denoised in 144/139/186 s over three runs; it now takes 34 s, with pixel-identical images. Krea-2 fp8 produced no step in two minutes; it now takes 60 s per image. Z-Image at 1536px is possible for the first time.
Four causes, addressed in order:
expandable_segments:True, set before torch is imported and only when no allocator variable is configured. Any explicitpytorch_cuda_alloc_confstill wins.torch.cuda.mem_get_infois the device total minus this process's own usage there: it ignores other processes, and Windows starts paging once the process passes its WDDM budget (15.09 of 15.92 GiB alone; lowered within a second when another GPU process starts). The cache's free figure is now capped by that budget throughD3DKMTQueryVideoMemoryInfo, and a worker says so when memory stays in system RAM across two sessions (PDHGPU Process Memory\Shared Usage). Both are best-effort and silent everywhere else.auto_tiled_decodesetting. The FLUX.1 autoencoder's working-memory constants follow the convolution backend, as the FLUX.2 ones already did: MIOpen needs 3600 B/pixel-byte to decode where cuDNN needs 2200.Items 3 and 4 change behaviour on every platform: math-kernel attention runs in chunks on Linux ROCm too, and decodes that would take most of the card are tiled everywhere. On CUDA the estimates and the fused path are untouched.
Documentation says only what is true: AMD remains supported on Linux only; the Windows ROCm notes are hints.
Related Issues / Discussions
None.
QA Instructions
Gates
uv run --no-sync pytest -n 6: 8905 passed, 174 skipped, 9 xfailed.uv tool run ruff@0.11.2 check/format --checkon the changed paths: clean.pnpm -C docs build: 326 pages, link validation clean. Generatedopenapi.json,schema.tsanddocs/src/generated/settings.jsonwere regenerated with their generators;generate_docs_json.pynow emits Path defaults with.as_posix(), so regenerating on Windows no longer flips the separators and failscheck-docs-data.-m slow, RX 9060 XT): 5 passed, 1 skipped. The skip istest_a_fused_kernel_is_still_wrong_for_the_wide_head: no fused backend runs that shape on gfx1200 at all, so there is nothing to compare.tests/app/invocations/test_flux2_working_memory.pyfail depending on the build (5 on gfx1200 / torch 2.12+rocm7.14 here); they fail identically on the base commit, as they assume a CUDA/Linux backend.End-to-end, from this branch's code with no runtime patches
RX 9060 XT, 16 GB (port 9091):
RTX 4090, 24 GB (port 9090), for the CUDA regression:
Measurements behind two design decisions
CurrentUsage. HIP under Windows keeps 1-2 GiB after a free plusempty_cache()and hands it to the next allocation, whileCurrentUsagestill counts it: 4 GiB allocated, then freed, leavesCurrentUsageat 1.14 GiB withmemory_reservedat 0 and torch's free figure back at its starting value; the next 2 GiB cost only 1 GiB of new usage. Capping againstCurrentUsagewould hide what an offload just freed.delfreed 0.00 GiB,empty_cache()freed 2.99 GiB), so the loop saw no progress and unloaded every unlocked model.Not verified
Linux ROCm (the query-row chunks in the VAE and the FLUX.1 MIOpen constant), multi-GPU under Windows ROCm, other Windows ROCm builds, MPS/XPU pre-tiling, and small CUDA cards on real hardware. The Qwen-Image VAE's ROCm constants are unchanged (measurable only on Windows here), and
auto_detect_slice_sizestill usestorch.cuda.mem_get_infodirectly.Review
Three independent read-only reviews ran over the full candidate with the same base and acceptance criteria: correctness/spec conformance, architecture/operational safety/performance, and test value/product quality. Material findings resolved:
settings.jsonregenerated on Windows flipped two path defaults and would have failed the docs job; the generator now emits POSIX paths.PYTORCH_CUDA_ALLOC_CONFinto later tests, which made a model-cache test fail depending on file order.AutoencoderKL, which is the exact failure this PR targets; it is now priced and pre-tiled like the Z-Image node.auto_tiled_decode=falsefor Wan and Anima, a Wancpu_onlytest that passed regardless of the code, PDH struct offsets plus a hardware test, and a test that the worker loop calls the warning at all.The final blocker-only review over the resulting candidate reported no blockers.
Compatibility / Rollout
auto_tiled_decode(default true); generated OpenAPI, frontend types and docs settings regenerated.wddm.pyis new and best-effort: every failure answers "unknown" and leaves the previous behaviour in place.Checklist
What's Newcopy (if doing a release after this PR)