[None][feat] enable GQA and cross-attention for attn2d by NVShreyas · Pull Request #14961 · NVIDIA/TensorRT-LLM

NVShreyas · 2026-06-04T15:06:54Z

Summary by CodeRabbit

Bug Fixes
- Enhanced runtime validation and consistency checks for attention mechanisms in distributed settings to ensure correct dimension and shape handling.
Tests
- Expanded test coverage for advanced attention configurations on multi-GPU distributed systems.

Description

Test Coverage

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.

coderabbitai · 2026-06-04T15:14:17Z

📝 Walkthrough

Walkthrough

This PR enhances 2D mesh-sharded attention to support Group Query Attention (different Q and KV head counts) and cross-attention (different Q and KV sequence lengths). The implementation adds separate shard sequence tracking, runtime validation, and comprehensive distributed tests validating correctness against full-sequence reference behavior.

Changes

2D Attention: Group Query Attention and Cross-Attention Support

Layer / File(s)	Summary
Attention2DAttention: Documentation and Shape Semantics `tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`	Tensor shape documentation updated to reflect Q sharded as `[B, S_q/P, H_q, D]` and K/V as `[B, S_kv/P, H_kv, D]`, clarifying row/col mesh all-gather mappings and divisibility constraints for independent Q/KV sequence lengths.
Attention2DAttention: GQA Configuration and Forward Logic `tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py`	Added `num_kv_heads` capture from inner backend, rewrote `forward` to use separate `shard_seq_q` and `shard_seq_kv`, added head dimension/count/shape validation, and compute mesh-derived global sequence lengths with divisibility enforcement.
Attention Module: Relaxed Sequence-Parallel Validation `tensorrt_llm/_torch/visual_gen/modules/attention.py`	Narrowed SEPARATE_QKV cross-attention validation to only reject Ring sequence parallelism (`ring_size > 1`); removed prior Attention2D-related sizing constraints to permit GQA-based configurations.
Test Reference: _LSEVanillaAttention GQA Support `tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py`	Extended test reference class with optional `num_kv_heads` parameter, added `_expand_kv_heads` head-replication helper, and integrated `enable_gqa` flag into forward/forward_with_lse attention computation.
Test Logic: GQA and Cross-Attention Scenarios `tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py`	Added `_logic_attn2d_gqa` testing 8 query heads vs. 2 KV heads with equal sequences, and `_logic_attn2d_cross_attention` testing asymmetric Q/KV sequence lengths (8 vs. 4 tokens) with GQA on 2x2 distributed mesh; both validate against full-sequence SDPA reference.
Test Class: Distributed Test Harness `tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py`	Added `TestAttn2DAttentionGQAAndCrossAttention` class with `test_attn2d_gqa` and `test_attn2d_cross_attention` methods wrapping the new logic functions through the distributed test runner.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested labels

VisualGen

Suggested reviewers

yibinl-nvidia
chang-l
brb-nv

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	PR description is incomplete—only the template structure is present with no implementation details, test coverage, or rationale filled in.	Fill in the Description and Test Coverage sections with specifics about GQA/cross-attention changes. Complete the PR checklist items relevant to this PR and confirm API compatibility status.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main feature addition: enabling GQA and cross-attention support for attn2d, which aligns with the core changes across all modified files.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

coderabbitai

🧹 Nitpick comments (1)

tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py (1)
403-491: Coverage is sufficient for this PR scope.

In tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py, coverage is sufficient for the new GQA and cross-attention Attention2D paths (distributed correctness vs SDPA reference).
Optional follow-up outside this PR: add one negative test for invalid num_heads/num_kv_heads ratio handling in _LSEVanillaAttention.

Also applies to: 531-540
🤖 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/multi_gpu/test_attn2d_attention.py` around
lines 403 - 491, Add a negative unit test to validate that _LSEVanillaAttention
rejects invalid num_heads / num_kv_heads ratios: call _LSEVanillaAttention with
mismatched num_kv_heads (e.g., num_heads=8, num_kv_heads=3 or 0) and assert it
raises the expected error (ValueError or a specific exception) rather than
silently proceeding; place the test alongside the existing multi-GPU tests (the
file test_attn2d_attention.py) and reference the constructor
_LSEVanillaAttention to locate the code path that should perform the validation.

🤖 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.

Nitpick comments:
In `@tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py`:
- Around line 403-491: Add a negative unit test to validate that
_LSEVanillaAttention rejects invalid num_heads / num_kv_heads ratios: call
_LSEVanillaAttention with mismatched num_kv_heads (e.g., num_heads=8,
num_kv_heads=3 or 0) and assert it raises the expected error (ValueError or a
specific exception) rather than silently proceeding; place the test alongside
the existing multi-GPU tests (the file test_attn2d_attention.py) and reference
the constructor _LSEVanillaAttention to locate the code path that should perform
the validation.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 662e9eda-102f-4236-988b-271899d66b06

📥 Commits

Reviewing files that changed from the base of the PR and between c17611c and 878508c.

📒 Files selected for processing (3)

tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py
tensorrt_llm/_torch/visual_gen/modules/attention.py
tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py

NVShreyas · 2026-06-05T19:21:35Z

/bot run

tensorrt-cicd · 2026-06-05T19:30:30Z

PR_Github #52412 [ run ] triggered by Bot. Commit: 067bc41 Link to invocation

tensorrt-cicd · 2026-06-05T21:30:58Z

PR_Github #52412 [ run ] completed with state SUCCESS. Commit: 067bc41
/LLM/main/L0_MergeRequest_PR pipeline #41705 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

NVShreyas · 2026-06-06T18:19:45Z

/bot run

tensorrt-cicd · 2026-06-06T18:25:20Z

PR_Github #52521 [ run ] triggered by Bot. Commit: d27f4ca Link to invocation

tensorrt-cicd · 2026-06-06T19:22:05Z

PR_Github #52521 [ run ] completed with state SUCCESS. Commit: d27f4ca
/LLM/main/L0_MergeRequest_PR pipeline #41811 completed with status: 'ABORTED'

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

Link to invocation

NVShreyas · 2026-06-07T16:32:56Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-07T16:39:14Z

PR_Github #52592 [ run ] triggered by Bot. Commit: d27f4ca Link to invocation

tensorrt-cicd · 2026-06-08T01:04:59Z

PR_Github #52592 [ run ] completed with state FAILURE. Commit: d27f4ca
/LLM/main/L0_MergeRequest_PR pipeline #41874 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

NVShreyas · 2026-06-08T14:02:58Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T14:09:34Z

PR_Github #52788 [ run ] triggered by Bot. Commit: c64d151 Link to invocation

tensorrt-cicd · 2026-06-08T22:08:55Z

PR_Github #52788 [ run ] completed with state SUCCESS. Commit: c64d151
/LLM/main/L0_MergeRequest_PR pipeline #42042 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

NVShreyas · 2026-06-08T22:19:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-08T22:24:51Z

PR_Github #52864 [ run ] triggered by Bot. Commit: c64d151 Link to invocation

tensorrt-cicd · 2026-06-08T23:23:17Z

PR_Github #52864 [ run ] completed with state SUCCESS. Commit: c64d151
/LLM/main/L0_MergeRequest_PR pipeline #42116 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

NVShreyas · 2026-06-09T14:17:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-09T14:23:01Z

PR_Github #53077 [ run ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-09T20:42:44Z

PR_Github #53077 [ run ] completed with state FAILURE. Commit: 4b73508
/LLM/main/L0_MergeRequest_PR pipeline #42289 completed with status: 'ABORTED'

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

Link to invocation

NVShreyas · 2026-06-10T14:41:26Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-10T14:48:23Z

PR_Github #53320 [ run ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-10T19:38:04Z

PR_Github #53320 [ run ] completed with state SUCCESS. Commit: 4b73508
/LLM/main/L0_MergeRequest_PR pipeline #42506 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

NVShreyas · 2026-06-11T02:00:59Z

/bot kill

tensorrt-cicd · 2026-06-11T02:09:36Z

PR_Github #53432 [ kill ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-11T02:09:39Z

PR_Github #53432 [ kill ] completed with state SUCCESS. Commit: 4b73508
Successfully killed previous jobs for commit 4b73508

Link to invocation

NVShreyas · 2026-06-11T02:25:41Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T02:32:32Z

PR_Github #53444 [ run ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-11T07:18:23Z

PR_Github #53444 [ run ] completed with state FAILURE. Commit: 4b73508
/LLM/main/L0_MergeRequest_PR pipeline #42611 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

NVShreyas · 2026-06-11T07:33:23Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T07:42:35Z

PR_Github #53519 [ run ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-11T13:33:30Z

PR_Github #53519 [ run ] completed with state SUCCESS. Commit: 4b73508
/LLM/main/L0_MergeRequest_PR pipeline #42677 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

NVShreyas · 2026-06-11T13:33:55Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-11T13:40:03Z

PR_Github #53581 [ run ] triggered by Bot. Commit: 4b73508 Link to invocation

tensorrt-cicd · 2026-06-11T14:54:53Z

PR_Github #53581 [ run ] completed with state SUCCESS. Commit: 4b73508
/LLM/main/L0_MergeRequest_PR pipeline #42728 completed with status: 'SUCCESS'

CI Report

Link to invocation

NVShreyas requested review from a team as code owners June 4, 2026 15:06

NVShreyas requested a review from brb-nv June 4, 2026 15:06

github-actions Bot assigned NVShreyas Jun 4, 2026

coderabbitai Bot reviewed Jun 4, 2026

View reviewed changes

NVShreyas force-pushed the user/shreyasm/attn2d-gqa-cross branch from 878508c to da6c8ab Compare June 5, 2026 15:25

NVShreyas force-pushed the user/shreyasm/attn2d-gqa-cross branch from d27f4ca to c64d151 Compare June 8, 2026 14:02

chang-l approved these changes Jun 8, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py Outdated

zhenhuaw-me approved these changes Jun 9, 2026

View reviewed changes

NVShreyas added 5 commits June 8, 2026 20:16

enable GQA and cross-attention for attn2d

f86852c

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

allow attn2d for cosmos3

454f06b

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

cosmos3 attn2d tests

ad61184

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

remove outdated test

a366b8a

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

only remove redundant checks

4b73508

Signed-off-by: Shreyas Misra <shreyasm@nvidia.com>

NVShreyas force-pushed the user/shreyasm/attn2d-gqa-cross branch from 88a2253 to 4b73508 Compare June 9, 2026 03:23

yuxianq approved these changes Jun 9, 2026

View reviewed changes

NVShreyas merged commit 54dec4f into NVIDIA:main Jun 11, 2026
7 checks passed

chang-l mentioned this pull request Jun 22, 2026

[None][CI] Waive flaky test_vbench_dimension_score_wan (nvbugs/6357628) #15529

Closed

Uh oh!

Conversation

NVShreyas commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

NVShreyas commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

tensorrt-cicd commented Jun 5, 2026

Uh oh!

NVShreyas commented Jun 6, 2026

Uh oh!

tensorrt-cicd commented Jun 6, 2026

Uh oh!

tensorrt-cicd commented Jun 6, 2026

Uh oh!

NVShreyas commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 7, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

NVShreyas commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

NVShreyas commented Jun 8, 2026

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

Uh oh!

tensorrt-cicd commented Jun 8, 2026

Uh oh!

NVShreyas commented Jun 9, 2026

Uh oh!

tensorrt-cicd commented Jun 9, 2026

Uh oh!

tensorrt-cicd commented Jun 9, 2026

Uh oh!

NVShreyas commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

tensorrt-cicd commented Jun 10, 2026

Uh oh!

NVShreyas commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

NVShreyas commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

Uh oh!

NVShreyas commented Jun 11, 2026

Uh oh!

tensorrt-cicd commented Jun 11, 2026

NVShreyas commented Jun 4, 2026 •

edited

Loading

coderabbitai Bot commented Jun 4, 2026 •

edited

Loading