Skip to content

[None][feat] Refactor visual_gen's attention handling: - #17449

Draft
xrq-phys wants to merge 1 commit into
NVIDIA:mainfrom
xrq-phys:ruqingx/fix/attn_meta+full-refactor
Draft

[None][feat] Refactor visual_gen's attention handling:#17449
xrq-phys wants to merge 1 commit into
NVIDIA:mainfrom
xrq-phys:ruqingx/fix/attn_meta+full-refactor

Conversation

@xrq-phys

@xrq-phys xrq-phys commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

  • Metadata as params
  • Wan-like: Self attention and cross attention must use separate metadata sites
  • Enable (quantized) TRTLLM backend for cross attention
    • Separate is_cross definition in thop vs. in python: thop needs specialized is_cross treatment only if KV cache is paged or ragged, while separateQkv cross routes through regular context attention path
  • VisualGen attention module:
    • qkv_format only links to linear to_qkv's output format. Never infer self/cross attention from it.
    • use an is_cross flag to clearly mark self/cross attention instead.
  • Pipeline manages attention metadata.
  • Pipeline registers attention metadata as CUDA graph keys.

Test Coverage

Tests adjusted against new metadata use.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

- Metadata as params
- Wan-like: Self attention and cross attention must use separate metadata sites
- Enable (quantized) TRTLLM backend for cross attention
 * Separate is_cross definition in thop vs. in python: thop needs specialized is_cross treatment only if KV cache is paged or ragged, while separateQkv cross routes through regular context attention path
- VisualGen attention module:
 * qkv_format only links to linear to_qkv's output format. Never infer self/cross attention from it.
 * use an is_cross flag to clearly mark self/cross attention instead.
- Pipeline manages attention metadata.
- Pipeline registers attention metadata as CUDA graph keys.

Signed-off-by: Ruqing Xu <7891482+xrq-phys@users.noreply.github.com>
@xrq-phys xrq-phys changed the title [none][feat] Fully refactor visual_gen's attention handling: [None][feat] Fully refactor visual_gen's attention handling: Aug 9, 2026
@xrq-phys

xrq-phys commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@xrq-phys xrq-phys changed the title [None][feat] Fully refactor visual_gen's attention handling: [None][feat] Refactor visual_gen's attention handling: Aug 9, 2026
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64867 [ run ] triggered by Bot. Commit: e0bea07 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #64867 [ run ] completed with state SUCCESS. Commit: e0bea07
/LLM/main/L0_MergeRequest_PR pipeline #52708 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants