Repository navigation
Studio: chunk large Qwen-Image-2.1 attention queries on ROCm - #12335
Conversation
|
Reached Codex review convergence at oobabooga#318. |
|
Review record for #12335, carried over from the mirror it was reviewed on: danielhanchen/unsloth-staging-review#242. Findings are quoted as posted and attributed to the account that posted them. The reactions shown are the triage recorded on the mirror; none were re-applied here. Commits are named as the mirror's, with this PR's equivalent where one was found. Round 1: reviewed mirror
|
|
Tested on a discrete RDNA4 card with Four runs per side: one pair first, then three more in alternating order (off, on, off, on, ...), each a fresh process rendering twice (cold, then warm).
The ranges don't overlap, so the roughly 13% warm speedup is outside run-to-run noise here. Images are near-identical but not bit-identical: off vs on differs by at most 17/255 (mean at most 0.025), while two off runs differ by at most 2/255, so chunking does change the numerics slightly. Setup: Radeon AI PRO R9700 (gfx1201), Windows 11, PyTorch 2.11 / ROCm 7.13, main at Recommend widening the automatic path to all ROCm GPUs. The other eligibility checks (BF16, at least 8192 query tokens, the 32-head layout) already keep it off where it does not apply. Happy to open the follow-up. |
Qwen-Image-2.1's native BF16 attention becomes a bottleneck at large image sizes on Radeon 8060S. With Q4_K_M and Auto memory at 1536x1536, the measured slowdown remains even when the transformer is resident. Splitting independent query rows reduces the attention runtime while retaining every key and value.
Fix
Install a pipeline-local Qwen-Image-2.1 processor that dispatches eligible queries in 512-row chunks. Enable it automatically only on the measured gfx1151 architecture, for native unmasked BF16 inference with at least 8192 query tokens and the model's 32-head/4096-wide attention layout.
Keep the stock processor for prefill, masks, smaller requests, other attention backends, parallel execution, gradients and compilation. Reuse Diffusers' QKV/cache preparation and output projections. Speed Off disables installation; UNSLOTH_QWEN_IMAGE_ROCM_QUERY_CHUNKING=0 opts out, and =1 allows explicit testing on other ROCm GPUs.
Tests
Live A/B against main at 4700dbc on Radeon 8060S (gfx1151), PyTorch 2.11 / ROCm 7.13, Qwen-Image-2.1 Q4_K_M, Auto memory, seed 42 and guidance 1.0. The planner sees 20 GiB total and 19 GiB free, with a 19 GiB PyTorch allocation cap. Both attention variants were tested without the encoder fix.
Both OS jobs pass. All 16 attention comparisons are pixel-identical, including the 40-step images and image-conditioned outputs. At 1024x1024, both operating systems make zero chunking calls and preserve all three images and warm runtime. Peak GPU allocation is unchanged. The 40-step median decode step falls from 48.99 to 12.38 seconds on Windows and from 16.95 to 10.59 seconds on Linux; the unmodified first step stays around 45 and 20 seconds respectively.
The original 292 related tests passed; after the log clarification, all 174 processor/speed tests pass again. Coverage includes full-key/value chunk boundaries, cached-prefix and projection parity, backend selection, device/runtime guards and local installation. Formatting and diff checks pass. Post-A/B changes are limited to the load-log wording; attention computation is unchanged.
Live A/B run, both OS jobs passed
Scope
The automatic path is limited to gfx1151. RX 7900 XT and other ROCm architectures remain untested. The hardware probes in the run artifacts report shared-memory GPUs and 124.95 GiB of host RAM on Linux / 127.15 GiB on Windows, so the allocation cap does not emulate a discrete GPU's transfer costs. This does not change Auto placement, the Windows encoder fallback or cold MIOpen/VAE initialization. The processor uses helpers from Studio's pinned Diffusers version.