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

tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md

@@ -353,6 +353,21 @@ request checks. For mixed non-MLA batches, the manager checks each active phase
 independently with `is_supported(..., phase=...)`; a phased library accepts only
 phases backed by its corresponding `run_*()` entry point.
 
+`Fmha` owns both entry points. Libraries declare shared capabilities through
+class attributes, such as `supports_skip_correction`, and override only
+`_is_available()` and `_is_supported()` for implementation-specific checks.
+`is_available()` rejects unsupported static capabilities before calling
+`_is_available()`. `is_supported()` provides the same boundary for shared
+request capability checks and delegates all inputs, including `phase`, to
+`_is_supported()`. Both hooks default to `True` when no additional restriction
+is needed. Parent-hook delegation must use `super()._is_available()` or
+`super()._is_supported()` to avoid re-entering the public wrapper.
+
+Availability requirements must be finalized before manager construction and
+remain invariant for its lifetime. Request-varying capability requirements
+must be represented in `FmhaManager._make_cache_key`, because a cache hit
+reuses the selected library without rechecking support.
+
 The FMHA package is split by role:
 
 - `fmha/interface.py` defines the `Fmha` runtime contract.

tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py

@@ -44,16 +44,7 @@ class CuteDslMlaFmha(PhasedFmha):
     """Blackwell CuTe DSL FMHA library for decode-only MLA."""
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
-        if (
-            getattr(attn, "skip_correction_threshold", 0.0) > 0.0
-            and not cls.supports_skip_correction
-        ):
-            logger.debug(
-                "CuTe DSL MLA FMHA is unavailable: skip-correction is enabled and unsupported."
-            )
-            return False
-
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not IS_CUTLASS_DSL_AVAILABLE:
             logger.debug("CuTe DSL MLA FMHA is unavailable: nvidia-cutlass-dsl is not installed.")
             return False
@@ -196,7 +187,7 @@ def _kernel_can_implement(
             )
         return True, ""
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],

tensorrt_llm/_torch/attention/backends/fmha/fallback.py

@@ -59,7 +59,7 @@ class FallbackFmha(Fmha):
     supports_skip_correction = True
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         sparse_algorithm = getattr(attn.sparse_params, "algorithm", None)
         if sparse_algorithm in ("deepseek_v4", "dsa"):
             if getattr(attn, "kv_cache_dtype", None) == "fp8_ds_mla":
@@ -68,7 +68,7 @@ def is_available(cls, attn: "TrtllmAttention") -> bool:
                 return False
         return True
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],

tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py

@@ -55,17 +55,7 @@ def __init__(self, attn: "TrtllmAttention"):
         )
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
-        if (
-            getattr(attn, "skip_correction_threshold", 0.0) > 0.0
-            and not cls.supports_skip_correction
-        ):
-            logger.debug(
-                "FlashInfer sparse-MLA FMHA is unavailable: skip-correction is "
-                "enabled and unsupported."
-            )
-            return False
-
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not attn.is_mla_enable or getattr(attn, "kv_cache_dtype", None) != "fp8_ds_mla":
             return False
         return is_flashinfer_sparse_mla_enabled(getattr(attn.sparse_params, "algorithm", None))

tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py

@@ -392,17 +392,7 @@ def __init__(self, attn: "TrtllmAttention") -> None:
         self._multi_ctas_kv_counter_buffer: Optional[torch.Tensor] = None
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
-        if (
-            getattr(attn, "skip_correction_threshold", 0.0) > 0.0
-            and not cls.supports_skip_correction
-        ):
-            logger.debug(
-                "FlashInfer TRTLLM-Gen FMHA is unavailable: skip-correction is "
-                "enabled and unsupported."
-            )
-            return False
-
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         if not IS_FLASHINFER_AVAILABLE:
             logger.debug("FlashInfer TRTLLM-Gen FMHA is unavailable: flashinfer is not installed.")
             return False
@@ -537,7 +527,7 @@ def _check_mla_generation_support(
 
         return True, ""
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],

tensorrt_llm/_torch/attention/backends/fmha/interface.py

@@ -16,11 +16,12 @@
 import weakref
 from abc import ABC, abstractmethod
 from enum import Enum
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, ClassVar, Optional, final
 
 import torch
 
 from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs
+from tensorrt_llm.logger import logger
 
 if TYPE_CHECKING:
     from tensorrt_llm._torch.attention.backends.trtllm import (
@@ -39,7 +40,7 @@ class FmhaPhase(str, Enum):
 class Fmha(ABC):
     """Common runtime contract for TRT-LLM attention FMHA libraries."""
 
-    supports_skip_correction = False
+    supports_skip_correction: ClassVar[bool] = False
 
     def __init__(self, attn: "TrtllmAttention"):
         self._attn_ref: weakref.ReferenceType["TrtllmAttention"] = weakref.ref(attn)
@@ -52,9 +53,14 @@ def attn(self) -> "TrtllmAttention":
         return attn
 
     @classmethod
+    @final
     def is_available(cls, attn: "TrtllmAttention") -> bool:
         """Return whether this library can serve the given attention layer.
 
+        Check shared capabilities before the implementation's
+        ``_is_available`` hook. Libraries declare their capabilities as class
+        attributes and override only the hook for additional static checks.
+
         Evaluated once per ``FmhaManager`` construction, currently at the end
         of ``TrtllmAttention.update_quant_config()``. Conditions must depend
         only on state finalized before manager construction and invariant for
@@ -63,8 +69,23 @@ def is_available(cls, attn: "TrtllmAttention") -> bool:
         it is not revalidated. Request-varying conditions belong in
         ``is_supported`` instead.
         """
+        if attn.skip_correction_threshold > 0.0 and not cls.supports_skip_correction:
+            logger.debug(
+                f"{cls.__name__} is unavailable: skip-correction is enabled and unsupported."
+            )
+            return False
+        return cls._is_available(attn)
+
+    @classmethod
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
+        """Check implementation-specific static restrictions after capability checks.
+
+        Delegate to ``super()._is_available(attn)`` to reuse a parent hook;
+        calling ``is_available`` here would re-enter the shared wrapper.
+        """
         return True
 
+    @final
     def is_supported(
         self,
         q: torch.Tensor,
@@ -77,12 +98,33 @@ def is_supported(
     ) -> bool:
         """Return whether this library supports the request or requested phase.
 
+        Shared request capability checks belong here, before delegating to
+        ``_is_supported``. Libraries override only that hook for their
+        request-specific restrictions.
+
         Forward-varying selection conditions must be represented in
         ``FmhaManager._make_cache_key``. Conditions omitted from that key must
         remain invariant for the attention instance. Size-based conditions
         must also preserve the same result throughout each FMHA cache grid
         cell or add the relevant boundary to the grid's candidate list.
         """
+        return self._is_supported(q, k, v, metadata, forward_args, phase=phase)
+
+    def _is_supported(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        v: Optional[torch.Tensor],
+        metadata: "TrtllmAttentionMetadata",
+        forward_args: AttentionForwardArgs,
+        *,
+        phase: Optional[FmhaPhase] = None,
+    ) -> bool:
+        """Check implementation-specific request restrictions after capability checks.
+
+        Delegate to ``super()._is_supported(...)`` to reuse a parent hook;
+        calling ``is_supported`` here would re-enter the shared wrapper.
+        """
         return True
 
     @abstractmethod

tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py

@@ -187,14 +187,7 @@ class MsaSparseGqaFmha(Fmha):
     """
 
     @classmethod
-    def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool:
-        if (
-            attn is not None
-            and getattr(attn, "skip_correction_threshold", 0.0) > 0.0
-            and not cls.supports_skip_correction
-        ):
-            return False
-
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         # fmha_sm100 runs only on the SM100 family and is packaged in the
         # wheel, so it is unavailable off SM100 or without the wheel.
         # Imported lazily because the minimax_m3 package init imports the trtllm

tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py

@@ -88,7 +88,7 @@ def __init__(self, attn: "TrtllmAttention") -> None:
         self._decode_workspace_required_bytes = 0
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         sm = get_sm_version()
         if sm not in (100, 103):
             logger.debug(f"PrimTS FMHA is unavailable: requires SM100 or SM103, got SM{sm}.")
@@ -145,7 +145,7 @@ def _missing_fused_nanobind_ops() -> list[str]:
         )
         return [name for name in required_ops if not hasattr(thop, name)]
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],

tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py

@@ -55,7 +55,7 @@ def __init__(self, attn: "TrtllmAttention"):
         self._multi_processor_count: Optional[int] = None
 
     @classmethod
-    def is_available(cls, attn: "TrtllmAttention") -> bool:
+    def _is_available(cls, attn: "TrtllmAttention") -> bool:
         required_ops = (
             "get_trtllm_gen_context_workspace_layout",
             "trtllm_gen_context_preprocess",
@@ -73,7 +73,7 @@ def is_available(cls, attn: "TrtllmAttention") -> bool:
             return False
         return True
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],

tests/unittest/_torch/attention/fmha_test_utils.py

@@ -32,6 +32,7 @@ def __init__(self, local_layer_idx: int = 0) -> None:
         self.num_kv_heads = 1
         self.predicted_tokens_per_seq = 1
         self.has_fp8_kv_cache = False
+        self.skip_correction_threshold = 0.0
         self.local_layer_idx = local_layer_idx
 
 
@@ -52,7 +53,7 @@ def __init__(
         self._workspace_size = workspace_size
         self._support_predicate = support_predicate
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: torch.Tensor | None,
@@ -120,7 +121,7 @@ def __init__(
         self._support_predicate = support_predicate
         self._request_support_predicate = request_support_predicate
 
-    def is_supported(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: torch.Tensor | None,

tests/unittest/_torch/attention/test_fmha_interface.py

@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import cast
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+from fmha_test_utils import FakeAttention
+
+from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha
+from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha, FmhaPhase
+from tensorrt_llm._torch.attention.backends.fmha.registry import FMHA_LIBS
+from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs
+from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata
+
+
+class _MinimalFmha(Fmha):
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor | None,
+        v: torch.Tensor | None,
+        metadata: TrtllmAttentionMetadata,
+        forward_args: AttentionForwardArgs,
+    ) -> None:
+        pass
+
+
+@pytest.mark.parametrize("fmha_cls", FMHA_LIBS.values(), ids=FMHA_LIBS.keys())
+@pytest.mark.parametrize("threshold", [0.0, 0.1])
+@pytest.mark.parametrize("implementation_available", [False, True])
+def test_availability_checks_capabilities_before_implementation(
+    fmha_cls: type[Fmha],
+    threshold: float,
+    implementation_available: bool,
+) -> None:
+    attn = cast(TrtllmAttention, FakeAttention())
+    attn.skip_correction_threshold = threshold
+    capability_supported = threshold == 0.0 or fmha_cls is FallbackFmha
+
+    with patch.object(fmha_cls, "_is_available", return_value=implementation_available) as hook:
+        assert fmha_cls.is_available(attn) is (capability_supported and implementation_available)
+        if capability_supported:
+            hook.assert_called_once_with(attn)
+        else:
+            hook.assert_not_called()
+
+
+@pytest.mark.parametrize("phase", [None, FmhaPhase.CONTEXT, FmhaPhase.GENERATION])
+@pytest.mark.parametrize("supported", [False, True])
+def test_support_forwards_request_and_phase(phase: FmhaPhase | None, supported: bool) -> None:
+    attn = cast(TrtllmAttention, FakeAttention())
+    fmha = _MinimalFmha(attn)
+    q, k, v = (torch.empty((2, 4)) for _ in range(3))
+    metadata = Mock(spec=TrtllmAttentionMetadata)
+    forward_args = AttentionForwardArgs()
+
+    with patch.object(fmha, "_is_supported", return_value=supported) as hook:
+        assert fmha.is_supported(q, k, v, metadata, forward_args, phase=phase) is supported
+        hook.assert_called_once_with(q, k, v, metadata, forward_args, phase=phase)
+
+
+def test_default_hooks_accept_requests() -> None:
+    attn = cast(TrtllmAttention, FakeAttention())
+    fmha = _MinimalFmha(attn)
+
+    assert _MinimalFmha.is_available(attn)
+    assert fmha.is_supported(
+        torch.empty((2, 4)), None, None, Mock(spec=TrtllmAttentionMetadata), AttentionForwardArgs()
+    )
+
+
+def test_inherited_availability_hook_uses_subclass_capabilities() -> None:
+    class _SkipCorrectionFmha(_MinimalFmha):
+        supports_skip_correction = True
+
+        @classmethod
+        def _is_available(cls, attn: TrtllmAttention) -> bool:
+            return super()._is_available(attn)
+
+    attn = cast(TrtllmAttention, FakeAttention())
+    attn.skip_correction_threshold = 0.1
+
+    assert not _MinimalFmha.is_available(attn)
+    assert _SkipCorrectionFmha.is_available(attn)

tests/unittest/_torch/attention/test_fmha_manager.py

@@ -868,6 +868,7 @@ def __init__(self, attn: TrtllmAttention) -> None:
 
     attn = TrtllmAttention.__new__(TrtllmAttention)
     attn.is_mla_enable = False
+    attn.skip_correction_threshold = 0.0
     metadata = _make_metadata(num_contexts=0, num_generations=1)
     forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.generation_only)
     q = torch.empty((1, 4))