Bug: TensorRT 11.4 fails to build IAttention followed by FP8 Q/DQ, breaking FLUX.2 FP8

Start by running the provided repro_attention_fp8_consumer.py with TensorRT 11.4 and the -quantize_ppg=off build route, then compare with TensorRT 11.1 as described. The issue identifies no repository file or test to change and says a TensorRT fix is needed; done means the affected network builds without disabling quantize propagation.

This was generated by AI during triage.

Affected area

Model resolution or TensorRT engine build

TensorRT-Model-Connect version

39565ca554122730cd57b7adbdbb4614ffa06d5e (main). The reproducer below is standalone and uses only the TensorRT Python API.

Installation method

Development container

Model and revision

black-forest-labs/FLUX.2-dev at 26afe3a78bb242c0a8bb181dcc8937bb16e5c66c, case flux-2-dev-fp8 (FP16 base with FP8 linears). The reproducer needs no checkpoint.

Environment

• OS: Ubuntu 24.04.5 (development container built from the repository Dockerfile)
• GPU: NVIDIA GB300
• Driver: 595.58.03
• CUDA: 13.4.3
• TensorRT: 11.4.0.106 fails. 11.1.0.106 builds the same network. Both come from the tensorrt-cu13 wheels.

Summary

Since the update to TensorRT 11.4 (#1811), the flux-2-dev-fp8 DiT engine no longer builds (TRT engine serialization failed for FLUX.2 DiT). The DiT's IAttention layers are not quantized, but each attention output feeds an FP8 linear (to_out) whose input is quantized to FP8. TensorRT 11.4 moves that FP8 quantize back into the attention. The lowered softmax then gets an FP8 input (bmm1_out), which has no implementation, and the build fails. TensorRT 11.1 builds the identical network.

The reproducer below fails with any FP8 Quantize/Dequantize consumer after an unquantized IAttention.

Minimal reproducer

# TensorRT 11.4.0.106 (current development container)
python3 repro_attention_fp8_consumer.py                                  # BUILD_FAILED
python3 repro_attention_fp8_consumer.py --no-matmul                      # BUILD_FAILED (Q/DQ only)
python3 repro_attention_fp8_consumer.py --build-route=-quantize_ppg=off  # BUILD_OK

# TensorRT 11.1.0.106, same script
pip install --no-deps --target ./trt111 tensorrt==11.1.0.106 tensorrt-cu13==11.1.0.106 \
  tensorrt-cu13-bindings==11.1.0.106 tensorrt-cu13-libs==11.1.0.106
PYTHONPATH=./trt111 LD_LIBRARY_PATH=./trt111/tensorrt_libs:$LD_LIBRARY_PATH \
  python3 repro_attention_fp8_consumer.py                                # BUILD_OK

repro_attention_fp8_consumer.py:

"""Minimal TensorRT 11.4 build failure: an FP8 Q/DQ consumer after IAttention.

Network (strongly typed, FP16):
    q, k, v [1, H, S, D] -> IAttention(SOFTMAX) -> transpose/reshape to [S, H*D]
    -> Quantize(FP8) -> Dequantize -> MatMul with an FP8 weight constant
    + Dequantize (transposed B)

The IAttention layer itself is not quantized. Prints the TensorRT version,
the configuration and whether the build succeeded, then the first errors.
Optional --build-route sets IBuilderConfig.build_route.
"""

import argparse

import numpy as np
import tensorrt as trt


class CollectingLogger(trt.ILogger):
    def __init__(self) -> None:
        super().__init__()
        self.errors: list[str] = []

    def log(self, severity, msg) -> None:
        if severity <= trt.ILogger.Severity.ERROR:
            self.errors.append(msg)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seq", type=int, default=256)
    parser.add_argument("--heads", type=int, default=2)
    parser.add_argument("--head-dim", type=int, default=64)
    parser.add_argument("--no-matmul", action="store_true", help="mark the dequantized tensor as output")
    parser.add_argument("--build-route", default="")
    parser.add_argument("--list-routes", action="store_true")
    args = parser.parse_args()

    heads, head_dim, seq = args.heads, args.head_dim, args.seq
    logger = CollectingLogger()
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    if args.list_routes:
        print(getattr(config, "all_build_routes", "<no all_build_routes attribute>"))
        return 0
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
    )
    keep_alive = []

    def scalar(value: float) -> trt.ITensor:
        weights = np.array(value, dtype=np.float32)
        keep_alive.append(weights)
        return network.add_constant((), trt.Weights(weights)).get_output(0)

    shape = (1, heads, seq, head_dim)
    query, key, value = (network.add_input(name, trt.float16, shape) for name in ("q", "k", "v"))
    attention = network.add_attention(
        query, key, value, trt.AttentionNormalizationOp.SOFTMAX, False
    )
    width = heads * head_dim
    rows = network.add_shuffle(attention.get_output(0))
    rows.first_transpose = (0, 2, 1, 3)
    rows.reshape_dims = (seq, width)
    activation_scale = scalar(0.05)
    quantized = network.add_quantize(rows.get_output(0), activation_scale, trt.DataType.FP8)
    output = network.add_dequantize(quantized.get_output(0), activation_scale, trt.float16).get_output(0)
    if not args.no_matmul:
        codes = np.full((width, width), 0x38, dtype=np.uint8)  # E4M3 1.0
        keep_alive.append(codes)
        weight = network.add_constant(
            (width, width), trt.Weights(trt.DataType.FP8, codes.ctypes.data, codes.size)
        )
        weight_dq = network.add_dequantize(weight.get_output(0), scalar(0.01), trt.float16)
        output = network.add_matrix_multiply(
            output, trt.MatrixOperation.NONE, weight_dq.get_output(0), trt.MatrixOperation.TRANSPOSE
        ).get_output(0)
    network.mark_output(output)

    if args.build_route:
        config.build_route = args.build_route
    plan = builder.build_serialized_network(network, config)
    print(
        f"tensorrt={trt.__version__} shape={shape} matmul={not args.no_matmul} "
        f"build_route={args.build_route or '-'} -> {'BUILD_OK' if plan else 'BUILD_FAILED'}"
    )
    for error in logger.errors[:3]:
        print("  ERROR:", " | ".join(error.split("\n"))[:600])
    return 0 if plan else 1


if __name__ == "__main__":
    raise SystemExit(main())

Observed behavior

TensorRT 11.4.0.106 fails the build. Bisection on 11.4.0.106:

Variant	Result
IAttention → rows → FP8 Q/DQ → FP8 linear (script default)	fails
Same without the linear (FP8 Q/DQ only)	fails
FLUX.2 shape [1, 48, 4608, 128]: FP8 QKV linear → IAttention → FP8 output linear	fails
Same, FP8 linear only before IAttention	builds
Same, FP8 linear only after IAttention	fails
IAttention.decomposable = False	fails
Sequence length 256, 1024 or 4608; FP16 or BF16; builder optimization level 1 or 3	fails
build_route = "-quantize_ppg=off" (script default and FLUX.2 shape, FP16 and BF16)	builds
-quantize_ppg:fc_fusion=off, -mha_fusion:v2=off, -mha_fusion:generate_attention_op=off, -reshape_ppg=off or -transpose_ppg=off	fails

TensorRT 11.1.0.106 builds every failing variant above.

Expected behavior

The network builds, as it does with TensorRT 11.1.0.106. Quantize propagation should not move an FP8 quantize into IAttention's softmax input when no implementation accepts it; the quantize can stay after the attention output.

Relevant logs

tensorrt=11.4.0.106 shape=(1, 2, 256, 64) matmul=True build_route=- -> BUILD_FAILED
  ERROR: Error Code: 9: Skipping tactic 0x0000000000000000 due to exception [type.cpp:186: infer_type] Could not infer output types for operation:   262: softmax: softmax_out__mye260-(f8[2,256,256][]so[], mem_prop=0, align=1) | bmm1_out__mye248-(f8[2,256,256][]so[], mem_prop=0, align=1), stream = 0 // __soft263, No matching rules found for input operand types; [myelin_graph.h:1295: attachExceptionMsgToGraph] MyelinCheckException: operation.cpp:1570: CHECK_OK(infer_shape_t::canonicalize_and_verify(op)) failed.
  ERROR: IBuilder::buildSerializedNetwork: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[(Unnamed Layer* 1) [AttentionInput]...(Unnamed Layer* 10) [Matrix Multiply]]}.)

The FLUX.2 DiT build reports the same softmax: softmax_out (f8[48,4608,4608]) | bmm1_out (f8[48,4608,4608]) error.

Impact and next steps

• flux-2-dev-fp8 cannot be built with TensorRT 11.4.
• build_route = "-quantize_ppg=off" is a candidate workaround. It turns off quantize propagation for the whole network, and its effect on FP8 GEMM fusion, accuracy and performance of the full FLUX.2 DiT has not been measured.
• This needs a TensorRT fix; the reproducer above is ready to route to the Myelin team.

Submission checks

• I searched open and closed issues and found no duplicate.
• This is not a security vulnerability.
• I removed secrets, credentials, private URLs, and restricted artifacts.