helexa/helexa
1
0
Fork 0
Code Issues 104 Pull Requests Actions Packages Projects Releases Wiki Activity

Z2 — memory discipline + high-res: flash-attn mask fix, tiled VAE decode (+ upstream candle PRs) #199

New Issue
Closed
opened 2026-07-29 14:27:45 +00:00 by grenade · 1 comment
grenade commented 2026-07-29 14:27:45 +00:00
Owner
Copy Link

Part of the text-to-image epic #197. Unlocks ≥1536² and shrinks the working set; both fixes were proven by hand on beast during the 2026-07-29 bench.

Context

Two upstream candle defects cap resolution today (bench findings #2 and #3 in #197):

  1. flash-attn is dead code in z_image: ZImageTransformer2DModel::forward always wraps attention masks in Some(...), and attention_cuda falls back to naive attention whenever a mask is present. At batch-1 the mask is all-ones. Passing None unlocked the full 9-step denoise at 1536² and 2048² on a 5090 (steps: 2.0s / 4.4s).
  2. VAE decode OOMs at ≥1536²: naive mid-block attention over a 36k+-token sequence (latent 192²+) generates tens of GB of f32 transients — the VRAM trace showed a +9.6GB spike into a >6GB failed allocation after the DiT was freed.

Scope

  • Mask elision: detect all-ones masks (or thread Option<&Tensor> through properly) so the flash path engages at batch-1. Correct general fix: skip the mask when every element is valid, keep the fallback for real padding.
  • Tiled VAE decode: decode latents in overlapping tiles with blending (diffusers enable_vae_tiling semantics), bounding decode transients regardless of resolution. Trigger above a latent-size threshold; full-frame below it.
  • Upstream both to candle as PRs; carry as vendored patches (workspace [patch] on the candle fork) until merged. Given mistral.rs-era upstream-responsiveness history, plan for the vendored state to persist; keep the patch surface minimal.
  • Resolution policy: with both fixes, publish supported bounds (multiple-of-16 dimension rule; practical ceiling set by flash-attn seq len + tile decode, not OOM). Reject out-of-bounds requests with a clean error, never an OOM-poisoned worker.
  • Re-bench 1024²: flash-attn should also improve the 0.68s/step baseline; record the delta.

Acceptance

  • 1536² and 2048² generations complete on a 32GB card within the staged-lifetime budget; a deliberately huge request is rejected cleanly (no worker poisoning); candle PRs opened and linked here.

Depends on

Z1 (#198).

Part of the text-to-image epic #197. Unlocks ≥1536² and shrinks the working set; both fixes were proven by hand on beast during the 2026-07-29 bench. ## Context Two upstream candle defects cap resolution today (bench findings #2 and #3 in #197): 1. **flash-attn is dead code in `z_image`**: `ZImageTransformer2DModel::forward` always wraps attention masks in `Some(...)`, and `attention_cuda` falls back to naive attention whenever a mask is present. At batch-1 the mask is all-ones. Passing `None` unlocked the full 9-step denoise at 1536² and 2048² on a 5090 (steps: 2.0s / 4.4s). 2. **VAE decode OOMs at ≥1536²**: naive mid-block attention over a 36k+-token sequence (latent 192²+) generates tens of GB of f32 transients — the VRAM trace showed a +9.6GB spike into a >6GB failed allocation *after* the DiT was freed. ## Scope - **Mask elision**: detect all-ones masks (or thread `Option<&Tensor>` through properly) so the flash path engages at batch-1. Correct general fix: skip the mask when every element is valid, keep the fallback for real padding. - **Tiled VAE decode**: decode latents in overlapping tiles with blending (diffusers `enable_vae_tiling` semantics), bounding decode transients regardless of resolution. Trigger above a latent-size threshold; full-frame below it. - **Upstream both to candle** as PRs; carry as vendored patches (workspace `[patch]` on the candle fork) until merged. Given mistral.rs-era upstream-responsiveness history, plan for the vendored state to persist; keep the patch surface minimal. - Resolution policy: with both fixes, publish supported bounds (multiple-of-16 dimension rule; practical ceiling set by flash-attn seq len + tile decode, not OOM). Reject out-of-bounds requests with a clean error, never an OOM-poisoned worker. - Re-bench 1024²: flash-attn should also improve the 0.68s/step baseline; record the delta. ## Acceptance - 1536² and 2048² generations complete on a 32GB card within the staged-lifetime budget; a deliberately huge request is rejected cleanly (no worker poisoning); candle PRs opened and linked here. ## Depends on Z1 (#198).
grenade added this to the Text-to-image (Z-Image) milestone 2026-07-29 14:27:45 +00:00
grenade added the p2-next label 2026-07-29 14:29:56 +00:00
grenade commented 2026-07-29 16:13:26 +00:00
Author
Owner
Copy Link

Landed in two parts: neuron-side tiled VAE decode via PR #206 (64-latent tiles, 16-latent seam overlap, weighted-accumulation blending — seam-free at 1024²/1536²/2048², verified visually), and the candle-side mask elision via PR #209 (pin to grenade/candle branch z-image-mask-elision; upstream PR huggingface/candle#3798 — drop the pin when it lands).

Live results on beast (5090, flash-attn build): 1536² = 21.2 units (denoise 7.2 s, decode 5.3 s), 2048² = 37.7 units (denoise 15.1 s, decode 8.4 s) — both OOM'd pre-fix. One additional finding fixed on the #204 branch: mixed-resolution request sequences fragment the cudarc pool until a generation that would fit a clean slate OOMs; the worker now trims the pool after every image job.

Resolution policy note: flash-attn ships on the blackwell flavour only (#95's staged rollout), so ≥1536² is a beast-tier capability today; ada/ampere serve ≤1024². Extending flash-attn to ada is #95's call, not this issue's.

Landed in two parts: neuron-side tiled VAE decode via PR #206 (64-latent tiles, 16-latent seam overlap, weighted-accumulation blending — seam-free at 1024²/1536²/2048², verified visually), and the candle-side mask elision via PR #209 (pin to `grenade/candle` branch `z-image-mask-elision`; upstream PR huggingface/candle#3798 — drop the pin when it lands). Live results on beast (5090, flash-attn build): **1536² = 21.2 units (denoise 7.2 s, decode 5.3 s), 2048² = 37.7 units (denoise 15.1 s, decode 8.4 s)** — both OOM'd pre-fix. One additional finding fixed on the #204 branch: mixed-resolution request sequences fragment the cudarc pool until a generation that would fit a clean slate OOMs; the worker now trims the pool after every image job. Resolution policy note: flash-attn ships on the **blackwell** flavour only (#95's staged rollout), so ≥1536² is a beast-tier capability today; ada/ampere serve ≤1024². Extending flash-attn to ada is #95's call, not this issue's.
Sign in to join this conversation.
No Branch/Tag Specified
Branches Tags
Labels
Clear labels
out-of-scope
Outside helexa's constraining principles: near-frontier models on consumer hardware, served transparently and predictably
 p1-now
Active pain or unblocking work — pick up first
 p2-next
Engine levers — start once P1 baselines exist
 p3-later
Dependent on earlier work, or demand-driven
No Label p2-next
Milestone
No items
No Milestone Text-to-image (Z-Image)
Projects
Clear projects
No project
Assignees
Clear assignees
grenade
No Assignees
1 Participants
Notifications
Due Date
No due date set.
Dependencies

No dependencies set.

Reference: helexa/helexa#199
Delete Branch "%!s()"

Deleting a branch is permanent. Although the deleted branch may continue to exist for a short time before it actually gets removed, it CANNOT be undone in most cases. Continue?