[TRTLLM-12807][feat] Add multiple FMHA library support to TRTLLM attention backend - #15204
Conversation
|
/bot run --disable-fail-fast --add-multi-gpu-test |
yuxianq
changed the title
📝 WalkthroughWalkthroughThis PR introduces a pluggable FMHA (Fused Multi-Head Attention) backend system to TensorRT-LLM's attention implementation. It replaces the hardcoded TRTLLM-Gen and thop.attention dispatch paths with an ordered registry of backend implementations selected via the ChangesFMHA Backend Architecture & Dispatch
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py (1)
689-695:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winDead code:
sink_token_lengthcheck is always False.
sink_token_lengthis hardcoded to0, making the subsequent checkif sink_token_length != 0unreachable. Either remove this dead code or replace with actual StreamingLLM detection.🧹 Proposed fix (remove dead code)
- sink_token_length = 0 - if sink_token_length != 0: - return ( - False, - f"[Generation] StreamingLLM " - f"(sink_token_length={sink_token_length}) is not supported.", - )🤖 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 `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py` around lines 689 - 695, The check for streaming support contains dead code: the local variable sink_token_length is hardcoded to 0 so the subsequent if (sink_token_length != 0) is unreachable; either remove the sink_token_length declaration and the entire conditional branch, or replace it with a real StreamingLLM detection (e.g., inspect the LLM instance or a passed-in parameter such as sink_token_length/streaming flag) so the code actually detects streaming LLMs; update references around the sink_token_length conditional in flashinfer_trtllm_gen.py accordingly and ensure any returned message still uses the same wording if you keep the branch.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/attention_backend/fmha/registry.py (1)
72-73: 💤 Low valueConsider caching the parsed result for hot-path efficiency.
get_enabled_fmha_lib_classes()re-parsesTLLM_FMHA_LIBSon every call. If the registry dispatch iterates through backends per-forward, this could add overhead. Consider caching the result (e.g., via@lru_cache(maxsize=1)or a module-level singleton) if the env var is not expected to change at runtime.🤖 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 `@tensorrt_llm/_torch/attention_backend/fmha/registry.py` around lines 72 - 73, get_enabled_fmha_lib_classes currently calls _parse_fmha_libs_env on every invocation which can be expensive on hot paths; modify it to cache the parsed result (e.g., decorate get_enabled_fmha_lib_classes with functools.lru_cache(maxsize=1) or compute once into a module-level variable) so subsequent calls return the cached list of FMHA classes from FMHA_LIBS without re-parsing the TLLM_FMHA_LIBS env; ensure the cached value is a list of FmhaCls as returned today and reference get_enabled_fmha_lib_classes, _parse_fmha_libs_env and FMHA_LIBS when locating the code to change.tensorrt_llm/_torch/attention_backend/trtllm.py (1)
1591-1593: 💤 Low valueOptional: Remove extraneous space in f-string format spec.
The format spec
{skipped_blocks / total_blocks * 100: .2f}includes a space before.2f, which renders as" 12.34%"with a leading space. Consider:.2ffor cleaner output unless the space is intentional for alignment.♻️ Proposed formatting tweak
- print(f"SKIP_SOFTMAX_STAT: layer{self.layer_idx}: " - f"{skipped_blocks} / {total_blocks}" - f" = {skipped_blocks / total_blocks * 100: .2f}%") + print(f"SKIP_SOFTMAX_STAT: layer{self.layer_idx}: " + f"{skipped_blocks} / {total_blocks}" + f" = {skipped_blocks / total_blocks * 100:.2f}%")🤖 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 `@tensorrt_llm/_torch/attention_backend/trtllm.py` around lines 1591 - 1593, The f-string used in the logging print currently uses a space in the format specifier "{skipped_blocks / total_blocks * 100: .2f}" which yields a leading space; update that f-string to use "{skipped_blocks / total_blocks * 100:.2f}" (referencing the same print that includes self.layer_idx, skipped_blocks and total_blocks) so the percentage renders without the extra leading space.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py`:
- Around line 689-695: The check for streaming support contains dead code: the
local variable sink_token_length is hardcoded to 0 so the subsequent if
(sink_token_length != 0) is unreachable; either remove the sink_token_length
declaration and the entire conditional branch, or replace it with a real
StreamingLLM detection (e.g., inspect the LLM instance or a passed-in parameter
such as sink_token_length/streaming flag) so the code actually detects streaming
LLMs; update references around the sink_token_length conditional in
flashinfer_trtllm_gen.py accordingly and ensure any returned message still uses
the same wording if you keep the branch.
---
Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/fmha/registry.py`:
- Around line 72-73: get_enabled_fmha_lib_classes currently calls
_parse_fmha_libs_env on every invocation which can be expensive on hot paths;
modify it to cache the parsed result (e.g., decorate
get_enabled_fmha_lib_classes with functools.lru_cache(maxsize=1) or compute once
into a module-level variable) so subsequent calls return the cached list of FMHA
classes from FMHA_LIBS without re-parsing the TLLM_FMHA_LIBS env; ensure the
cached value is a list of FmhaCls as returned today and reference
get_enabled_fmha_lib_classes, _parse_fmha_libs_env and FMHA_LIBS when locating
the code to change.
In `@tensorrt_llm/_torch/attention_backend/trtllm.py`:
- Around line 1591-1593: The f-string used in the logging print currently uses a
space in the format specifier "{skipped_blocks / total_blocks * 100: .2f}" which
yields a leading space; update that f-string to use "{skipped_blocks /
total_blocks * 100:.2f}" (referencing the same print that includes
self.layer_idx, skipped_blocks and total_blocks) so the percentage renders
without the extra leading space.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0b24c50a-926b-48db-b498-d0d3fa0c8dd1
📒 Files selected for processing (9)
tensorrt_llm/_torch/attention_backend/fmha/__init__.pytensorrt_llm/_torch/attention_backend/fmha/fallback.pytensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.pytensorrt_llm/_torch/attention_backend/fmha/interface.pytensorrt_llm/_torch/attention_backend/fmha/phased.pytensorrt_llm/_torch/attention_backend/fmha/registry.pytensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdtests/unittest/_torch/attention_backend/test_attention_op_sync.py
|
PR_Github #53256 [ run ] triggered by Bot. Commit: |
5afc19b to
36d5f74
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
36d5f74 to
2b6564c
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
If |
|
And what about |
2b6564c to
aad1559
Compare
|
/bot run --disable-fail-fast --add-multi-gpu-test |
|
PR_Github #53277 [ run ] triggered by Bot. Commit: |
|
PR_Github #54751 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54813 [ run ] triggered by Bot. Commit: |
|
PR_Github #54813 [ run ] completed with state
|
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
ff1973c to
79d1e09
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54900 [ run ] triggered by Bot. Commit: |
|
PR_Github #54900 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54919 [ run ] triggered by Bot. Commit: |
|
PR_Github #54919 [ run ] completed with state |
…ntion backend (NVIDIA#15204) Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
…ntion backend (NVIDIA#15204) Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>
5 participants
Description
JIRA: https://jirasw.nvidia.com/browse/TRTLLM-12807
Refactor the TRTLLM PyTorch attention FMHA implementation into explicit FMHA library classes under
tensorrt_llm/_torch/attention_backend/fmha/.This change:
Fmhainterface andPhasedFmhahelper for context/generation and MHA/MLA dispatch.FlashInferTrtllmGenFmha.thop.attentionpath intoFallbackFmha.TLLM_FMHA_LIBS.attention_backend/trtllm_gen.pycompatibility wrapper.Test Coverage
python3 -m pytest -q tests/unittest/_torch/attention_backend/test_attention_op_sync.pytests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_bf16[latency]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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.