Repository navigation
Studio: chunk large Qwen-Image-2.1 attention queries on ROCm - #242
Closed
danielhanchen wants to merge 1 commit into
Closed
danielhanchen wants to merge 1 commit into
danielhanchen wants to merge 1 commit into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
Owner
Author
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.
Mirror of unslothai/unsloth#12335 by @oobabooga, opened for review only.
Do not merge here. Fixes agreed on this mirror must be applied to the source PR.
Source description digest at mirror time:
a644c8dfb261