Repository navigation
Validate runtime blob shapes in GRU, RNN, LSTM, DeformableConv2D, Gemm and MultiHeadAttention, and check int8 scale records - #7152
Conversation
…m and MultiHeadAttention, and check int8 scale records
- GRU, RNN, LSTM: validate input width against weight row width and
validate explicit hidden/cell state blob extents before forward loops.
- DeformableConv2D: validate input channel count against weights and
check offset/mask channel and spatial dimensions.
- Gemm: validate K and constant M/N/K compatibility before packing.
- MultiHeadAttention: validate q/k/v widths against qdim/kdim/vdim and
require matching k/v sequence lengths.
- GRU, RNN, LSTM, Gemm, MultiHeadAttention: return -100 from load_model
when an int8 scale blob is empty.
- Add tests/test_{gru,rnn,lstm,deformableconv2d,gemm,multiheadattention}_invalid.cpp.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: de25c0b083
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (q_blob.w != qdim || k_blob.w != kdim || v_blob.w != vdim) | ||
| return -1; | ||
|
|
||
| if (k_blob.h * k_blob.elempack != v_blob.h * v_blob.elempack) |
There was a problem hiding this comment.
When CPU cross-attention runs with kv_cache enabled, distinct q/k inputs, and nonempty cache slots (for example, five unmasked bottoms with empty current k/v at 1/2 and cached tensors at 3/4), forward() sets append_kv false and never reads the current k/v blobs. These unconditional width and length checks now return -1 before that branch, breaking cache-reuse calls that previously needed only q and the cached tensors; gate current-k/v validation on whether the cache is empty or being appended.
Useful? React with 👍 / 👎.
… in int8 scale tests
- src/layer/multiheadattention.cpp: gate live k_blob and v_blob shape checks
in check_input_shapes() on append_kv so cross-attention kv_cache reuse
with populated cached_xk/cached_xv is accepted, and validate past_xk_blob
and past_xv_blob dimensions and sequence-length agreement.
- tests/test_multiheadattention_invalid.cpp: add
test_multiheadattention_kv_cache_cross() covering cross-attention kv_cache
population, reuse with empty live k/v, and mismatched past cache rejection.
- tests/test_{gemm,gru,lstm,multiheadattention,rnn}_invalid.cpp: assert
ret == expected_ret directly in test_*_int8_scales().
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 54f33e011e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
With kv_cache enabled and caller-provided nonempty caches, these are 3-D Mats whose sequence axis is h and whose packed axis is c. This comparison instead uses h * elempack and never checks dims, w == embed_dim / num_heads, or c * elempack == num_heads; a cache with too few columns or heads therefore passes and is later indexed via channel(q).row(j)[l] (or copied into a num_heads allocation), causing out-of-bounds access, while the optimized forwards now reject otherwise valid equal-h K/V caches whose channel packings differ. Compare h directly and validate the logical head count and head width for both caches.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Updated MultiHeadAttention::check_input_shapes to validate that non-empty past_xk_blob and past_xv_blob have dims == 3, w == embed_dim / num_heads, past_xk_blob.h == past_xv_blob.h, and c * elempack == num_heads, and added test_multiheadattention_kv_cache_shape coverage in tests/test_multiheadattention_invalid.cpp.
…apes - src/layer/multiheadattention.cpp: require non-empty past_xk_blob and past_xv_blob to have dims == 3, w == embed_dim / num_heads, past_xk_blob.h == past_xv_blob.h, and c * elempack == num_heads. - tests/test_multiheadattention_invalid.cpp: add test_multiheadattention_kv_cache_shape() covering valid 3D past_k/past_v and mismatched empty, h, w, c, and dims.
Summary
These six layers index their weights and recurrent state with sizes derived from the loaded parameters, but use the extents of their runtime blobs as given. When a model feeds a blob whose shape disagrees with the weights (for example from a
MemoryDataorReshapeproducer),forwardreads outside the blob or weight allocation, and for the recurrent state it also writes outside it. Separately, several int8 scale records are loaded without an.empty()check and then dereferenced.GRU,RNN,LSTM:forwardrejects an input whose width differs from the weight row width (weight_data_size / num_directionsdivided bynum_outputfor RNN, bynum_output * 3for GRU and byhidden_size * 4for LSTM), and explicit hidden / cell state blobs that are narrower thannum_output/hidden_sizeor have fewer rows than directions.hidden_state[q]is written for everyq < num_output, so a narrower state blob is overwritten past its end.DeformableConv2D: the input channel count must equal the weights' input channels, the offset blob needs at least2 * kernel_w * kernel_hchannels ofout_h x out_w, and the mask blob at leastkernel_w * kernel_hchannels ofout_h x out_w.Gemm: K of A must not exceed K of B, and must equalconstantKwhen A or B is constant. When a constant C is broadcast over M or N, the runtime M / N must equalconstantM/constantN. A runtime B with more K than A is still accepted, because only the first K elements are read andtest_sdpauses that shape.MultiHeadAttention: the q width must equalqdim; whenappend_kvis enabled, the k and v widths must equalkdimandvdimand k and v must have the same sequence length; and when pastkv_cacheblobs are supplied, they must be 3D with widthembed_dim / num_heads, channel countnum_heads(c * elempack == num_heads) and matching sequence lengths.load_modelofGRU,RNN,LSTM,GemmandMultiHeadAttentionreturns -100 when a scale record is missing, instead of continuing with an emptyMat(or calling[0]on it for theGemmB scale and theMultiHeadAttentionoutput scale).The shape checks live in a protected
check_input_shapes()helper on each base layer. It is called at the top offorward()in the base implementation and in the x86, arm, loongarch, mips and riscv overrides, underNCNN_VALIDATION, and returns -1. The Vulkan implementations are not changed.This touches no file changed by #7148 and applies on top of it.
Tests
tests/test_{gru,rnn,lstm,deformableconv2d,gemm,multiheadattention}_invalid.cpp. Each runs the naive and thecreate_layer_cpulayer on a valid control and on mismatched shapes, and the int8 cases load a model with and without its scale records.forwardreturns 0 where -1 is expected, or AddressSanitizer reports a heap-buffer-overflow.ctest195/195 (189 existing plus the 6 new tests) in both. No AddressSanitizer report in the ASan run.test_gru,test_rnn,test_lstm,test_deformableconv2d,test_gemmandtest_multiheadattention. The one exception istest_rnnandtest_lstmin the ASan build, where*_int8_with_hiddenreports a value mismatch with bf16 storage; it fails the same way without this change.benchmark/load identically with and without this change. The 35 vision models, includingvision_transformer(which usesGemmandMultiHeadAttention), produce byte-identical outputs with zero weights and with patterned weights.Note: This patch was prepared with AI assistance and verified by human review and local test execution.