[None][feat] Add TransformerEngine FP8 attention backend for VisualGen

requirements.txt

@@ -89,6 +89,7 @@ mcp<2.0.0 # 2.0.0 removed mcp.server.fastmcp; scaffolding still uses the 1.x API
 apache-tvm-ffi==0.1.13.post2 # required by CUTLASS DSL 4.6.1's TVM-FFI provider
 torch-c-dlpack-ext==0.1.3 # used for reduce nvidia-cutlass-dsl host overhead, optional package for improved torch tensor calling perf
 flash-attn-4==4.0.0b19
+transformer_engine==2.15.0 # required by the TE FP8 attention backend for VisualGen
 mistral-common>=1.10.0
 torchao>=0.14.1,<0.16.0
 cuda-core

tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py

@@ -32,6 +32,7 @@
 from .flash_attn4 import FlashAttn4Attention
 from .interface import AttentionBackend, AttentionTensorLayout
 from .parallel import Attention2DAttention, RingAttention, UlyssesAttention, wrap_parallel_attention
+from .te import TEAttention
 from .trtllm import TrtllmAttention, TrtllmAttentionMetadata
 from .utils import create_attention, get_visual_gen_attention_backend
 from .vanilla import VanillaAttention
@@ -45,6 +46,7 @@
     "CuTeDSLAttention",
     "VSAAttention",
     "FlashAttn4Attention",
+    "TEAttention",
     "TrtllmAttention",
     "TrtllmAttentionMetadata",
     "UlyssesAttention",

tensorrt_llm/_torch/visual_gen/attention_backend/te.py

@@ -0,0 +1,161 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-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.
+"""TransformerEngine FP8 attention backend for visual generation (diffusion) models.
+
+Uses TransformerEngine's ``DotProductAttention`` under ``fp8_autocast`` with
+``Float8CurrentScaling(fp8_dpa=True, fp8_mha=True)``.  Operates in NHD layout
+([B, S, H, D]) which maps directly to TE's ``qkv_format="bshd"`` -- no
+transpose overhead.
+
+``forward`` is decorated with ``@torch.compiler.disable`` because TE FP8
+modules graph-break under torch.compile.
+"""
+
+import math
+from typing import Any, Optional, Tuple
+
+import torch
+
+from ...attention_backend.interface import PredefinedAttentionMask
+from .interface import AttentionBackend, AttentionTensorLayout
+
+try:
+    from transformer_engine.common.recipe import Float8CurrentScaling
+    from transformer_engine.pytorch import DotProductAttention, fp8_autocast
+
+    _TE_AVAILABLE = True
+except ImportError:
+    _TE_AVAILABLE = False
+
+
+class TEAttention(AttentionBackend):
+    """FP8 attention via TransformerEngine ``DotProductAttention``.
+
+    FP8 is always enabled -- this backend exists to get FP8 attention.
+    No KV cache: diffusion models recompute attention each denoising step.
+    For BF16 attention use ``VANILLA``.
+
+    Does not support ``forward_with_lse``: TE's ``DotProductAttention`` does not
+    expose softmax stats, so the LSE would have to be recomputed from a separate
+    O(S^2) fp32 score matrix -- strictly more work than the attention itself.
+    That rules TE out as an inner backend for ``Attention2DAttention`` and
+    ``RingAttention``, which raise at construction on ``support_lse() is False``.
+    ``UlyssesAttention`` does not need LSE and still works.
+    """
+
+    def __init__(
+        self,
+        layer_idx: int = 0,
+        num_heads: int = 8,
+        head_dim: int = 64,
+        num_kv_heads: Optional[int] = None,
+        dtype: Optional[torch.dtype] = None,
+        quant_attention_config: Optional[Any] = None,
+        fp8_group: Optional[Any] = None,
+        **kwargs,
+    ):
+        if not _TE_AVAILABLE:
+            raise ImportError(
+                "TransformerEngine is required for the TE attention backend. "
+                "Install transformer_engine before using backend='TE'."
+            )
+        if quant_attention_config is not None:
+            # TE drives its own FP8 recipe; swallowing this would hide a config mistake.
+            raise NotImplementedError(
+                "TE attention backend does not honor quant_attention_config -- it always "
+                "runs FP8 via its own Float8CurrentScaling recipe. Drop quant_attention_config "
+                "or pick a backend that applies it."
+            )
+        self.layer_idx = layer_idx
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.num_kv_heads = num_kv_heads or num_heads
+        self.dtype = dtype
+        self.scale = 1.0 / math.sqrt(head_dim)
+        # Current scaling: TE gates delayed scaling's update on grad, so it never calibrates here.
+        self.recipe = Float8CurrentScaling(fp8_dpa=True, fp8_mha=True)
+        # None = no amax reduction: per-rank scales, no world-wide collective per call.
+        self.fp8_group = fp8_group
+        # Cache per trait: avoids rebuilding a DotProductAttention on every forward.
+        self._attn_ops: dict[tuple, Any] = {}
+
+    def _get_attn_op(self, num_gqa_groups: Optional[int], attn_mask_type: str) -> Any:
+        traits = (self.num_heads, self.head_dim, num_gqa_groups, attn_mask_type)
+        if traits not in self._attn_ops:
+            op = DotProductAttention(
+                self.num_heads,
+                self.head_dim,
+                num_gqa_groups=num_gqa_groups,
+                attn_mask_type=attn_mask_type,
+                softmax_scale=self.scale,
+                qkv_format="bshd",
+            )
+            # eval(): cuDNN's FP8 gate demands head_dim == 128 on sm < 100 while is_training.
+            self._attn_ops[traits] = op.eval()
+        return self._attn_ops[traits]
+
+    def _parse_inputs(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attention_mask: PredefinedAttentionMask,
+        key_padding_mask: Optional[torch.Tensor],
+    ) -> Tuple[Optional[int], str]:
+        if key_padding_mask is not None:
+            raise NotImplementedError("TE attention backend does not yet support key_padding_mask.")
+        if attention_mask == PredefinedAttentionMask.CAUSAL:
+            attn_mask_type = "causal"
+        elif attention_mask is None or attention_mask == PredefinedAttentionMask.FULL:
+            # None means "no mask" across the backends, and Attention2D forwards it as-is.
+            attn_mask_type = "no_mask"
+        else:
+            raise NotImplementedError(
+                f"TE attention backend does not support attention_mask={attention_mask!r}. "
+                "Only PredefinedAttentionMask.FULL and CAUSAL are supported."
+            )
+        enable_gqa = self.num_heads != self.num_kv_heads
+        num_gqa_groups = k.shape[-2] if enable_gqa else None
+        return num_gqa_groups, attn_mask_type
+
+    @torch.compiler.disable
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        *,
+        attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """FP8 self/cross attention. q/k/v shape: [B, S, H, D]. Returns [B, S, H, D]."""
+        num_gqa_groups, attn_mask_type = self._parse_inputs(q, k, attention_mask, key_padding_mask)
+        attn_op = self._get_attn_op(num_gqa_groups, attn_mask_type)
+        with fp8_autocast(enabled=True, fp8_recipe=self.recipe, fp8_group=self.fp8_group):
+            # TE returns [B, S, H*D]; restore to [B, S, H, D].
+            out = attn_op(q, k, v, attention_mask=None)
+        return out.unflatten(-1, (self.num_heads, self.head_dim))
+
+    @property
+    def preferred_layout(self) -> AttentionTensorLayout:
+        return AttentionTensorLayout.NHD
+
+    @classmethod
+    def support_fused_qkv(cls) -> bool:
+        return False
+
+    @classmethod
+    def support_lse(cls) -> bool:
+        return False

tensorrt_llm/_torch/visual_gen/attention_backend/utils.py

@@ -37,7 +37,7 @@ def get_visual_gen_attention_backend(
     Get diffusion attention backend class by name.
 
     Args:
-        backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL")
+        backend_name: Backend identifier ("VANILLA", "TRTLLM", "FA4", "CUTEDSL", "TE")
 
     Returns:
         Diffusion attention backend class
@@ -51,10 +51,12 @@ def get_visual_gen_attention_backend(
                  Requires flash-attn package with cute interface
         - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA
                       from AttentionConfig.sparse_attention_config.
+        - "TE": TransformerEngine FP8 attention; requires transformer_engine package
     """
     # Lazy imports to avoid circular dependency
     from .cute_dsl import CuTeDSLAttention
     from .flash_attn4 import FlashAttn4Attention
+    from .te import TEAttention
     from .trtllm import TrtllmAttention
     from .vanilla import VanillaAttention
 
@@ -68,6 +70,8 @@ def get_visual_gen_attention_backend(
         return FlashAttn4Attention
     elif backend_name == "CUTEDSL":
         return CuTeDSLAttention
+    elif backend_name == "TE":
+        return TEAttention
     else:
         # Default to VANILLA for maximum compatibility
         return VanillaAttention

tensorrt_llm/visual_gen/args.py

@@ -97,10 +97,10 @@ class QuantAttentionConfig(StrictBaseModel):
 class AttentionConfig(StrictBaseModel):
     """Configuration for Attention layers."""
 
-    backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL"] = Field(
+    backend: Literal["VANILLA", "TRTLLM", "FA4", "CUTEDSL", "TE"] = Field(
         "VANILLA",
         status="prototype",
-        description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4, CUTEDSL",
+        description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, FA4, CUTEDSL, TE (TransformerEngine FP8)",
     )
     quant_attention_config: Optional[QuantAttentionConfig] = Field(
         None,

tests/unittest/_torch/visual_gen/test_attention_te.py

@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-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.
+"""Tests for the TransformerEngine FP8 attention backend (TEAttention).
+
+Three tests, each covering one failure mode of the backend:
+  1. numerics of ``forward``            -- output matches a full-precision reference
+  2. mask handling                      -- FULL / None / CAUSAL / unsupported
+  3. object contract                    -- layout, no-LSE, op cache, config rejection
+
+FP8 tolerances: TEAttention quantizes Q/K/V, so it sits at roughly 0.08 relative
+error against an fp32 reference (measured on GB300 across S=128..1024 for
+B=1 H=10 D=128). Comparisons against a full-precision reference therefore use
+cosine similarity plus a loose relative error rather than elementwise
+assert_close.
+"""
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask
+from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout
+from tensorrt_llm._torch.visual_gen.attention_backend.te import TEAttention
+
+try:
+    import transformer_engine  # noqa: F401
+
+    _te_available = True
+except ImportError:
+    _te_available = False
+
+pytestmark = pytest.mark.skipif(
+    not _te_available or not torch.cuda.is_available(),
+    reason="transformer_engine and GPU required",
+)
+
+FP8_COS_SIM = 0.99
+FP8_REL_ERR = 0.12
+
+
+@pytest.fixture
+def make_te_attn():
+    def _make(num_heads=4, head_dim=64, num_kv_heads=None):
+        return TEAttention(
+            layer_idx=0,
+            num_heads=num_heads,
+            head_dim=head_dim,
+            num_kv_heads=num_kv_heads or num_heads,
+        )
+
+    return _make
+
+
+def _qkv(B, S, H, D, Hkv=None, seed=0):
+    torch.manual_seed(seed)
+    kw = dict(device="cuda", dtype=torch.bfloat16)
+    return (
+        torch.randn(B, S, H, D, **kw),
+        torch.randn(B, S, Hkv or H, D, **kw),
+        torch.randn(B, S, Hkv or H, D, **kw),
+    )
+
+
+def _reference(q, k, v, D, is_causal=False, enable_gqa=False):
+    """Full-precision SDPA reference in [B, S, H, D] layout."""
+    out = F.scaled_dot_product_attention(
+        q.transpose(1, 2).float(),
+        k.transpose(1, 2).float(),
+        v.transpose(1, 2).float(),
+        scale=1.0 / (D**0.5),
+        is_causal=is_causal,
+        enable_gqa=enable_gqa,
+    )
+    return out.transpose(1, 2).to(q.dtype)
+
+
+def _assert_close_fp8(out, ref, what):
+    o, r = out.reshape(-1).float(), ref.reshape(-1).float()
+    cos = F.cosine_similarity(o, r, dim=0).item()
+    rel = (torch.linalg.vector_norm(o - r) / torch.linalg.vector_norm(r)).item()
+    assert cos > FP8_COS_SIM, f"{what}: cosine similarity {cos:.4f} <= {FP8_COS_SIM}"
+    assert rel < FP8_REL_ERR, f"{what}: relative error {rel:.4f} >= {FP8_REL_ERR}"
+
+
+@pytest.mark.parametrize("S,Hkv", [(64, None), (256, None), (1024, None), (256, 2)])
+def test_forward_matches_reference(make_te_attn, S, Hkv):
+    """forward() returns a well-formed [B, S, H, D] tensor matching full-precision SDPA.
+
+    Parametrized over sequence length and over GQA (Hkv=2 against H=8), because
+    the GQA path takes a different branch in _parse_inputs and previously did not
+    agree with the non-GQA one.
+    """
+    B, H, D = 1, 8, 64
+    attn = make_te_attn(num_heads=H, head_dim=D, num_kv_heads=Hkv)
+    q, k, v = _qkv(B, S, H, D, Hkv, seed=42)
+
+    with torch.no_grad():
+        out = attn(q, k, v)
+
+    assert out.shape == (B, S, H, D)
+    assert torch.isfinite(out).all(), "output contains NaN or Inf"
+    _assert_close_fp8(out, _reference(q, k, v, D, enable_gqa=Hkv is not None), f"S={S} Hkv={Hkv}")
+
+
+def test_attention_mask_handling(make_te_attn):
+    """FULL and None mean no mask, CAUSAL masks, anything else is rejected.
+
+    None is the no-mask convention across the backends -- VanillaAttention derives
+    causality as `attention_mask == CAUSAL`, and Attention2DAttention explicitly
+    permits None and forwards it to the inner backend unchanged.
+    """
+    B, S, H, D = 1, 64, 4, 64
+    attn = make_te_attn(num_heads=H, head_dim=D)
+    q, k, v = _qkv(B, S, H, D, seed=3)
+
+    with torch.no_grad():
+        out_full = attn(q, k, v, attention_mask=PredefinedAttentionMask.FULL)
+        out_none = attn(q, k, v, attention_mask=None)
+        out_causal = attn(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL)
+
+    # FP8 tolerance, not exact: current scaling rederives the scale on every call.
+    _assert_close_fp8(out_none, out_full, "attention_mask=None must mean FULL")
+    _assert_close_fp8(out_causal, _reference(q, k, v, D, is_causal=True), "causal")
+    assert not torch.allclose(out_full, out_causal, atol=1e-3), "causal must differ from full"
+
+    with pytest.raises(NotImplementedError, match="key_padding_mask"):
+        attn(q, k, v, key_padding_mask=torch.ones(B, S, device="cuda", dtype=torch.bool))
+    with pytest.raises(NotImplementedError, match="attention_mask"):
+        attn(q, k, v, attention_mask="bidirectional")
+
+
+def test_backend_contract(make_te_attn):
+    """Layout/LSE declarations, per-trait op caching, and config rejection."""
+    attn = make_te_attn(num_heads=4, head_dim=64)
+    assert attn.preferred_layout == AttentionTensorLayout.NHD
+    assert TEAttention.support_fused_qkv() is False
+
+    # No LSE: this is the flag Attention2D/Ring read to reject TE as an inner backend.
+    assert TEAttention.support_lse() is False
+
+    # Each trait keeps its own instance; rebuilding on a mask switch would reset FP8 calibration.
+    q, k, v = _qkv(1, 64, 4, 64, seed=5)
+    with torch.no_grad():
+        attn(q, k, v, attention_mask=PredefinedAttentionMask.FULL)
+        attn(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL)
+    op_full = attn._get_attn_op(None, "no_mask")
+    op_causal = attn._get_attn_op(None, "causal")
+    assert op_full is not op_causal
+    assert len(attn._attn_ops) == 2
+    with torch.no_grad():  # switching back reuses, never rebuilds
+        attn(q, k, v, attention_mask=PredefinedAttentionMask.FULL)
+    assert attn._get_attn_op(None, "no_mask") is op_full
+    assert len(attn._attn_ops) == 2
+
+    # TE runs its own FP8 recipe and must not silently absorb quant_attention_config.
+    with pytest.raises(NotImplementedError, match="quant_attention_config"):
+        TEAttention(num_heads=4, head_dim=64, quant_attention_config=object())