EmbLayerNormVarSeqlenPlugin: dangling pointer and out-of-bounds read during serialization

Il bug si trova in getFieldsToSerialize() in plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp, dove il campo output_fp16 punta a un bool locale. Aggiungi un membro int32_t persistente alla classe nel file .h corrispondente e fai in modo che il campo punti a esso, seguendo lo schema int32_t mOutputFp16 di embLayerNormPlugin.cpp. Il lavoro è completato quando il valore serializzato è 00 00 00 00 per FP32 e 01 00 00 00 per FP16, e i passaggi build, inspect e run dello script di riproduzione non si interrompono più.

Description

EmbLayerNormVarSeqlenPluginBase::getFieldsToSerialize() (plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp, used by CustomEmbLayerNormPluginDynamic versions 4 (HFace) and 5 (MTron)) serializes the output_fp16 field incorrectly:

bool output_fp16 = mType == DataType::kHALF;
mDataToSerialize.emplace_back("output_fp16", &output_fp16, PluginFieldType::kINT32, 1);

There are two separate defects:

1. Dangling pointer: output_fp16 is a local variable, but TensorRT reads fields[i].data only after getFieldsToSerialize() has returned. By then the variable no longer exists.
2. Out-of-bounds read: the field is declared as kINT32 with length 1, so TensorRT copies 4 bytes from a 1-byte bool. The creator reads it back with static_cast<int32_t const*>(data)[0] != 0, so any nonzero garbage in those 4 bytes decodes as true.

ASan (detect_stack_use_after_return=1, plugin built from source with -fsanitize=address, real TensorRT 11.4 engine build) confirms both defects:

ERROR: AddressSanitizer: stack-use-after-return on address 0x7d70ca804030
READ of size 4 at 0x7d70ca804030 thread T0
    #0 in memcpy
    #1 (libnvinfer.so.11+0x2c121de)
    ...
Address 0x7d70ca804030 is located in stack of thread T0 at offset 48 in frame
    #0 in nvinfer1::plugin::bert::EmbLayerNormVarSeqlenPluginBase::getFieldsToSerialize()
    [48, 49) 'output_fp16' (line 286) <== Memory access at offset 48 partially overflows this variable

The serialized engine contains the plugin fields more than once, and every copy is written from the same dangling pointer. The script below prints two of them (official TensorRT 11.4 wheel, unfixed plugin).

The trailing field record ([int32 namelen][name\0][int32 type][int32 length][data], printed by build; same bytes for v4 and v5). This copy is not the one used at deserialization:

FP32 expected: 00 00 00 00
FP32 actual:   00 01 00 00   (int32 = 256)

FP16 expected: 01 00 00 00
FP16 actual:   01 01 00 00   (int32 = 257)

The value TensorRT actually passes to createPlugin(kRUNTIME) when the engine is deserialized (printed by inspect). This is the copy the creator decodes:

FP32 v5: c9 fd 7f 00   (int32 = 8388041 -> decoded as FP16)
FP32 v4: cd fc 7f 00   (int32 = 8387789 -> decoded as FP16)
FP16 v5: 55 fe 7f 00   (int32 = 8388181 -> decoded as FP16, correct only by accident)
FP16 v4: e7 fd 7f 00   (int32 = 8388071 -> decoded as FP16, correct only by accident)

Expected: 00 00 00 00 for FP32 and 01 00 00 00 for FP16. The garbage is stable within one engine file but changes between builds. As a result, a plugin built as FP32 is recreated as FP16 after deserialization, and the first inference aborts the process:

[TRT] [F] Assertion failed: outputs[0].type == mType
/_src/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp:354   (v5; :325 for v4)
Aborting...

Scope: The serialization defect affects both plugin versions. End-to-end inference recovery was verified with MTron v5. HFace v4 also exhibits a separate output-dimension assertion, which is outside the scope of this issue.

The non-varseqlen plugin (embLayerNormPlugin.cpp) does this correctly by pointing at int32_t members (mOutputFp16, mUseFullMask), as do the other V3 plugins that serialize flags as kINT32.

Suggested fix: add an int32_t mOutputFp16{} member, set it inside getFieldsToSerialize(), and point the field at it. The proposed fix preserves the serialized field schema and correctly serializes newly built engines without changing the expected field representation. Previously serialized FP32 engines with a corrupted output_fp16 field may need to be rebuilt.

I tested this with a source build against TensorRT 11.4: both the trailing field record and the value passed to createPlugin(kRUNTIME) become exactly 00 00 00 00 / 01 00 00 00 for v4 and v5, MTron v5 FP32 and FP16 engines deserialize and match a NumPy reference, and ASan reports no errors during the build or the run.

--- a/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.h
+++ b/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.h
     DataType mType{};
     DataType mMaskType{};
+    int32_t mOutputFp16{}; // backing storage for the serialized 'output_fp16' field
--- a/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp
+++ b/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp
-    bool output_fp16 = mType == DataType::kHALF;
-    mDataToSerialize.emplace_back("output_fp16", &output_fp16, PluginFieldType::kINT32, 1);
+    mOutputFp16 = mType == DataType::kHALF ? 1 : 0;
+    mDataToSerialize.emplace_back("output_fp16", &mOutputFp16, PluginFieldType::kINT32, 1);
 ...
-    if (output_fp16)
+    if (mOutputFp16 != 0)

I'm happy to open a PR with this fix and a regression test if that would help.

Environment

TensorRT Version: 11.4.0.106 (tensorrt-cu13 pip wheel; OSS main at 16ed0636 "TensorRT 11.4 OSS Release")

NVIDIA GPU: GeForce RTX 5070 Ti (SM 12.0)

NVIDIA Driver Version: 580.178.04

CUDA Version: 13.0.88

CUDNN Version: N/A (not used by this plugin)

Operating System: Ubuntu 24.04.4 LTS (kernel 7.0.0-34-generic), GCC 13.3

Python Version (if applicable): 3.12.13

Tensorflow Version (if applicable): N/A

PyTorch Version (if applicable): N/A

Baremetal or Container (if so, version): Baremetal

Relevant Files

Model link: N/A. The repro builds a one-layer network with the plugin and random weights (script below).

Steps To Reproduce

Commands or scripts:

pip install tensorrt-cu13==11.4.0.106 numpy
python repro_output_fp16.py build 5 fp32_v5.engine   # MTron (v5)
python repro_output_fp16.py inspect fp32_v5.engine   # fresh process: value passed to createPlugin(kRUNTIME)
python repro_output_fp16.py run fp32_v5.engine       # fresh process: deserialize + run
# same for HFace: build 4 fp32_v4.engine / inspect fp32_v4.engine / run fp32_v4.engine

Output (TensorRT default-stream warning omitted):

$ python repro_output_fp16.py build 5 fp32_v5.engine
build: output dtype=DataType.FLOAT; output_fp16 in trailing field record: type=5 length=1 bytes=00 01 00 00 -> int32=256 (expected 0)

$ python repro_output_fp16.py inspect fp32_v5.engine
inspect: output_fp16 passed to createPlugin(TensorRTPhase.RUNTIME): type=5 length=1 bytes=c9 fd 7f 00 -> int32=8388041

$ python repro_output_fp16.py run fp32_v5.engine
[TRT] [F] Assertion failed: outputs[0].type == mType
/_src/plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp:354
Aborting...
$ echo $?
1

v4 gives the same trailing-record bytes, passes cd fc 7f 00 to createPlugin(kRUNTIME), and aborts at :325. The FP16 values above were obtained with the same script after changing FLAG to np.array([1], dtype=np.int32): build prints bytes=01 01 00 00 -> int32=257 (expected 1) and inspect prints bytes=55 fe 7f 00 -> int32=8388181.

"""CustomEmbLayerNormPluginDynamic (IPluginV3, v4 HFace / v5 MTron): output_fp16 serialization bug.

  python repro_output_fp16.py build 5 fp32_v5.engine   # build FP32 engine, print the output_fp16 field record bytes
  python repro_output_fp16.py inspect fp32_v5.engine   # print the output_fp16 value passed to createPlugin(kRUNTIME)
  python repro_output_fp16.py run fp32_v5.engine       # deserialize in a fresh process and run
"""
import ctypes
import struct
import sys

import numpy as np
import tensorrt as trt

H, WORD_V, POS_V, TOK_V = 128, 64, 128, 2
rng = np.random.default_rng(0)
BETA = rng.standard_normal(H).astype(np.float32)
GAMMA = rng.standard_normal(H).astype(np.float32)
WORD = rng.standard_normal(WORD_V * H).astype(np.float32)
POS = rng.standard_normal(POS_V * H).astype(np.float32)
TOK = rng.standard_normal(TOK_V * H).astype(np.float32)
FLAG = np.array([0], dtype=np.int32)  # output_fp16 = 0 -> FP32 plugin

LOGGER = trt.Logger(trt.Logger.WARNING)
trt.init_libnvinfer_plugins(LOGGER, "")
NAME = "CustomEmbLayerNormPluginDynamic"


def build(version, path):
    F, T = trt.PluginField, trt.PluginFieldType
    fc = trt.PluginFieldCollection([
        F("output_fp16", FLAG, T.INT32),
        F("bert_embeddings_layernorm_beta", BETA, T.FLOAT32),
        F("bert_embeddings_layernorm_gamma", GAMMA, T.FLOAT32),
        F("bert_embeddings_word_embeddings", WORD, T.FLOAT32),
        F("bert_embeddings_token_type_embeddings", TOK, T.FLOAT32),
        F("bert_embeddings_position_embeddings", POS, T.FLOAT32),
    ])
    creator = trt.get_plugin_registry().get_creator(NAME, version, "")
    assert creator is not None, f"{NAME} version {version} not registered"
    plugin = creator.create_plugin("emb", fc, trt.TensorRTPhase.BUILD)
    assert plugin is not None, "create_plugin failed"

    builder = trt.Builder(LOGGER)
    net = builder.create_network(0)
    ins = [net.add_input(n, trt.int32, (-1,)) for n in ("input_ids", "segment_ids", "cu_seqlens", "max_seqlen")]
    out = net.add_plugin_v3(ins, [], plugin).get_output(0)
    out.name = "embeddings"
    net.mark_output(out)

    cfg = builder.create_builder_config()
    prof = builder.create_optimization_profile()
    shapes = {
        "input_ids": (1, 12, 64),
        "segment_ids": (1, 12, 64),
        "cu_seqlens": (2, 3, 9),
        "max_seqlen": (1, 7, 128),
    }
    for n, (lo, opt, hi) in shapes.items():
        prof.set_shape(n, (lo,), (opt,), (hi,))
    cfg.add_optimization_profile(prof)

    blob = builder.build_serialized_network(net, cfg)
    assert blob is not None, "engine build failed"
    data = bytes(blob)
    with open(path, "wb") as f:
        f.write(data)

    # The engine contains the plugin fields more than once. This is the trailing field record
    # [int32 namelen][name\0][int32 type][int32 length][data]; it is NOT the copy TensorRT passes to
    # createPlugin(kRUNTIME) at deserialization (see `inspect`).
    field_name = b"output_fp16\x00"
    offset = data.rfind(field_name)
    assert offset != -1, "output_fp16 field not found in serialized engine"
    i = offset + len(field_name)
    ftype, length = struct.unpack_from("<ii", data, i)
    assert length == 1, f"unexpected output_fp16 length {length}"
    raw = data[i + 8:i + 12]
    print(f"build: output dtype={out.dtype}; output_fp16 in trailing field record: type={ftype} length={length} "
          f"bytes={raw.hex(' ')} -> int32={struct.unpack('<i', raw)[0]} (expected {int(FLAG[0])})")


class SpyCreator(trt.IPluginCreatorV3One):
    """Stands in for the plugin creator at runtime and prints the fields TensorRT passes to createPlugin."""

    def __init__(self, version):
        trt.IPluginCreatorV3One.__init__(self)
        self.name = NAME
        self.plugin_namespace = ""
        self.plugin_version = version
        self.field_names = trt.PluginFieldCollection([])

    def create_plugin(self, name, fc, phase):
        for f in fc:
            if f.name == "output_fp16":
                raw = f.data.tobytes()
                print(f"inspect: output_fp16 passed to createPlugin({phase}): type={int(f.type)} length={f.size} "
                      f"bytes={raw.hex(' ')} -> int32={struct.unpack('<i', raw)[0]}", flush=True)
        return None  # stop here: the engine is only inspected, never run


def inspect(path):
    registry = trt.get_plugin_registry()
    spies = [SpyCreator(v) for v in ("4", "5")]
    for spy in spies:
        assert registry.deregister_creator(registry.get_creator(NAME, spy.plugin_version, ""))
        assert registry.register_creator(spy, "")
    quiet = trt.Logger(trt.Logger.INTERNAL_ERROR)  # the spy returning None makes TensorRT log an error
    with open(path, "rb") as f:
        eng = trt.Runtime(quiet).deserialize_cuda_engine(f.read())
    assert eng is not None, "deserialize failed"
    eng.create_execution_context()  # TensorRT calls createPlugin(kRUNTIME) here


def run(path):
    with open(path, "rb") as f:
        eng = trt.Runtime(LOGGER).deserialize_cuda_engine(f.read())
    assert eng is not None, "deserialize failed"
    ctx = eng.create_execution_context()

    cudart = ctypes.CDLL("libcudart.so.13")
    cudart.cudaMalloc.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]
    cudart.cudaMemcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]

    def dev(nbytes, host=None):
        p = ctypes.c_void_p()
        assert cudart.cudaMalloc(ctypes.byref(p), nbytes) == 0
        if host is not None:
            assert cudart.cudaMemcpy(p, host.ctypes.data, host.nbytes, 1) == 0
        return p.value

    inputs = {
        "input_ids": np.arange(12, dtype=np.int32) % WORD_V,
        "segment_ids": np.zeros(12, np.int32),
        "cu_seqlens": np.array([0, 5, 12], np.int32),
        "max_seqlen": np.zeros(7, np.int32),
    }
    for n, a in inputs.items():
        ctx.set_input_shape(n, a.shape)
        ctx.set_tensor_address(n, dev(a.nbytes, a))
    ctx.set_tensor_address("embeddings", dev(12 * H * 4))
    print("run: execute_async_v3 ->", ctx.execute_async_v3(0))


if __name__ == "__main__":
    if sys.argv[1] == "build":
        build(sys.argv[2], sys.argv[3])
    elif sys.argv[1] == "inspect":
        inspect(sys.argv[2])
    else:
        run(sys.argv[2])

Have you tried the latest release?: Yes, 11.4.0.106 is the latest release, and the code on main is unchanged. The pattern has been in the file since the 10.4 GA update.

Can this model run on other frameworks?: N/A. This is a plugin serialization bug, not a model issue.