[None][perf] Skip redundant Qwen Image attention masks

[pst2154]

Description\nSkip Qwen-Image joint attention mask construction when the text mask is all valid. This keeps all-valid prompt batches on the normal VisualGen attention backend path instead of falling into the masked SDPA fallback.\n\n## Changes\n- Return None from Qwen prompt encoding when all prompt tokens are valid.\n- Treat all-true encoder_hidden_states_mask as unmasked in the transformer direct-call path.\n- Add a unit test covering all-true mask output equivalence with the unmasked path.\n\n## Testing\n- git diff --check\n- python3 -m py_compile tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py\n- Targeted pytest blocked locally: ModuleNotFoundError: No module named 'mpi4py' from tests/unittest/conftest.py.

Summary by CodeRabbit

• Bug Fixes

  • Improved image-generation handling so fully valid prompt masks are treated as unmasked, reducing unnecessary masking during inference.
  • Updated attention behavior to skip mask construction when it isn’t needed, helping results stay consistent with unmasked runs.

• Tests

  • Added coverage confirming that an all-true text mask produces the same output as no mask.

[coderabbitai]

📝 Walkthrough

Walkthrough

Changes make the prompt embeddings mask optional, returning None when all tokens are valid instead of a full mask tensor. The transformer forward pass conditionally skips building/concatenating the block attention mask when the text mask is entirely True. A new unit test verifies output equivalence between None and all-True masks.

Changes

Optional Attention Mask Handling

Layer / File(s)	Summary
Prompt encoding returns optional mask tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py	_encode_prompt return type updated to Optional[torch.Tensor]; returns (prompt_embeds, None) when all mask entries are valid instead of always returning the full mask.
Conditional transformer attention mask construction tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py	forward only builds and concatenates block_attention_mask when encoder_hidden_states_mask is not entirely True; otherwise leaves it None.
Test coverage for all-True mask equivalence tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py	Adds imports for torch and QwenImageTransformer2DModel, and a new test asserting outputs match when passing None vs. an all-True mask.

Estimated code review effort: 2 (Simple) | ~10 minutes

Sequence Diagram(s)

sequenceDiagram
    participant Pipeline as QwenImagePipeline
    participant Encode as _encode_prompt
    participant Transformer as QwenImageTransformer2DModel

    Pipeline->>Encode: request prompt embeddings
    Encode->>Encode: compute prompt_embeds_mask
    alt all tokens valid
        Encode-->>Pipeline: (prompt_embeds, None)
    else some tokens masked
        Encode-->>Pipeline: (prompt_embeds, mask tensor)
    end
    Pipeline->>Transformer: forward(encoder_hidden_states_mask)
    alt mask is None or all True
        Transformer->>Transformer: skip block_attention_mask concat
    else mask has False entries
        Transformer->>Transformer: build and concat block_attention_mask
    end
    Transformer-->>Pipeline: output tensor

Loading

Suggested reviewers: rakib-hasan, aswinvisva, yechank-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and matches the PR’s performance-focused change to skip redundant Qwen Image attention masks.
Description check	✅ Passed	It covers the issue, changes, and testing, with only minor template-section differences.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py (1)

1044-1055: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift

Avoid data-dependent control flow on the mask inside the graph-captured forward. encoder_hidden_states_mask.all() reads a CUDA tensor value every step and can bake the wrong branch into a CUDA graph for same-shaped masks with different padding patterns. Hoist the all-ones decision out of the captured region or key the graph on a static flag instead.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py`
around lines 1044 - 1055, The mask handling in the Qwen image forward path uses
data-dependent branching on encoder_hidden_states_mask via the all() check,
which is unsafe in graph-captured execution. Update the logic in
transformer_qwen_image.py around the block_attention_mask construction to avoid
reading the mask value inside the captured forward; instead, move the all-ones
decision outside the captured region or derive it from a static flag/shape
metadata so the branch is stable for CUDA graph replay.

🧹 Nitpick comments (1)

tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py (1)

235-271: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add pipeline-level coverage for _encode_prompt's None-return behavior.

This test only validates transformer-level equivalence between None and an all-True mask. The other half of the PR — _encode_prompt returning None when all prompt tokens are valid (and a real mask when they aren't) — has no direct unit test here. Consider adding a test in this file (or a nearby pipeline test module) that stubs self.tokenizer/self.text_encoder to produce controlled attention_mask values and asserts _encode_prompt returns None vs. a tensor accordingly. As per path instructions, please assess whether this gap needs to be closed in this PR or tracked as a follow-up.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py` around
lines 235 - 271, Add direct pipeline-level coverage for
QwenImagePipeline._encode_prompt to verify its attention-mask None-return
behavior, since the current test only checks QwenImageTransformer2DModel
equivalence for all-True masks. Create a focused test in this file or a nearby
pipeline test that stubs self.tokenizer and self.text_encoder to control the
produced attention_mask, then assert _encode_prompt returns None when all prompt
tokens are valid and returns a tensor when any token is masked. Use the
_encode_prompt method and the pipeline’s tokenizer/text_encoder members as the
main symbols to locate and validate the behavior.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py`:
- Around line 1044-1055: The mask handling in the Qwen image forward path uses
data-dependent branching on encoder_hidden_states_mask via the all() check,
which is unsafe in graph-captured execution. Update the logic in
transformer_qwen_image.py around the block_attention_mask construction to avoid
reading the mask value inside the captured forward; instead, move the all-ones
decision outside the captured region or derive it from a static flag/shape
metadata so the branch is stable for CUDA graph replay.

---

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py`:
- Around line 235-271: Add direct pipeline-level coverage for
QwenImagePipeline._encode_prompt to verify its attention-mask None-return
behavior, since the current test only checks QwenImageTransformer2DModel
equivalence for all-True masks. Create a focused test in this file or a nearby
pipeline test that stubs self.tokenizer and self.text_encoder to control the
produced attention_mask, then assert _encode_prompt returns None when all prompt
tokens are valid and returns a tensor when any token is masked. Use the
_encode_prompt method and the pipeline’s tokenizer/text_encoder members as the
main symbols to locate and validate the behavior.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7713dbe1-d063-4f82-bb6a-b05c9541b856

📥 Commits

Reviewing files that changed from the base of the PR and between 5042300 and b059835.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image.py
• tensorrt_llm/_torch/visual_gen/models/qwen_image/transformer_qwen_image.py
• tests/unittest/_torch/visual_gen/test_qwen_image_pipeline_config.py

[pst2154]

Closing this in favor of the combined clean Qwen VisualGen perf PR: #16142. That PR includes this mask fastpath work plus the related QK-norm/RoPE, adaLN, and NVFP4 GELU fixes in one reviewable branch.