[None][feat] Add fused T5 encoder self-attention path in bertAttentionPlugin

[zhanghuanzj]

New Features

[None][perf] Add fused single-kernel T5 encoder self-attention, ~2.5x faster than the unfused path

Description

GPU-side optimization for the T5 / mT5 encoder self-attention pipeline in bertAttentionPlugin. The change is gated behind an env var / build flag and is disabled by default, falling back to the original unfused path whenever the shape/arch is unsupported.

Fused T5 Attention

A new kernel family (fusedT5AttentionKernels.{h,cu}, exposed via FusedT5AttentionRunner) collapses the entire T5 encoder attention into a single launch, referencing the same "fuse everything that touches the [B,H,S,S] score matrix" idea used by fused MHA kernels.

Motivation. The current encoder self-attention runs as ~6 separate ops — QKV split → cuBLAS Q·Kᵀ → invokeAddRelativeAttentionBiasUnaligned → invokeMaskedSoftmax → cuBLAS S·V → transpose. Each step round-trips through HBM, and the dominant cost is the [B,H,S,S] intermediate score tensor plus the T5 relative-position bias. For T5-base at B=32, S=512 the score tensor alone is ~6 MB/batch written and re-read every layer.

Design. One CTA per (q_tile, batch, head); everything stays in SMEM/registers between phases:

Phase	Work	Notes
Load	Q tile → SMEM, per-head bucket bias [num_buckets] → SMEM, Q → WMMA regs	Q fragment loaded once, reused across all K tiles
1	S = Q·Kᵀ (WMMA, fp32 acc)	head_size is a compile-time template param
2	scale + T5 bias via LUT lookup + online softmax	bias += sBias[bucketTable[k-q + S-1]], single K-loop
3	rescale accumulator, load V (reuses K SMEM), O += P·V (WMMA)
Final	out = O / rowSum	supports padded and packed (cu_seqlens) layouts

Engineering details.

• Bucket LUT instead of dense bias. T5 relative bias only depends on k - q, so an offline [2S-1] int16 LUT (bit-exact with HF _relative_position_bucket, bidirectional) + a [H, num_buckets] bias table fully replace the [1,H,S,S] explicit tensor.
• Per-plugin device LUT. ensureFusedT5BucketTable (re)builds a plugin-owned device buffer on demand; this replaces a constant-memory design so multiple T5 models with different shapes can coexist in one process. Clone nulls out the copied raw pointer to avoid a double cudaFree; terminate frees it.
• WMMA fast path for SM80/86/89/90; a SIMT reference fallback covers SM70/75 for correctness only.
• Gating. Engages only for implicit relative-bias mode (mMaxDistance > 0), head_size ∈ {32,64,128}, seq_len ≤ 2048, fp16/bf16, SM ≥ 70, and within a conservative SMEM budget; otherwise isSupported() returns false and the caller uses the legacy path.

Enablement (any one of):

• trtllm-build --use_fused_t5_attention enable (added to the cli_plugin_args allowlist);
• env var TRTLLM_ENABLE_FUSED_T5_ATTENTION=1 (must be visible at both build and runtime; the C++ gate is read in the plugin's initialize());
• Python PluginConfig(use_fused_t5_attention=True), which propagates to the env var via a pydantic validator.

Expected gains.

• Intermediate [B,H,S,S] score tensor and the [1,H,S,S] explicit bias: eliminated from HBM.
• Kernel/GEMM launches per attention: ~6 → 1.
• Q loaded from SMEM once and cached in WMMA registers across all K tiles.
• On A100, T5-base encoder (H=12, D=64, num_buckets=32, B=32, S=512): ~0.85 ms → ~0.34 ms (~2.5x).

Test Coverage

• Numerical correctness — cpp/tests/unit_tests/kernels/fusedT5AttentionKernelsTest.cu compares the fused kernel against a reference implementation across head_size ∈ {32,64,128}, multiple seq_len / num_buckets / batch / num_heads, both fp16 and bf16.
• Bucket LUT — bit-exact check against HF _relative_position_bucket (bidirectional).
• Layouts — both padded and packed (removePadding + cu_seqlens) inputs.
• Gating / Fallback — isSupported / isShapeSupported return false on unsupported shape/arch, and the plugin uses the legacy path.
• End-to-end — standard T5 encoder flow (convert_checkpoint.py → trtllm-build ... --bert_attention_plugin float16 --use_fused_t5_attention enable) validated against the unfused baseline within fp16/bf16 tolerance.

Notes

• Scope is intentionally T5 encoder self-attention (bidirectional). T5 decoder causal self-attention still uses gpt_attention_plugin and is untouched.
• The kernel is a specialized fast path, not a general attention op: seq_len ≤ 2048 (LUT is O(2S-1)) and only implicit relative-bias mode is fused. This trades generality for a smaller, faster kernel; anything outside the envelope silently falls back.
• Default-off and env/flag-gated, so behavior is unchanged unless explicitly enabled.

Summary by CodeRabbit

• New Features

  • Added optional fused T5 encoder self-attention for supported configurations.
  • Added the use_fused_t5_attention plugin option, disabled by default.
  • Added support for padded and packed inputs with automatic capability checks.
  • Enabled configuration through TRTLLM_ENABLE_FUSED_T5_ATTENTION=1.

• Tests

  • Added CUDA unit tests covering numerical accuracy, relative-position bucketing, supported configurations, and enablement behavior.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds an optional fused T5 encoder self-attention implementation with WMMA and SIMT CUDA paths, plugin integration, environment and Python configuration controls, bucket-table management, and CUDA unit tests.

Changes

Fused T5 attention

Layer / File(s)	Summary
Configuration and public API cpp/tensorrt_llm/common/envUtils.*, cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.h, tensorrt_llm/plugin/plugin.py	Adds the environment accessor, PluginConfig.use_fused_t5_attention, CLI wiring, fused attention parameters, and runner interfaces.
Bucket table and kernel execution cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.cu	Implements T5 bucket mapping, bias extraction, bucket-table initialization, WMMA and SIMT kernels, runtime dispatch, and support validation.
Bert attention integration cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.*	Adds conditional fused execution, bucket-table allocation and cleanup, clone handling, workspace sizing, and fallback to the existing attention path.
Kernel validation cpp/tests/unit_tests/kernels/fusedT5AttentionKernelsTest.cu, cpp/tests/unit_tests/kernels/CMakeLists.txt	Adds bucket mapping, capability, padded/packed layout, and fp16/bf16 numerical parity tests.

Estimated code review effort: 5 (Critical) | ~120 minutes

Suggested reviewers: yuxianq, lfr-0531, jieli-matrix, heyuhhh, zhaoyangwang-nvidia

Sequence Diagram(s)

sequenceDiagram
  participant PluginConfig
  participant BertAttentionPlugin
  participant FusedT5AttentionRunner
  participant CUDAKernel
  PluginConfig->>BertAttentionPlugin: enable fused T5 configuration
  BertAttentionPlugin->>FusedT5AttentionRunner: check runtime support and shape
  BertAttentionPlugin->>FusedT5AttentionRunner: submit QKV, bias, and bucket table
  FusedT5AttentionRunner->>CUDAKernel: dispatch WMMA or SIMT attention
  CUDAKernel-->>BertAttentionPlugin: write attention output

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 31.03% which is insufficient. The required threshold is 80.00%.	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.
Title check	✅ Passed	The title is concise, specific, and matches the main change: adding a fused T5 encoder self-attention path.
Description check	✅ Passed	The description covers the feature, enablement, tests, and fallback behavior, with only minor template sections omitted.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 8

🤖 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.

Inline comments:
In `@cpp/tensorrt_llm/common/envUtils.cpp`:
- Around line 283-288: Unify fused T5 environment parsing so all callers use the
same semantics: update getEnvEnableFusedT5Attention() and the fused T5 kernel
logic to rely on a shared helper that recognizes the supported enabled values,
including “1”, “true”, and “TRUE”. Remove duplicated parsing from
fusedT5AttentionKernels.cu and ensure bertAttentionPlugin and
FusedT5AttentionRunner::isSupported() consume this common result.

In `@cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.cu`:
- Line 2: Update the copyright year in the header comment of
fusedT5AttentionKernels.cu from 2020-2025 to 2020-2026, keeping the existing
NVIDIA copyright text unchanged.
- Around line 641-652: The local isEnabledByEnv() parser is inconsistent with
the shared configuration behavior. Replace its getenv/string comparison logic
with the shared common::getEnvEnableFusedT5Attention() accessor, or update
common::getBoolEnv() consistently if true values must be supported globally,
ensuring all fused T5 attention call sites use the same parsing rules.
- Around line 685-691: Move the cudaFuncSetAttribute call out of the
per-dispatch launchOne path and initialize the max dynamic shared-memory size
once for each fusedT5AttentionWmmaKernel<T, kHeadSize> specialization, or
protect the attribute update and kernel launch with synchronization so
concurrent dispatches cannot race. Preserve the existing kernel launch behavior
and ensure every specialization is configured before use.

In `@cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.h`:
- Line 2: Update the copyright header in fusedT5AttentionKernels.h to include
the current year 2026, changing the existing 2020-2025 range to 2020-2026.

In `@cpp/tests/unit_tests/kernels/fusedT5AttentionKernelsTest.cu`:
- Line 2: Update the copyright header in fusedT5AttentionKernelsTest.cu to use
the year range 2020-2026, matching the current-year requirement used by other
new files.
- Around line 273-347: Device allocations in the test leak when ASSERT_TRUE or
ASSERT_EQ returns early. Update the test’s device-buffer ownership in the
surrounding test function to use RAII, such as std::unique_ptr with a cudaFree
deleter or a scoped cleanup guard for dQkv, dBias, dOut, dTbl, dLen, and dCu,
and remove the reliance on trailing cudaFree calls so cleanup is unconditional.
- Around line 211-218: Update runNumericParityCase and the fused T5 attention
tests to execute numeric parity cases on pre-SM80 GPUs using the SIMT fallback
instead of skipping them; retain the SM80+ WMMA coverage and add or adapt a
dedicated SIMT parity test covering the same cases for SM70–79.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 61e447ba-99ae-4d0d-b6f1-82ada91f940b

📥 Commits

Reviewing files that changed from the base of the PR and between 6b9f6dc and 7404e8f.

📒 Files selected for processing (9)

• cpp/tensorrt_llm/common/envUtils.cpp
• cpp/tensorrt_llm/common/envUtils.h
• cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.cu
• cpp/tensorrt_llm/kernels/fusedT5AttentionKernels.h
• cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp
• cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h
• cpp/tests/unit_tests/kernels/CMakeLists.txt
• cpp/tests/unit_tests/kernels/fusedT5AttentionKernelsTest.cu
• tensorrt_llm/plugin/plugin.py