[wu6u3tw]

@coderabbitai summary

Description

• What: new VisualGen attention backend AttentionConfig(backend="TE"): TransformerEngine DotProductAttention (cuDNN fused FP8 attention) with a DelayedScaling(fp8_dpa=True, fp8_mha=True) recipe, NHD layout.
• Hardware: Blackwell and Rubin (SM 100 / 103 / 107). Supersedes the earlier Blackwell-only TE backend in [None][feat] Add TransformerEngine FP8 attention backend for VisualGen #17849.
• Inference scaling fix: DelayedScaling only updates its scales when autograd is enabled, so in grad-free inference they stay at 1.0. The attention output of a softmax spread over many keys is small and loses precision below E4M3's normal range at that scale. The backend writes the forward Q/K/V, O and S scales from the current tensors before every call.
• Safety check: on each module's first call, the backend verifies TE's fp8_meta slot layout against the recorded amax and raises on mismatch instead of mis-scaling.
• API: adds "TE" to AttentionConfig.backend (public VisualGen API, needs team review). quant_attention_config is rejected for TE, since TE owns its FP8 recipe.
• Lossy: FP8 attention; the accuracy study is available for the review sync.

Test Coverage

• New tests/unittest/_torch/visual_gen/test_attention_te.py:
  • config accepted / quant_attention_config rejected (CPU; added to l0_cpu.yml)
  • backend wiring and NHD layout
  • accuracy vs FP32 SDPA: MHA, GQA, cross-attention, causal
  • per-call scales track inputs in grad-free inference (no_grad, inference_mode)
  • negative control: stale scales lose accuracy on the same inputs
  • key_padding_mask rejected
• GPU tests added to l0_b200.yml.
• Ran on Rubin (SM 10.7, TE 2.20): 12 passed, 1 skipped (causal on a cross-attention shape), 0 failed.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

🤖 Generated with Claude Code

[YihuiLu512]

Heads-up: #20214 moved visual_gen/attention_backend/ to visual_gen/attention/backends/; the old package is now a deprecated shim that only forwards package-root exports.

After a rebase, your utils.py edits follow the rename to attention/backends/utils.py, while the new te.py stays at the old attention_backend/te.py. As a result, from .te import TEAttention raises ImportError, and neither a merge conflict nor pre-merge CI will flag it.

When you rebase, could you put te.py under tensorrt_llm/_torch/visual_gen/attention/backends/?