diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/te.py b/tensorrt_llm/_torch/visual_gen/attention_backend/te.py new file mode 100644 index 000000000000..d2ddd3b0ce84 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/te.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +Runs TransformerEngine's ``DotProductAttention`` (cuDNN fused FP8 attention) +under ``fp8_autocast`` with a ``DelayedScaling(fp8_dpa=True, fp8_mha=True)`` +recipe, in NHD layout (``qkv_format="bshd"``, no transpose). + +Inference-correct FP8 scales: ``DelayedScaling`` only refreshes its scales from +the amax history when autograd is enabled (TE runs the scale update in the +autocast exit only under ``torch.is_grad_enabled()``). Diffusion inference runs +grad-free, so the scales would stay at their initial 1.0. E4M3's relative +precision does not depend on the scale inside its normal range, but values +below it lose precision. That hits the attention output when softmax is spread +over many keys: each output row averages many V rows, so its magnitude is about +std(v) / sqrt(seq_len), e.g. ~4e-3 at 75,600 tokens. + +This backend therefore writes the forward scales before every call from the +current tensors: Q/K/V = 448 / amax(q, k, v); O = 448 / amax(v), since each +output row is a convex combination of V rows, so |O| <= amax(v); S = 448, +since softmax probabilities lie in [0, 1]. +""" + +import math +from typing import Optional + +import torch + +from ...attention.backends.interface import PredefinedAttentionMask +from .interface import AttentionBackend, AttentionTensorLayout + +try: + from transformer_engine.common.recipe import DelayedScaling + from transformer_engine.pytorch import DotProductAttention, fp8_autocast +except ImportError: # TE absent: fail at construction, not at import. + DotProductAttention = None + DelayedScaling = None + fp8_autocast = None + +_FP8_E4M3_MAX = 448.0 +# Forward fp8_meta scale slots of DotProductAttention: Q/K/V input, output O, +# softmax S. Verified at runtime on each module's first call (see _verify_slots), +# so a TE layout change raises instead of silently mis-scaling. +_SLOT_QKV, _SLOT_O, _SLOT_S = 2, 3, 8 + + +class TEAttention(AttentionBackend): + """FP8 attention via TransformerEngine ``DotProductAttention``. + + No KV cache (diffusion recomputes every step). Always FP8; use another + backend for BF16 attention. Configure with ``AttentionConfig(backend="TE")`` + and no ``quant_attention_config``: TE owns its FP8 recipe. + """ + + 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, + **kwargs, + ): + if DotProductAttention is None: + raise ImportError( + "TransformerEngine is required for the TE attention backend " + "(transformer_engine.pytorch.DotProductAttention is not importable)." + ) + 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) + self.recipe = DelayedScaling(fp8_dpa=True, fp8_mha=True) + # DotProductAttention holds FP8 metadata; rebuild only when the + # (GQA groups, mask) traits change. + self._attn_op = None + self._traits = None + self._slots_verified = False + + def _lazy_init(self, num_gqa_groups: Optional[int], attn_mask_type: str) -> None: + traits = (num_gqa_groups, attn_mask_type) + if traits != self._traits: + self._attn_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", + ) + self._traits = traits + self._slots_verified = False + + @torch.no_grad() + def _verify_slots(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + """Check the slot map against the amax TE recorded on its first call.""" + amax = self._attn_op.fp8_meta["scaling_fwd"].amax_history[0].float() + want_qkv = max(q.abs().amax().item(), k.abs().amax().item(), v.abs().amax().item()) + got_qkv, got_s = amax[_SLOT_QKV].item(), amax[_SLOT_S].item() + if abs(got_qkv - want_qkv) > 1e-3 * max(want_qkv, 1.0) or not 0.0 < got_s <= 1.0 + 1e-3: + raise RuntimeError( + f"TE fp8_meta slot layout mismatch (layer {self.layer_idx}): " + f"slot {_SLOT_QKV}={got_qkv} vs amax(q, k, v)={want_qkv}; " + f"slot {_SLOT_S}={got_s} (softmax, expected in (0, 1]). " + f"Recorded amax: {amax.tolist()}" + ) + self._slots_verified = True + + @torch.no_grad() + def _set_forward_scales(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + """Write this call's forward FP8 scales in place (on device, no host sync).""" + meta = self._attn_op.fp8_meta["scaling_fwd"] + amax_v = v.abs().amax().float() + amax_qkv = torch.maximum(torch.maximum(q.abs().amax(), k.abs().amax()).float(), amax_v) + scale = meta.scale + scale[_SLOT_QKV] = _FP8_E4M3_MAX / amax_qkv.clamp_min(1e-12) + scale[_SLOT_O] = _FP8_E4M3_MAX / amax_v.clamp_min(1e-12) + scale[_SLOT_S] = _FP8_E4M3_MAX + if getattr(meta, "scale_inv", None) is not None: # older TE keeps it separately + meta.scale_inv.copy_(1.0 / scale) + + @torch.compiler.disable + def forward( + self, + q: torch.Tensor, + k: torch.Tensor | None = None, + v: torch.Tensor | None = None, + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + key_padding_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """FP8 attention. q/k/v are NHD ([B, S, H, D]); returns [B, S, H, D].""" + if k is None or v is None: + raise ValueError("TEAttention does not support fused QKV; pass q, k and v.") + if q.dim() != 4 or q.shape[-1] != self.head_dim: + raise ValueError( + f"Expected q of shape [B, S, H, D={self.head_dim}], got {tuple(q.shape)}" + ) + if key_padding_mask is not None: + raise NotImplementedError("TEAttention does not support key_padding_mask.") + + is_causal = attention_mask == PredefinedAttentionMask.CAUSAL + num_gqa_groups = k.shape[-2] if self.num_heads != self.num_kv_heads else None + self._lazy_init(num_gqa_groups, "causal" if is_causal else "no_mask") + + if "scaling_fwd" not in getattr(self._attn_op, "fp8_meta", {}): + # TE builds fp8_meta inside the first FP8 forward. Run once to + # create it (output discarded), then scale correctly below. + with fp8_autocast(enabled=True, fp8_recipe=self.recipe): + self._attn_op(q, k, v, attention_mask=None) + if not self._slots_verified: + self._verify_slots(q, k, v) + self._set_forward_scales(q, k, v) + + with fp8_autocast(enabled=True, fp8_recipe=self.recipe): + out = self._attn_op(q, k, v, attention_mask=None) + # TE returns [B, S, H*D]; restore [B, S, H, D]. + return out.unflatten(-1, (self.num_heads, self.head_dim)) + + @property + def preferred_layout(self) -> AttentionTensorLayout: + return AttentionTensorLayout.NHD diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index e3109e354e50..e25f40ca5e7a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -38,7 +38,7 @@ def get_visual_gen_attention_backend( Get diffusion attention backend class by name and sparse algorithm. Args: - backend_name: Backend identifier ("VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4") + backend_name: Backend identifier ("VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4", "TE") sparse_algorithm: ``algorithm`` of the sparse attention config, if any. VSA ("vsa") and SOL ("sol_attn") have their own TRTLLM and CUTEDSL classes; other algorithms, such as skip-softmax, run in the dense class. @@ -102,6 +102,10 @@ def get_visual_gen_attention_backend( from .cudnn import CuDNNAttention return CuDNNAttention + elif backend_name == "TE": + from .te import TEAttention + + return TEAttention else: # Default to VANILLA for maximum compatibility from .vanilla import VanillaAttention @@ -130,7 +134,7 @@ def create_attention( internally, simplifying the forward() call. Args: - backend: Backend identifier ("VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4") + backend: Backend identifier ("VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4", "TE") layer_idx: Layer index in the model num_heads: Number of attention heads head_dim: Dimension per head diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index b08d322f95a0..4242524ebb08 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -107,10 +107,11 @@ class QuantAttentionConfig(StrictBaseModel): class AttentionConfig(StrictBaseModel): """Configuration for Attention layers.""" - backend: Literal["VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4"] = Field( + backend: Literal["VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4", "TE"] = Field( "VANILLA", status="prototype", - description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, CUDNN, FLASHINFER, CUTEDSL, FA4", + description="Attention backend: VANILLA (PyTorch SDPA), TRTLLM, CUDNN, FLASHINFER, CUTEDSL, " + "FA4, TE (TransformerEngine FP8)", ) quant_attention_config: Optional[QuantAttentionConfig] = Field( None, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2ca039299f9c..bb7f3b3c702c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -312,6 +312,7 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py + - unittest/_torch/visual_gen/test_attention_te.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_trtllm_attention_metadata.py - unittest/_torch/visual_gen/test_attention_fa4.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 5e57f3d4f85f..830ebad7d16b 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -96,6 +96,8 @@ l0_cpu: - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py - unittest/_torch/visual_gen/test_attention_flashinfer.py::test_flashinfer_backend_is_registered - unittest/_torch/visual_gen/test_attention_integration.py + - unittest/_torch/visual_gen/test_attention_te.py::test_te_backend_config_accepted + - unittest/_torch/visual_gen/test_attention_te.py::test_te_backend_rejects_quant_attention_config - unittest/_torch/visual_gen/test_cache_dit.py - unittest/_torch/visual_gen/test_executor_lifecycle.py - unittest/_torch/visual_gen/test_flux_infer.py diff --git a/tests/unittest/_torch/visual_gen/test_attention_te.py b/tests/unittest/_torch/visual_gen/test_attention_te.py new file mode 100644 index 000000000000..81b597769b46 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_te.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 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 visual-gen attention backend.""" + +import pytest +import torch +import torch.nn.functional as F +from pydantic import ValidationError + +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.visual_gen.attention_backend import create_attention +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig + +try: + import transformer_engine.pytorch # noqa: F401 + + _te_available = True +except ImportError: + _te_available = False + +# (name, batch, num_heads, num_kv_heads, seq_len_q, seq_len_kv, head_dim) +SHAPES = [ + ("mha", 2, 8, 8, 512, 512, 128), + ("gqa", 1, 16, 4, 1024, 1024, 128), + ("cross", 1, 8, 8, 1024, 512, 128), +] + + +# Set from Rubin measurements (see the flat-softmax tests below). +MAX_FLAT_SOFTMAX_REL_ERR = 0.06 +STALE_TO_FIXED_MIN_RATIO = 3.0 +# Small V keeps the attention output (~std(v) / sqrt(seq)) far below E4M3's normal +# range at scale 1.0; per-call scales are invariant to this magnitude. +FLAT_SOFTMAX_V_AMP = 0.1 + + +def _require_te_fp8() -> None: + if not torch.cuda.is_available(): + pytest.skip("TE FP8 attention requires CUDA.") + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("TE FP8 fused attention test targets Blackwell/Rubin-class GPUs (sm100+).") + if not _te_available: + pytest.fail("TransformerEngine is not importable; the TE backend tests need it.") + + +def _make_qkv(batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim, amp=1.0, v_amp=1.0): + torch.manual_seed(0) + dev = torch.device("cuda") + q = amp * torch.randn(batch, seq_q, num_heads, head_dim, device=dev, dtype=torch.bfloat16) + k = amp * torch.randn(batch, seq_kv, num_kv_heads, head_dim, device=dev, dtype=torch.bfloat16) + v = v_amp * torch.randn(batch, seq_kv, num_kv_heads, head_dim, device=dev, dtype=torch.bfloat16) + return q, k, v + + +def _reference(q, k, v, is_causal): + """FP32 SDPA in NHD layout.""" + rep = q.shape[2] // k.shape[2] + k = k.repeat_interleave(rep, dim=2) + v = v.repeat_interleave(rep, dim=2) + out = F.scaled_dot_product_attention( + q.float().transpose(1, 2), + k.float().transpose(1, 2), + v.float().transpose(1, 2), + is_causal=is_causal, + ) + return out.transpose(1, 2) + + +def _rel_err(out, ref): + return ((out.float() - ref).norm() / ref.norm()).item() + + +def _make_te(num_heads, num_kv_heads, head_dim): + return create_attention( + backend="TE", + layer_idx=0, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + dtype=torch.bfloat16, + attention_config=AttentionConfig(backend="TE"), + ) + + +def test_te_backend_config_accepted(): + assert AttentionConfig(backend="TE").backend == "TE" + + +def test_te_backend_rejects_quant_attention_config(): + """TE owns its FP8 recipe; an explicit quant_attention_config is a config mistake.""" + with pytest.raises(ValidationError, match="quant_attention_config requires backend"): + AttentionConfig( + backend="TE", quant_attention_config=QuantAttentionConfig(qk_dtype="fp8", v_dtype="fp8") + ) + + +def test_te_backend_wiring(): + _require_te_fp8() + attention = _make_te(8, 8, 128) + assert type(attention).__name__ == "TEAttention" + assert attention.preferred_layout == AttentionTensorLayout.NHD + assert not attention.support_fused_qkv() and not attention.support_lse() + + +@pytest.mark.parametrize("shape", SHAPES, ids=[s[0] for s in SHAPES]) +@pytest.mark.parametrize("is_causal", [False, True]) +def test_te_fp8_attention_accuracy(shape, is_causal): + _require_te_fp8() + _, batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim = shape + if is_causal and seq_q != seq_kv: + pytest.skip("causal attention is only exercised for self-attention shapes") + q, k, v = _make_qkv(batch, num_heads, num_kv_heads, seq_q, seq_kv, head_dim) + attention = _make_te(num_heads, num_kv_heads, head_dim) + mask = PredefinedAttentionMask.CAUSAL if is_causal else PredefinedAttentionMask.FULL + + with torch.no_grad(): + out = attention.forward(q, k, v, attention_mask=mask) + ref = _reference(q, k, v, is_causal) + + assert out.shape == q.shape + cos = F.cosine_similarity(out.float().flatten(), ref.flatten(), dim=0).item() + assert cos > 0.995, f"cosine {cos} vs FP32 reference" + + +def _run_te(attention, q, k, v, grad_mode): + ctx = torch.no_grad() if grad_mode == "no_grad" else torch.inference_mode() + with ctx: + return attention.forward(q, k, v) + + +@pytest.mark.parametrize("grad_mode", ["no_grad", "inference_mode"]) +def test_te_fp8_scales_track_inputs_without_autograd(grad_mode): + """Grad-free inference must write per-call FP8 scales, not keep the initial 1.0. + + Small Q/K give a near-uniform softmax over 16k keys, so each output row + averages many V rows and is small (~std(v) / sqrt(seq)). At scale 1.0 such + outputs fall into E4M3's subnormal range and lose precision; per-call + scales keep them in range. (Large Q/K instead give a peaked softmax whose + error is dominated by E4M3 rounding of Q and K, which no scale changes, so + that regime cannot detect the bug.) The two calls differ 2x in magnitude, + so the checked scales must come from the current inputs. + """ + _require_te_fp8() + num_heads, head_dim, seq = 8, 128, 16384 + attention = _make_te(num_heads, num_heads, head_dim) + + for amp in (0.25, 0.5): + q, k, v = _make_qkv( + 1, num_heads, num_heads, seq, seq, head_dim, amp=amp, v_amp=FLAT_SOFTMAX_V_AMP + ) + err = _rel_err(_run_te(attention, q, k, v, grad_mode), _reference(q, k, v, False)) + assert err < MAX_FLAT_SOFTMAX_REL_ERR, f"amp={amp}: relative error {err}" + + scale = attention._attn_op.fp8_meta["scaling_fwd"].scale.float() + amax_qkv = max(t.abs().amax().item() for t in (q, k, v)) + assert scale[2].item() == pytest.approx(448.0 / amax_qkv, rel=1e-3) + assert scale[3].item() == pytest.approx(448.0 / v.abs().amax().item(), rel=1e-3) + assert scale[8].item() == pytest.approx(448.0) + + +def test_te_fp8_stale_scales_lose_flat_softmax(monkeypatch): + """Negative control: without per-call scales the same inputs lose accuracy. + + Shows the flat-softmax regime above actually detects stale scales. + """ + _require_te_fp8() + num_heads, head_dim, seq = 8, 128, 16384 + q, k, v = _make_qkv( + 1, num_heads, num_heads, seq, seq, head_dim, amp=0.25, v_amp=FLAT_SOFTMAX_V_AMP + ) + ref = _reference(q, k, v, False) + fixed = _rel_err(_run_te(_make_te(num_heads, num_heads, head_dim), q, k, v, "no_grad"), ref) + + from tensorrt_llm._torch.visual_gen.attention_backend import te as te_module + + monkeypatch.setattr(te_module.TEAttention, "_set_forward_scales", lambda self, q, k, v: None) + stale = _rel_err(_run_te(_make_te(num_heads, num_heads, head_dim), q, k, v, "no_grad"), ref) + assert stale > STALE_TO_FIXED_MIN_RATIO * fixed, f"stale {stale} vs per-call {fixed}" + + +def test_te_rejects_key_padding_mask(): + _require_te_fp8() + q, k, v = _make_qkv(1, 8, 8, 256, 256, 128) + attention = _make_te(8, 8, 128) + with pytest.raises(NotImplementedError, match="key_padding_mask"): + attention.forward( + q, k, v, key_padding_mask=torch.ones(1, 256, dtype=torch.bool, device="cuda") + )