[None][refactor] Centralize FMHA availability and support capability checks

[yuxianq]

Dev Engineer Review

• Centralizes skip_correction checks in Fmha.is_available().
• Splits public is_supported() from the _is_supported() implementation hook.
• Updates FMHA backends to use _is_available() and _is_supported().
• Adds supports_skip_correction with a default of False.
• Documents hook delegation, capability lifetimes, and cache-key requirements.
• Main regression risk: external implementations that override the former public hooks must migrate.
• Reported validation passed: 165 B200 tests and all commit hooks.

QA Engineer Review

• Modified FMHA unit tests and fixtures.
• Added coverage for capability rejection, availability ordering, request and phase forwarding, inherited defaults, and fallback behavior.
• Removed the FMHA manager capability-construction test.
• No changed tests require integration CI or manual-QA test-list entries.
• Coverage verdict: sufficient.

Per-File QA Perspective

• tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md: Documents hook behavior and cache requirements. Verify the guidance matches runtime behavior.
• tensorrt_llm/_torch/attention/backends/fmha/interface.py: Adds centralized capability gating and support delegation. Verify rejection ordering and argument forwarding.
• tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py: Renames FMHA hooks. Verify selection behavior.
• tensorrt_llm/_torch/attention/backends/fmha/fallback.py: Renames fallback hooks. Verify fallback behavior.
• tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py: Removes local skip-correction gating. Verify base-class rejection and remaining MLA checks.
• tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py: Renames hooks and removes local gating. Verify centralized checks and FlashInfer behavior.
• tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py: Changes _is_available() to require attention context. Verify correct context propagation.
• tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py: Renames FMHA hooks. Verify existing selection constraints.
• tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py: Renames FMHA hooks. Verify custom-mask selection and support checks.
• tests/unittest/_torch/attention/fmha_test_utils.py: Updates fixtures and fake hooks. Covers interface compatibility; no integration test-list entry applies.
• tests/unittest/_torch/attention/test_fmha_interface.py: Covers capability gating, forwarding, defaults, and inheritance; no integration test-list entry applies.
• tests/unittest/_torch/attention/test_fmha_manager.py: Removes manager capability-construction coverage; no integration test-list entry applies.
• tests/unittest/_torch/attention/test_prims_ts_fmha.py: Updates skip-correction fixture state; no integration test-list entry applies.

Description

FMHA capability checks were repeated in individual implementations, making it easy for a library to miss a newly introduced requirement. Centralize skip-correction gating in Fmha.is_available() before any implementation-specific probe, and migrate every registered library to the _is_available() hook. Libraries opt in through declarations such as supports_skip_correction = True.

Also split is_supported() into a final base-class entry point and an _is_supported() implementation hook, preserving all request arguments and the phase keyword. This provides the shared entry point for future request capabilities, including the block-sparse check in #18815. The block-sparse gate itself belongs to that integration. Document capability lifetimes, cache-key requirements, and parent-hook delegation.

Test Coverage

Before the test-only deletion in 67732df4ca, 165 tests passed on B200 with existing native build artifacts across the following suites:

• tests/unittest/_torch/attention/test_fmha_interface.py
• tests/unittest/_torch/attention/test_fmha_registry.py
• tests/unittest/_torch/attention/test_fmha_manager.py
• tests/unittest/_torch/attention/test_combined_fmha.py
• tests/unittest/_torch/attention/test_prims_ts_fmha.py
• tests/unittest/_torch/attention/test_attention_mla.py::test_attention_mla_skip_correction (BF16 and FP8 KV cache)

The first five suites passed 163 cases. The two MLA cases initially failed to locate CUDA JIT headers; both passed after adding the test virtualenv's pip to PATH so the JIT could resolve the editable installation. No source workaround was needed.

The manager construction capability test was subsequently removed, leaving 164 previously validated cases. Runtime tests were not rerun for this test-only deletion. All commit hooks passed for the deletion, including lint/formatting, vendored-source verification, test-list validation, and DCO.

Coverage includes capability rejection before probing, availability result propagation for all registered libraries, request/phase forwarding, inherited defaults, and existing selection/cache/composition behavior.

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 453375ea-e84e-423d-bed7-6f892b8c5e8f

📥 Commits

Reviewing files that changed from the base of the PR and between 022a3fb and 67732df.

📒 Files selected for processing (1)

• tests/unittest/_torch/attention/test_fmha_manager.py

💤 Files with no reviewable changes (1)

• tests/unittest/_torch/attention/test_fmha_manager.py

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

Walkthrough

The FMHA interface now centralizes capability checks and delegates backend-specific availability and request support checks to private hooks. Backends, tests, fixtures, and the developer guide use the updated contract.

Changes

FMHA hook contract

Layer / File(s)	Summary
Centralized FMHA validation tensorrt_llm/_torch/attention/backends/fmha/interface.py	The public wrappers validate shared capabilities before calling _is_available() and _is_supported().
Backend hook migration tensorrt_llm/_torch/attention/backends/fmha/*	FMHA backends rename availability and support hooks and remove duplicated skip-correction rejection checks.
Contract validation and documentation tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md, tests/unittest/_torch/attention/*	Tests cover capability gating, argument forwarding, default hooks, and updated test fixtures. The developer guide documents the updated contract.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant FMHAManager
  participant Fmha
  participant BackendHook
  FMHAManager->>Fmha: Check availability
  Fmha->>Fmha: Validate shared capabilities
  Fmha->>BackendHook: Call _is_available()
  BackendHook-->>Fmha: Return availability
  FMHAManager->>Fmha: Check request support
  Fmha->>BackendHook: Call _is_supported()
  BackendHook-->>Fmha: Return support result

Loading

Merge Risk: ⚪ Minimal · up to 67732

FMHA capability validation is centralized behind shared public hooks with backend-specific private implementations. No current merge-blocking risk is identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.16% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 38 functions across 12 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Description check	✅ Passed	The description explains the centralization of FMHA capability checks, the hook changes, documentation updates, and test coverage. It includes the required Description, Test Coverage, and PR Checklist…
Title check	✅ Passed	The title follows the required ticket and type format and clearly summarizes the primary refactor: centralizing FMHA availability and support capability checks.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #72685 [ run ] triggered by Bot. Commit: 67732df Link to invocation

[tensorrt-cicd]

PR_Github #72685 [ run ] completed with state SUCCESS. Commit: 67732df
/LLM/main/L0_MergeRequest_PR pipeline #59678 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

[yihwang-nv]

Thanks, LGTM!

Do we want to strictly prohibit derived FMHA classes from overriding is_available and is_support? If yes, should we add a test to ensure those methods are absent from each registered FMHA lib in a separate PR?

[github-actions]

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

[yihwang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #72746 [ run ] triggered by Bot. Commit: 67732df Link to invocation

[tensorrt-cicd]

PR_Github #72746 [ run ] completed with state SUCCESS. Commit: 67732df
/LLM/main/L0_MergeRequest_PR pipeline #59734 completed with status: 'SUCCESS'

CI Report

Link to invocation