Qwen Image 2.1 transformer is incompatible with torch.compile #14821

Activity

1. abel1502

   added

   bugSomething isn't workingSomething isn't working

   on Sep 20, 2026
2. chinmayirm commented on Sep 22, 2026

   

   chinmayirm

   on Sep 22, 2026

   Last edited by chinmayirm

   More actions

   Hi @abel1502, thanks for the report and the repro.

   I agree with the diagnosis. pipe.transformer.compile(..., fullgraph=True) traces QwenImage21Transformer2DModel.forward, which calls QwenImage21Rope.forward, and that does is_image_token = image_pad_mask.tolist() (transformer_qwenimage21.py around line 686). The traceback matches the source on main. The inductor flags in the repro are unrelated.

   Root cause:

   image_pad_mask is bool. It is torch.repeat_interleave(img_mask[0], repeats), and the pipeline builds img_mask as a bool tensor. Dynamo only traces Tensor.tolist() for int8, int16, int32, and int64 (gb0109), so a bool mask fails before the rest of the forward is traced. fullgraph=True turns that into a hard error. This is QwenImage21Rope, which runs for both attention processors.

   Casting the mask with .to(torch.int32).tolist() only satisfies the dtype check. Integer tolist() is lowered to one .item() per token, and the Python list.index / range walk then branches on those values. That is still a host sync and still not one graph. This file already treats tolist() as a device sync in _qwenimage21_prefix_segments.

   The same forward has three more host-dependent spots, so fixing RoPE alone still leaves this repro unable to compile:

   1. build_token_metadata uses nonzero, so the index vector length depends on mask values.
   2. use_kv_cache defaults to True, so step 0 is "extract" and later steps are "cached". prefix_len = int((~target_token_mask).sum()) syncs a scalar on every step and then slices with it.
   3. On the default QwenImage21AttnProcessor, _qwenimage21_prefix_segments calls tolist() on the whole prefix.

   compile_repeated_blocks() does not cover this. RoPE runs in the parent forward, and _repeated_blocks is only QwenImage21TransformerBlock.

   Proposed approach:

   img_shapes already gives each block's (frame, height, width) as Python ints. The mask only says where those blocks sit between text tokens. The current walk is:

   • text tokens get frame = height = width, advancing by 1
   • an image block shares one frame id and a zero-centered height/width grid
   • after the block, the position advances by max(height, width)

   Build that with static-shape tensor ops. Let mask be 0/1 and seq_len = mask.shape[0]:

   • img_before = cumsum(mask) - mask
   • text_before = cumsum(1 - mask) - (1 - mask)
   • for each block, once seen image tokens have been consumed, add max(height, width) wherever img_before >= seen
   • frame_index = text_before + that advance

   Checked against the current Python walk for text/image/text, adjacent images, and a trailing text tail. For example, mask F F T T T T F with one 2x2 block produces frame ids [0, 1, 2, 2, 2, 2, 4].

   The centered grid does not need the mask. Per block:

   start = -(length - length // 2)
   torch.arange(start, start + length)

   Repeat height across width and width down height, then concatenate in block order. masked_scatter writes that static vector onto the True positions. height_index and width_index start as clones of frame_index, so text tokens keep the frame id on all three axes. torch._check(mask.sum() == image_token_count) keeps the scatter length static for Dynamo. Eager keeps today's ValueError when img_shapes and the mask disagree.

   The other three spots, in the same change:

   1. The image-token count is already sum(math.prod(shape) for shape in img_shapes). Build image_ids and target_token_mask with masked_scatter of a static id vector, and use the same torch._check instead of nonzero.
   2. The pipeline appends the target as a suffix, so prefix_len = seq_len - math.prod(img_shapes[0][-1]) matches int((~target_token_mask).sum()) with no .item().
   3. Prefix segment count is static: one text gap and one image block per condition image, then the text before the target. Block starts come from argmax on the image cumsum. The SDPA loop stays a Python loop over that static count, with 0-dim tensor slice bounds. A zero-length text gap (two image blocks with no text between them, which build_token_metadata already calls out) is an empty slice, and torch.cat drops it.

   QwenImage21Rope also rebinds self.freqs with .to(device) on every forward. Those tables are fixed at init and are not buffers, so they stay unregistered. QwenEmbedRope._get_device_freqs already caches this transfer behind lru_cache_unless_export. I would use that here so the first call moves them and later calls do not.

   Tests:

   On the tiny config already in tests/models/transformers/test_models_transformer_qwenimage21.py:

   • Eager RoPE matches the current Python implementation for text/image/text, adjacent images, odd and even spatial sizes, and a target-only sequence.
   • torch.compile(model, fullgraph=True) matches eager, then a second call with error_on_recompile=True.
   • Same for QwenImage21AttnProcessor and QwenImage21FlexAttnProcessor.
   • KV cache "extract" then "cached" under fullgraph=True.
   • Add TorchCompileTesterMixin to this model. The Qwen-Image 1.x transformer tests already have it. This 2.1 file does not, which is consistent with compile being broken.

   I would be glad to open a PR for this once the approach looks right.

   ❤️React with ❤️1Reacted by abel1502