[bug] MultiHeadAttention::load_param divides by zero when num_heads is 0, crashing Net::load_param on any param file that sets 1=0

[WhosReadyToFlyOnaZipLine]

MultiHeadAttention::load_param divides by zero when num_heads is 0, crashing Net::load_param on any param file that sets 1=0

ncnn_attachment.zip

error log

Floating point exception (core dumped)

Under a debugger the faulting instruction is the integer division in

src/layer/multiheadattention.cpp:59
scale = pd.get(6, 1.f / sqrtf(embed_dim / num_heads));

context

• ncnn at c189d88, still present on current master
• Built exactly as your own x86-64 CI builds it, .github/workflows/linux-x64-cpu-clang.yml:
  -DNCNN_AVX=OFF -DNCNN_AVX2=OFF, Release, clang
• No sanitizer and no coverage instrumentation in the build that crashes
• Linux x86-64

how to reproduce

1. Save these four lines as crash.param:

7767517
2 2
Input                    in    0 1 in
MultiHeadAttention       mha   1 1 in out 0=64 1=0

2. Load it:

#include "net.h"
int main(int argc, char** argv)
{
    ncnn::Net net;
    net.load_param(argv[1]);       // crash.param
    return 0;
}

3. Floating point exception. Deterministic: 5 runs out of 5 on a plain Release build.

more

Root cause. MultiHeadAttention::load_param reads both operands from the param file
and checks neither:

int MultiHeadAttention::load_param(const ParamDict& pd)
{
    embed_dim = pd.get(0, 0);
    num_heads = pd.get(1, 1);
    ...
    scale = pd.get(6, 1.f / sqrtf(embed_dim / num_heads));   // line 59

embed_dim and num_heads are both int, so embed_dim / num_heads is an integer
division, and num_heads == 0 raises SIGFPE before sqrtf is ever called.

The default of 1 does not protect it. pd.get(1, 1) returns 1 only when key 1 is
absent. A param file that sets 1=0 supplies zero, and zero is what is used. Omitting
1= entirely does not crash - I checked, it loads cleanly.

Supplying scale does not protect it either. This is the part worth noting: the
expression 1.f / sqrtf(embed_dim / num_heads) is a C++ default argument, so it is
evaluated unconditionally at the call site, whether or not key 6 is present. I confirmed
this - ... 0=64 1=0 6=1.250000e-01 crashes identically.

Suggested fix, rejecting the value where it is read rather than at the division:

num_heads = pd.get(1, 1);
if (num_heads <= 0)
{
    NCNN_LOGE("MultiHeadAttention invalid num_heads %d", num_heads);
    return -1;
}

A negative num_heads is worth rejecting in the same check; it does not trap, but it
produces a negative head dimension that nothing downstream expects.

Same pattern, different layer. #6911 reports the same shape of defect in
ConvolutionDepthWise::load_param — an integer divide by a param value that is not
checked. This one is in MultiHeadAttention and is not covered by that report or by
#6913, #6914 or #6915, none of which touch multiheadattention.cpp. It may be worth
sweeping the other load_param implementations for the same construct while you are in
there.

Found by fuzzing the param reader. Happy to send more detail if it is useful.

[dosubot]

Hi @WhosReadyToFlyOnaZipLine! I'm Dosu and I'm helping the 腾讯开源's Org team.

Your analysis is correct and well-documented. The crash is confirmed in the current code — load_param reads num_heads on line 54 and immediately uses it as an integer divisor in the default argument expression on line 59, with no guard in between.

The weight_block_quantize path does have a num_heads <= 0 check at line 82, but that branch is only reached after the division has already faulted. The Vulkan backend's load_param simply delegates to the base class, so it inherits the same vulnerability.

Your suggested fix — checking num_heads <= 0 immediately after pd.get(1, 1) and returning -1 — is the right approach. It matches the pattern already used in the weight_block_quantize branch and aligns with how issue #6911 (ConvolutionDepthWise divide-by-zero) was addressed via PR #6912.

One additional note: a value of embed_dim == 0 with a valid num_heads wouldn't trap (integer division of 0 by N is 0), but embed_dim % num_heads != 0 would produce a meaningless head dimension. You may want to extend the check to catch that too, as the weight_block_quantize branch on line 82 already validates both conditions.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.