fix _flash_3_varlen_hub mask handling

[zhtmike]

What does this PR do?

As discussed in #13809, we fixed the mask handling in separate PR.

This PR is to fix the non-contiguous mask handling in #14114

Fixes # (issue)
Fix#14114

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you self-review the diff against .ai/review-rules.md?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

@sayakpaul

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[sayakpaul]

Thanks. How can we test it?

[zhtmike]

Thanks. How can we test it?

Seem there is no such test for different backend, should we write a Mixin for different attention backend, like AttentionBackendTesterMixin?

[sayakpaul]

Or maybe a standalone reproducer that passes on this PR branch and fails on main?

[zhtmike]

Or maybe a standalone reproducer that passes on this PR branch and fails on main?

can be verified in #14114

[sayakpaul]

I guess what I am more worried about is whether this PR is breaking existing functionalities and other supported pipelines.

[github-actions]

Hi @zhtmike, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Fixes #1234) to the PR description so the issue is linked. See the contribution guide for more details. If this PR intentionally does not fix a tracked issue, a maintainer can add the no-issue-needed label to silence this reminder.

[zhtmike]

I guess what I am more worried about is whether this PR is breaking existing functionalities and other supported pipelines.

I agree with this. The attention backend lacks test coverage, so it is better to protect it with test cases.

The current fix is for the _flash_3_varlen_hub backend only. It can now correctly handle non-mask, contiguous mask, and non-contiguous mask cases. In my view, it should be safe.

I have also updated the PR content according to the GitHub Action review. Thanks.

[sayakpaul]

Let's introduce a test suite for this and add that to concerned models then.

[zhtmike]

Sorry, I just realized that there is already an AttentionBackendTesterMixin in the latest main branch. We don’t need to introduce a new suite.

I have added an extra TestQwenImageTransformerAttentionBackend to guard against the error. The local test has passed. (The main branch will this test will fail)

I have also tested with TestFluxTransformerAttentionBackend, it also passed.

Thanks!

[sayakpaul]

I ran pytest tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerAttentionBackend and they are passing on 2 H100s.

Can we also verify the image output from a QwenImage pipeline?

[sayakpaul]

I don't understand. Isn't _flash_3_hub what we're fixing in this PR? Or is it like only their varlen variants should be used when using QwenImage?

[zhtmike]

we are fixing _flash_3_varlen_hub backend for mask handling. QwenImage does not support "flash_hub" and "_flash_3_hub", because they are not supporting mask input currently.

[sayakpaul]

Yes makes snese.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zhtmike]

I ran pytest tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerAttentionBackend and they are passing on 2 H100s.

Can we also verify the image output from a QwenImage pipeline?

I ran with pipe.transformer.set_attention_backend("_flash_3_varlen_hub"), the output image is expected.

[sayakpaul]

Possible to provide the full code?

[zhtmike]

Possible to provide the full code?

Sure, we can use this code snippet

importtorchimporttorch.distributedasdistimportargparseimportosfromdiffusersimportQwenImagePipelinefromdiffusersimportContextParallelConfigdefparse_args():
parser=argparse.ArgumentParser(
description="Test Qwen-Image with Context Parallelism")
returnparser.parse_args()
args=parse_args()
ifdist.is_available():
dist.init_process_group(backend="nccl")
rank=dist.get_rank()
device=torch.device("cuda", rank%torch.cuda.device_count())
world_size=dist.get_world_size()
torch.cuda.set_device(device)
else:
rank=0device=torch.device("cuda"iftorch.cuda.is_available() else"cpu")
world_size=1model_id="Qwen/Qwen-Image"pipe=QwenImagePipeline.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
)
pipe.to(device)
pipe.transformer.set_attention_backend("_flash_3_varlen_hub")
ifworld_size>1:
fromdiffusersimportQwenImageTransformer2DModelassertisinstance(pipe.transformer, QwenImageTransformer2DModel)
pipe.transformer.enable_parallelism(config=ContextParallelConfig(
ulysses_degree=world_size))
pipe.set_progress_bar_config(disable=rank!=0)
positive_magic= {
"en": ", Ultra HD, 4K, cinematic composition.", # for english prompt"zh": ", 超清，4K，电影级构图.", # for chinese prompt
}
prompts= [
"A coffee shop entrance features a chalkboard sign reading "'"Qwen Coffee 😊 $2 per cup," with a neon light beside it ''displaying "通义千问". Next to it hangs a poster showing a '"beautiful Chinese woman, and beneath the poster is written "'"π≈3.1415926-53589793-23846264-33832795-02384197". '"Ultra HD, 4K, cinematic composition",
"A cute cat with long hair sitting on a sofa, Ultra HD, 4K, cinematic composition."
]
inputs= {
"prompt": [p+positive_magic["en"] forpinprompts],
"generator": torch.Generator(device="cpu").manual_seed(0),
"true_cfg_scale": 4.0,
"negative_prompt": " ",
"num_inference_steps": 50,
"num_images_per_prompt": 1,
"height": 1024,
"width": 1024,
}
withtorch.inference_mode():
output=pipe(**inputs)
fori, output_imageinenumerate(output.images):
ifworld_size>1:
save_path=f"output_image_ulysses{world_size}_{i}.png"else:
save_path=f"output_image_{i}.png"ifrank==0:
output_image.save(save_path)
print(f"image saved at {save_path}")
ifdist.is_initialized():
dist.destroy_process_group()

with command

torchrun --nproc_per_node=1 --local-ranks-filter=0 test.py > run_single.log 2>&1

to test with _flash_3_varlen_hub backend.

[sayakpaul]

Would it also make sense to make a note about this in the QwenImage pipeline docs (that one must use the varlen version of FA and FA3 backends)?