[https://nvbugs/6336801][fix] Fix attention op sync test

[yuxianq]

Summary by CodeRabbit

• Bug Fixes
  • Improved attention call handling so additional supported options are correctly recognized and validated.
  • Fixed test coverage for attention synchronization paths by re-enabling a previously skipped test case.
  • Added support for more attention-related settings in the validation flow, reducing false failures for valid configurations.

Description

Fix the attention op synchronization tests tracked by https://nvbugs/6336801.

• Remove the in-file NVBug waiver so both synchronization tests run again.
• Update the THOP/source alias allowlist for the skip-softmax threshold scale factors and remove the stale beam-width alias.
• Exclude skip_mla_rope_generation from the THOP call-site check because TrtllmAttention.forward consumes this test-only MLA option before the fallback call.

The current waives.txt has no entry for NVBug 6336801; the remaining waiver was the in-file skip marker removed by this change.

Test Coverage

• Targeted NVBug tests: 2 passed
• tests/unittest/_torch/attention/test_attention_op_sync.py: 8 passed
• Pre-commit hooks: passed

PR Checklist

• Please check this box once you have read and are following the contribution guidelines.

GitHub Bot Help

The requested CI command will be posted after the PR is created.

[yuxianq]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This change extends allowlists used to verify thop.attention kwarg mapping by adding an excluded field (skip_mla_rope_generation) and two aliased fields (skip_softmax_threshold_scale_factor_decode/prefill). It also removes a skip marker gating attention op sync tests, enabling them to run.

Changes

Attention Kwarg Allowlist and Test Re-enablement

Layer / File(s)	Summary
Excluded/aliased kwarg field updates tensorrt_llm/_torch/attention_backend/fmha/fallback.py, tests/unittest/_torch/attention/test_attention_op_sync.py	Adds skip_mla_rope_generation to _THOP_EXCLUDED_FIELDS and adds skip_softmax_threshold_scale_factor_decode/_prefill to _THOP_KWARG_SOURCE_ALIASES.
Skip marker removal for sync tests tests/unittest/_torch/attention/test_attention_op_sync.py	Removes the pytest import, the _THOP_SYNC_NVBUG skip constant, and its decorator usage so the affected tests execute instead of being skipped.

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

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and accurately identifies the NVBugs 6336801 attention-op sync fix.
Description check	✅ Passed	The description includes the required sections and clearly summarizes the issue, fix, and tests.
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]

🧹 Nitpick comments (1)

tests/unittest/_torch/attention/test_attention_op_sync.py (1)

87-88: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Skip marker removal restores intended test coverage.

Removing _THOP_SYNC_NVBUG and its decorator re-enables test_attr_kwarg_names_match_source_leaf_attrs_except_allowlisted_aliases and test_every_forward_args_field_is_consumed, restoring coverage for the thop kwarg/source alias contract and the forward-args consumption contract. This aligns with the PR's stated goal of fixing NVBug 6336801 without a corresponding waives.txt entry.

As per path instructions, this restores sufficient coverage for the affected kwarg-mapping and field-consumption checks; no additional test files appear necessary for this cohort.

Also applies to: 465-466, 566-576

🤖 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/attention/test_attention_op_sync.py` around lines 87 -
88, Remove the `_THOP_SYNC_NVBUG` skip marker and its decorator from the
affected attention op sync tests so
`test_attr_kwarg_names_match_source_leaf_attrs_except_allowlisted_aliases` and
`test_every_forward_args_field_is_consumed` run again; use the existing symbols
in `test_attention_op_sync.py` around `_HEADER` and the skipped test definitions
to locate the block, and keep the test coverage aligned with the thop
kwarg/source alias and forward-args consumption contracts.

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.

Nitpick comments:
In `@tests/unittest/_torch/attention/test_attention_op_sync.py`:
- Around line 87-88: Remove the `_THOP_SYNC_NVBUG` skip marker and its decorator
from the affected attention op sync tests so
`test_attr_kwarg_names_match_source_leaf_attrs_except_allowlisted_aliases` and
`test_every_forward_args_field_is_consumed` run again; use the existing symbols
in `test_attention_op_sync.py` around `_HEADER` and the skipped test definitions
to locate the block, and keep the test coverage aligned with the thop
kwarg/source alias and forward-args consumption contracts.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ea542267-23c7-49f7-bfd6-09fdb4383dd1

📥 Commits

Reviewing files that changed from the base of the PR and between 6c22551 and a4d5de8.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/attention_backend/fmha/fallback.py
• tests/unittest/_torch/attention/test_attention_op_sync.py

[tensorrt-cicd]

PR_Github #57677 [ run ] triggered by Bot. Commit: a4d5de8 Link to invocation

[tensorrt-cicd]

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

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57718 [ run ] triggered by Bot. Commit: 1d5d00c Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57735 [ run ] triggered by Bot. Commit: 7300577 Link to invocation

[tensorrt-cicd]

PR_Github #57718 [ run ] completed with state ABORTED. Commit: 1d5d00c

Link to invocation

[tensorrt-cicd]

PR_Github #57735 [ run ] completed with state SUCCESS. Commit: 7300577
/LLM/main/L0_MergeRequest_PR pipeline #46444 completed with status: 'SUCCESS'

CI Report

Link to invocation

[pengbowang-nv]

LGTM