buildWtsEngine assertion when an MXFP8 block scale (E8M0) is refittable under REFIT_INDIVIDUAL #4883

"""Repro: TensorRT cannot build a REFIT_INDIVIDUAL plan in which an MXFP8 block scale (E8M0) is refittable.

One MXFP8 Linear in the exact form ModelOpt exports (TRT_MXFP8DynamicQuantize on the activation,
TRT_MXFP8DequantizeLinear on an FP8 weight initializer with a UINT8 block-scale initializer holding the E8M0
exponents), built with REFIT_INDIVIDUAL + STRIP_PLAN and different `mark_weights_refittable` sets:

    weight_refittable   FP8 weight marked refittable, scale baked                 -> builds
    scale_refittable    FP8 weight + scale marked refittable                      -> fails:
        [TRT] [E] [wtsEngineUtils.cpp::buildWtsEngine::123] Error Code 2: Internal Error
              (Assertion engine failed. In buildWtsEngine at /_src/optimizer/builder/wtsEngineUtils.cpp:123)
    scale_workaround    scale replaced through INetworkDefinition by
                        float32 constant (refittable) * float32 [1,1] input (=1.0) -> Cast(E8M0) -> DQ -> builds

Also fails without STRIP_PLAN, and when the scale only becomes E8M0 through constant folding (float32 constant ->
Cast(E8M0) -> DQ): any refittable weight whose type is E8M0 when `buildWtsEngine` runs asserts. The main engine build
succeeds; nothing is logged before the assertion, even at VERBOSE. The workaround costs a scale re-layout per GEMM at
run time (~6% on pi0.5). Use case: one stripped plan, refit per fine-tuned checkpoint with that checkpoint's own scales.

Seen on: TensorRT 10.16.1.11, Jetson Thor (sm110), Python API.    python trt_e8m0_refit_repro.py [--verbose]
"""

import argparse

import numpy as np
import tensorrt as trt
from onnx import TensorProto, helper

TOKENS, IN_FEATURES, OUT_FEATURES, BLOCK_SIZE = 64, 256, 128, 32
WEIGHT_NAME, SCALE_NAME, SCALE_FLOAT_NAME, ONE_INPUT_NAME = "w_fp8", "w_scale", "w_scale_f32", "one"
WEIGHT_DQ_LAYER_NAME = "node_of_w_dq"  # the parser names layers node_of_<first output>
DQ_SCALE_INPUT_INDEX = 1
E8M0_BIAS = 127
OPSET = 20
CASES = {"weight_refittable": False, "scale_refittable": True, "scale_workaround": None}


def make_onnx() -> bytes:
    rng = np.random.default_rng(0)
    weight = rng.integers(0, 0x7E, (OUT_FEATURES, IN_FEATURES), dtype=np.uint8)  # any E4M3 byte but NaN
    scale = rng.integers(E8M0_BIAS - 4, E8M0_BIAS + 4, (OUT_FEATURES, IN_FEATURES // BLOCK_SIZE), dtype=np.uint8)
    initializers = [
        helper.make_tensor(WEIGHT_NAME, TensorProto.FLOAT8E4M3FN, weight.shape, weight.tobytes(), raw=True),
        helper.make_tensor(SCALE_NAME, TensorProto.UINT8, scale.shape, scale.tobytes(), raw=True),
    ]
    mxfp8 = dict(domain="trt", axis=-1, block_size=BLOCK_SIZE)
    nodes = [
        helper.make_node("TRT_MXFP8DynamicQuantize", ["x"], ["x_fp8", "x_scale"], output_dtype=TensorProto.FLOAT8E4M3FN, **mxfp8),
        helper.make_node("TRT_MXFP8DequantizeLinear", ["x_fp8", "x_scale"], ["x_dq"], output_dtype=TensorProto.BFLOAT16, **mxfp8),
        helper.make_node("TRT_MXFP8DequantizeLinear", [WEIGHT_NAME, SCALE_NAME], ["w_dq"], output_dtype=TensorProto.BFLOAT16, **mxfp8),
        helper.make_node("Transpose", ["w_dq"], ["w_t"], perm=[1, 0]),
        helper.make_node("MatMul", ["x_dq", "w_t"], ["y"]),
    ]
    graph = helper.make_graph(
        nodes, "mxfp8_linear",
        [helper.make_tensor_value_info("x", TensorProto.BFLOAT16, [TOKENS, IN_FEATURES])],
        [helper.make_tensor_value_info("y", TensorProto.BFLOAT16, [TOKENS, OUT_FEATURES])],
        initializer=initializers,
    )
    model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET), helper.make_opsetid("trt", 1)])
    return model.SerializeToString()


def apply_workaround(network: trt.INetworkDefinition) -> None:
    """float32 constant (refittable under SCALE_FLOAT_NAME) * runtime 1.0 -> Cast(E8M0), fed to the weight DQ."""
    dq_layer = next(network.get_layer(i) for i in range(network.num_layers) if network.get_layer(i).name == WEIGHT_DQ_LAYER_NAME)
    shape = tuple(dq_layer.get_input(DQ_SCALE_INPUT_INDEX).shape)
    constant = network.add_constant(shape, trt.Weights(np.ones(shape, dtype=np.float32)))
    network.set_weights_name(constant.weights, SCALE_FLOAT_NAME)
    one = network.add_input(ONE_INPUT_NAME, trt.DataType.FLOAT, (1,) * len(shape))
    product = network.add_elementwise(constant.get_output(0), one, trt.ElementWiseOperation.PROD)
    cast = network.add_cast(product.get_output(0), trt.DataType.E8M0)
    dq_layer.set_input(DQ_SCALE_INPUT_INDEX, cast.get_output(0))


def build(case: str, onnx_bytes: bytes, trt_logger: trt.Logger) -> str:
    builder = trt.Builder(trt_logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))
    if not trt.OnnxParser(network, trt_logger).parse(onnx_bytes):
        return "PARSE FAILED"
    scale_refittable = CASES[case]
    refittable = [WEIGHT_NAME]
    if scale_refittable is None:
        apply_workaround(network)
        refittable.append(SCALE_FLOAT_NAME)
    elif scale_refittable:
        refittable.append(SCALE_NAME)
    for name in refittable:
        if not network.mark_weights_refittable(name):
            return f"mark_weights_refittable({name!r}) FAILED"

    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.REFIT_INDIVIDUAL)
    config.set_flag(trt.BuilderFlag.STRIP_PLAN)
    profile = builder.create_optimization_profile()
    for i in range(network.num_inputs):
        shape = list(network.get_input(i).shape)
        profile.set_shape(network.get_input(i).name, min=shape, opt=shape, max=shape)
    config.add_optimization_profile(profile)
    return "OK" if builder.build_serialized_network(network, config) is not None else "BUILD FAILED"


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--verbose", action="store_true", help="full TensorRT builder log (trt.Logger.VERBOSE)")
    args = parser.parse_args()
    trt_logger = trt.Logger(trt.Logger.VERBOSE if args.verbose else trt.Logger.WARNING)
    print(f"TensorRT {trt.__version__}")
    onnx_bytes = make_onnx()
    results = {}
    for case in CASES:
        print(f"=== {case}: refittable = {[WEIGHT_NAME] + {False: [], True: [SCALE_NAME], None: [SCALE_FLOAT_NAME]}[CASES[case]]}")
        results[case] = build(case, onnx_bytes, trt_logger)
        print(f"=== {case}: {results[case]}")
    print("\nsummary:", results)