diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index e56bbb5fa24a..71e17050f946 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/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. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 9584d131e245..2400cafd6e0f 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/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], diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index f03ae9a01dc1..9e85674590f0 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/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], diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py index 2924cc440fc0..03b14aea057e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py +++ b/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)) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 5c30fea5f7b4..d4648893dd68 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/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], diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index 1e5bce7161aa..0691517cbae3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/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 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py index ee85aaf02e0d..63a5388a0004 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py +++ b/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 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index afad6479067d..f89df659d774 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/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], diff --git a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py index 4391ef94bb22..86e377ffbb48 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py +++ b/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], diff --git a/tests/unittest/_torch/attention/fmha_test_utils.py b/tests/unittest/_torch/attention/fmha_test_utils.py index f109d523a170..a20a99dbf156 100644 --- a/tests/unittest/_torch/attention/fmha_test_utils.py +++ b/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, diff --git a/tests/unittest/_torch/attention/test_fmha_interface.py b/tests/unittest/_torch/attention/test_fmha_interface.py new file mode 100644 index 000000000000..a103251a442c --- /dev/null +++ b/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) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 17dd1beda5ce..5827e15bee17 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/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)) diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 87dbd7364f4b..a9eb5e61ebd6 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -86,6 +86,7 @@ def __init__( self.v_head_dim = 128 if is_mla else None self.predicted_tokens_per_seq = 1 self.sparse_params = None + self.skip_correction_threshold = 0.0 self.position_embedding_type = 0 self.quant_mode = 0 self.q_scaling = 1.0