Skip to content

[Enhancement] Attention backend selection: name every fallback, reject invalid configs early #19615

Description

@DhineshPonnarasan

🚀 The feature, motivation and pitch

TensorRT-LLM's PyTorch path selects an attention backend from VANILLA / TRTLLM / FLASHINFER via the attn_backend configuration. The current selection behavior has several gaps that make it hard to answer "which backend am I actually running, and why":

  • An invalid or unknown attn_backend value is not rejected. get_attention_backend() logs a generic warning (Falling back to TRTLLM attention backend) and selects TRTLLM, so a typo can silently change the executed backend — along with its supported features and performance characteristics.
  • A genuinely unavailable backend (e.g. FLASHINFER requested where the package is not installed) takes the same undifferentiated fallback path and is indistinguishable from an invalid name.
  • Some backend-dependent constraints are detected late. The attention-sinks and relative-attention-bias TRTLLM requirements are asserted during forward execution, after the cost of a full build; this class of late failure has already surfaced in user reports ([Bug]: GPT-OSS-20B fails on RTX Pro 6000 when using FlashInfer backend #13156).
  • Some backend-dependent fallback/degradation messages do not clearly identify the backend or reason involved (e.g. rope-fusion handling), and at least one eligibility gate disables with no message.
  • Backend selection is redundantly resolved on the per-layer construction path (Attention.__init__ resolves the backend class, then create_attention resolves it again internally).

Requested behavior, stated as observable behavior rather than implementation:

  • Validate backend names at configuration time against the supported set, with an actionable error listing valid values for invalid input.
  • Make every backend fallback identify the requested backend, the reason it could not be selected, and the selected fallback.
  • Move backend-dependent checks earlier (configuration/construction) where technically possible, while retaining genuinely runtime-dependent checks at runtime with improved diagnostics.
  • Eliminate redundant backend dispatch during attention construction.
  • Add no-GPU regression coverage for backend selection decisions.
  • Preserve existing behavior for all valid configurations, and explicitly document the intentional change that invalid input raises instead of silently selecting TRTLLM.

Alternatives

Several alternatives were considered.

Keeping the current generic warning-only behavior does not address invalid configuration input or late validation — the warning neither names the problem nor prevents a wasted build.

Introducing a large MoE-style or general attention resolution framework would be larger than necessary for the current backend-selection problem, which involves three backend classes and a small set of knowable conditions.

Extending only the existing _resolve_* feature-level checks would not address name validation, availability handling, or selection-time decisions because those checks operate inside an already-selected backend.

Redesigning FMHA/FmhaManager or FlashInfer internal kernel dispatch is also out of scope because inner-kernel selection is a separate layer with different inputs and constraints.

This request is therefore scoped to the existing attention backend selection/configuration boundary, leaving those layers untouched.

Additional context

Verified against upstream main at commit 14729d4a3a (synchronized with origin/main):

  • The fallback warning in tensorrt_llm/_torch/attention/backends/utils.py remains generic.
  • attn_backend in tensorrt_llm/llmapi/llm_args.py remains a string field without supported-value validation.
  • The attention-sinks / relative-attention-bias checks in tensorrt_llm/_torch/attention/attention.py remain forward-time.
  • Redundant backend resolution remains in the construction path.
  • No dedicated selection/fallback test coverage was found; existing tests exercise backends rather than the selection decisions.

Related upstream PRs and issues — including #13156, #14635, #14942, #15417, #17800, #18025, #18520, and the FMHA selection work — were reviewed, and no exact or substantial duplicate of this proposal was found. They are cited here as precedent and context.

Testing plan:

  • Parameterized no-GPU unit-test matrix covering backend × sparse configuration × FlashInfer availability × MLA/chunked and mocked capability conditions.
  • Existing attention and llmapi test suites.
  • Small-model integration runs confirming unchanged valid-configuration behavior.
  • T4 validation of applicable FP16 selection and degradation paths.

I have access to a Tesla T4 (SM75, 16 GB) and will validate the applicable FP16 selection and degradation paths on real hardware. Newer-architecture-specific behavior will rely on mocked capability tests and existing CI hardware.

This proposal does not attempt to restore SM75/T4 kernel support.

Before submitting a new issue...

  • Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

Inference runtime<NV>General operational aspects of TRTLLM execution not in other categories.feature requestNew feature or request. This includes new model, dtype, functionality support

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions