Skip to content

VAE: tile the encoder when the device cannot hold its graph, and retry with smaller tiles when a graph does not fit - #125

Open
timoncool wants to merge 4 commits into
ServeurpersoCom:masterfrom
timoncool:vae-tile-fallback
Open

timoncool wants to merge 4 commits into
ServeurpersoCom:masterfrom
timoncool:vae-tile-fallback

Conversation

@timoncool

@timoncool timoncool commented Oct 22, 2026 •

Copy link
Copy Markdown

The decoder checks the device's free memory before it builds a graph and tiles at 256 frames when a graph of min(T, chunk) frames would not fit (about 10 MB of activations a frame; HOT-Step #186: a 726-frame graph asked for 6.78 GB and took a 4 GB card down). The encoder had no such check: a cover, a timbre reference or an understand pass over a long source built one graph per 1024-frame tile, or one for the whole source when it was shorter than the chunk, whatever the card had free.

  • The encoder gets the same free-memory check and tiles at 256 frames when the graph would not fit. Its estimate is measured: the encoder's compute buffer (ggml_backend_sched_get_buffer_size after a successful ggml_backend_sched_alloc_graph) grows by exactly 6,389,760 bytes a latent frame with no constant part (96 to 1024-frame tiles, --vae-chunk 128/256/512/1024 on a 173 s song, RTX 4090, CUDA, F32 activations); the check takes that plus a tenth.
  • When ggml_backend_sched_alloc_graph fails, the decoder and the encoder no longer give up: they retry the whole pass with half the tile, down to 64 frames. Only a 64-frame tile that still does not fit is FATAL ("graph alloc failed for N frames, the smallest tile").
  • With enough memory nothing changes: same tiles, same output.

Seen on a 6-minute source (8986 latent frames) with a graph for the whole track, which is what a 1024-frame tile is to a small card: before, [VAE-Enc] Graph: 565 nodes, T_audio=17253120 then CUDA error: invalid argument; after, the encoder tiles at 256, 71 tiles, and the cover renders. On the 173 s song with --vae-chunk 9000: A 4333-frame encode graph needs about 28.36 GB, 22.34 GB free: tiling at 256, 34 tiles.

Known limit of the check: on Windows with CUDA (WDDM), the free memory the backend reports does not count what other processes hold, since WDDM pages their memory out. The estimate then sees almost the whole card, so the check does not tile for memory taken by other programs, and as long as WDDM pages instead of refusing, the graph allocation does not fail either, so the retry does not start. Vulkan reports a budget that does include other processes.

Summary by CodeRabbit

  • Bug Fixes
    • Improved encoding and decoding of large inputs on devices with limited graphics memory. Processing can now reduce tile sizes and retry when a tile does not fit, helping prevent memory-related failures.
    • Large workloads are more likely to complete without requiring users to adjust settings manually.

… chunk) frames - the untiled short track and the 1024-frame tile alike, about 10 MB a frame (after HOT-Step #186)
…and both halve the tile down to 64 when a graph does not fit
@coderabbitai

coderabbitai Bot commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Source mode: Auto
  • Repository: Unavailable
  • Branch: Unavailable
  • File: No configuration file from this repository applied
  • Commit: Unavailable
  • Layers: defaults
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 64c3bc35-0b4a-43ae-932b-3ac098ddbf3d
📥 Commits

Reviewing files that changed from the base of the PR and between de858d8 and 98303bd.

📒 Files selected for processing: 2
  • src/vae-enc.h
  • src/vae.h

Included review availability: This review used your included allowance. Your plan provides up to 2 included reviews per hour; 0 remain after this review.

📝 Walkthrough

Encoder and decoder graph allocation failures now return VAE_GRAPH_NO_FIT, which their APIs propagate. Tiled encoding and decoding use separate tile helpers and wrappers that can reduce tile size and retry after graph-capacity failures. The wrappers can also use reported device memory to adjust initial tile sizing.

Changes: Encoder and decoder graph-capacity handling

Layer / File(s) Summary
Graph-capacity result propagation
src/vae.h, src/vae-enc.h
Graph allocation failures now reset the scheduler and return VAE_GRAPH_NO_FIT. Encode and decode propagate this result.
Tile helper result handling
src/vae.h, src/vae-enc.h
Tile processing uses separate helpers with explicit tile parameters. The helpers propagate VAE_GRAPH_NO_FIT from graph computation.
Memory-based sizing and retries
src/vae.h, src/vae-enc.h
The wrappers can reduce initial tile sizing based on reported device memory. When a tile graph does not fit, they halve the tile size until the minimum limit is reached.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant vae_ggml_decode_tiled
  participant vae_ggml_decode_tiles
  participant vae_ggml_decode
  participant Backend
  Caller->>vae_ggml_decode_tiled: Request tiled decode
  vae_ggml_decode_tiled->>Backend: Query free device memory
  vae_ggml_decode_tiled->>vae_ggml_decode_tiles: Process tiles with selected size and overlap
  vae_ggml_decode_tiles->>vae_ggml_decode: Compute a tile
  vae_ggml_decode->>Backend: Allocate graph
  Backend-->>vae_ggml_decode: Return allocation result
  vae_ggml_decode-->>vae_ggml_decode_tiles: Return tile result
  vae_ggml_decode_tiles-->>vae_ggml_decode_tiled: Return VAE_GRAPH_NO_FIT or another result
  vae_ggml_decode_tiled->>vae_ggml_decode_tiles: Retry with a smaller tile after VAE_GRAPH_NO_FIT
Loading

Suggested reviewers: serveurpersocom

Merge Risk: ⚪ Minimal · up to 98303

No merge-blocking issue was identified in the graph-capacity retry paths; the change is ready for normal checks.

Pre-merge checks: Passed 5
Check name Status Explanation
Description Check Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check Passed The title clearly describes the main changes: encoder tiling based on device memory and retries with smaller tiles when graph allocation fails.
Docstring Coverage Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 9 functions across 2 files.
Linked Issues check Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create a new PR
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

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