sdpa flash attention x86 optimization, llm models enable gqa, dispatch for avx2 bfloat2float - #6923
sdpa flash attention x86 optimization, llm models enable gqa, dispatch for avx2 bfloat2float#6923nihui wants to merge 33 commits into
Conversation
|
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #6923 +/- ##
==========================================
+ Coverage 95.46% 95.81% +0.35%
==========================================
Files 960 930 -30
Lines 459398 447646 -11752
==========================================
- Hits 438578 428927 -9651
+ Misses 20820 18719 -2101 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
nihui
changed the title
|
i7-7700 4t
|
![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: bd165b91f6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b6bebe8a75
ℹ️ 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".
| ``` | ||
|
|
||
| Set the same allocator on every extractor belonging to the session. The session owns it, and it must outlive every cache `Mat`. The sequence-length hint controls the first reservation but is not a hard limit; the cache still grows if necessary. Without a hint, ncnn uses a moderate initial reservation and geometric growth. | ||
| Set the same allocator on every extractor belonging to the session. The session owns it, and it must outlive every cache `Mat`. The sequence-length hint controls the first reservation but is not a hard limit; the cache still grows if necessary. The hint may also be used without a dedicated allocator. Without a hint, ncnn uses a moderate initial reservation and geometric growth. |
There was a problem hiding this comment.
When users follow this new guidance without setting a dedicated allocator, the hint is ignored by CPU MultiHeadAttention, generic CPU SDPA, and Vulkan SDPA: their create_or_grow_kvcache implementations only call kvcache_capacity inside an if (opt.kvcache_allocator) or if (opt.kvcache_vkallocator) block and otherwise allocate exactly new_seqlen. Those paths therefore still reallocate and copy the cache on every decoding step, contrary to this paragraph; either implement allocator-free reservation on those backends or qualify the documentation as x86-optimized-SDPA-only.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 134e7812f9
ℹ️ 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.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1c461e845
ℹ️ 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.
🟡 Changes recommended
KV-cache extraction, batch validation, and AVX2 dequantization error propagation have unresolved correctness issues.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds optimized x86 SDPA/GQA execution, allocator-independent KV-cache handling, and AVX2 BF16 runtime dispatch across x86 layers.
Changes:
- Introduces optimized FP32/BF16 SDPA prefill, decode, and KV-cache paths.
- Adds AVX2 BF16 dispatch wrappers across x86 operators.
- Expands KV-cache, GQA, mask, allocator, OOM, and benchmark coverage.
File summaries
| File | Description |
|---|---|
tests/test_sdpa.cpp |
Expands SDPA mask and GQA cases. |
tests/test_sdpa_oom.cpp |
Adds KV-cache allocator OOM tests. |
tests/test_sdpa_kvcache.cpp |
Reworks multi-step cache validation. |
tests/test_sdpa_kvcache_session.cpp |
Expands extractor/session cache tests. |
tests/test_multiheadattention_kvcache.cpp |
Tests batched-cache rejection. |
tests/test_multiheadattention_kvcache_allocator.cpp |
Tests optional cache allocators. |
tests/test_multiheadattention_block_quant.cpp |
Updates output allocator expectations. |
tests/perf/perf_sdpa_prefill.cpp |
Revises prefill benchmark configurations. |
tests/perf/perf_sdpa_decode.cpp |
Revises decode benchmark configurations. |
src/net.cpp |
Preserves extracted cache storage layouts. |
src/layer/x86/unaryop_x86.cpp |
Dispatches BF16 unary operations to AVX2. |
src/layer/x86/unaryop_x86_avx2.cpp |
Adds AVX2 unary operation wrappers. |
src/layer/x86/tanh_x86_avx2.cpp |
Adds AVX2 BF16 tanh wrapper. |
src/layer/x86/tanh_bf16s.h |
Adds AVX2 tanh dispatch. |
src/layer/x86/swish_x86_avx2.cpp |
Adds AVX2 BF16 swish wrapper. |
src/layer/x86/swish_bf16s.h |
Adds AVX2 swish dispatch. |
src/layer/x86/softmax_x86_avx2.cpp |
Adds AVX2 BF16 softmax wrappers. |
src/layer/x86/softmax_bf16s.h |
Adds AVX2 softmax dispatch. |
src/layer/x86/sigmoid_x86_avx2.cpp |
Adds AVX2 BF16 sigmoid wrapper. |
src/layer/x86/sigmoid_bf16s.h |
Adds AVX2 sigmoid dispatch. |
src/layer/x86/selu_x86_avx2.cpp |
Adds AVX2 BF16 SELU wrapper. |
src/layer/x86/selu_bf16s.h |
Adds AVX2 SELU dispatch. |
src/layer/x86/sdpa_x86.cpp |
Integrates optimized x86 SDPA kernels. |
src/layer/x86/sdpa_x86_avx512bf16.cpp |
Adds AVX-512 BF16 SDPA wrappers. |
src/layer/x86/sdpa_x86_avx2.cpp |
Adds AVX2 BF16 SDPA wrappers. |
src/layer/x86/sdpa_kvcache.h |
Implements packed KV-cache growth. |
src/layer/x86/scale_x86_avx2.cpp |
Adds AVX2 BF16 scale wrappers. |
src/layer/x86/scale_bf16s.h |
Adds AVX2 scale dispatch. |
src/layer/x86/rotaryembed_x86_avx2.cpp |
Adds AVX2 rotary embedding wrapper. |
src/layer/x86/rotaryembed_bf16s.h |
Adds AVX2 rotary dispatch. |
src/layer/x86/rmsnorm_x86_avx2.cpp |
Adds AVX2 RMSNorm wrapper. |
src/layer/x86/rmsnorm_bf16s.h |
Adds AVX2 RMSNorm dispatch. |
src/layer/x86/relu_x86_avx2.cpp |
Adds AVX2 BF16 ReLU wrapper. |
src/layer/x86/relu_bf16s.h |
Adds AVX2 ReLU dispatch. |
src/layer/x86/quantize_x86_avx2.cpp |
Adds AVX2 BF16 quantization wrapper. |
src/layer/x86/quantize_bf16s.h |
Adds AVX2 quantization dispatch. |
src/layer/x86/prelu_x86_avx2.cpp |
Adds AVX2 BF16 PReLU wrappers. |
src/layer/x86/prelu_bf16s.h |
Adds AVX2 PReLU dispatch. |
src/layer/x86/pooling_x86_avx2.cpp |
Adds AVX2 BF16 pooling wrappers. |
src/layer/x86/pooling_bf16s.h |
Adds AVX2 pooling dispatch. |
src/layer/x86/multiheadattention_x86.cpp |
Rejects batched x86 KV cache. |
src/layer/x86/mish_x86_avx2.cpp |
Adds AVX2 BF16 Mish wrapper. |
src/layer/x86/mish_bf16s.h |
Adds AVX2 Mish dispatch. |
src/layer/x86/layernorm_x86_avx2.cpp |
Adds AVX2 LayerNorm wrapper. |
src/layer/x86/layernorm_bf16s.h |
Adds AVX2 LayerNorm dispatch. |
src/layer/x86/interp_x86_avx2.cpp |
Adds AVX2 BF16 interpolation wrapper. |
src/layer/x86/interp_bf16s.h |
Adds AVX2 interpolation dispatch. |
src/layer/x86/instancenorm_x86_avx2.cpp |
Adds AVX2 InstanceNorm wrappers. |
src/layer/x86/instancenorm_bf16s.h |
Adds AVX2 InstanceNorm dispatch. |
src/layer/x86/innerproduct_x86_avx2.cpp |
Adds AVX2 BF16 inner-product wrappers. |
src/layer/x86/innerproduct_gemm_bf16s.h |
Adds AVX2 inner-product GEMM dispatch. |
src/layer/x86/innerproduct_bf16s.h |
Adds AVX2 inner-product dispatch. |
src/layer/x86/hardswish_x86_avx2.cpp |
Adds AVX2 HardSwish wrapper. |
src/layer/x86/hardswish_bf16s.h |
Adds AVX2 HardSwish dispatch. |
src/layer/x86/hardsigmoid_x86_avx2.cpp |
Adds AVX2 HardSigmoid wrapper. |
src/layer/x86/hardsigmoid_bf16s.h |
Adds AVX2 HardSigmoid dispatch. |
src/layer/x86/groupnorm_x86_avx2.cpp |
Adds AVX2 GroupNorm wrapper. |
src/layer/x86/groupnorm_bf16s.h |
Adds AVX2 GroupNorm dispatch. |
src/layer/x86/gemm_x86_avx2.cpp |
Adds AVX2 BF16 GEMM wrappers. |
src/layer/x86/gemm_bf16s.h |
Adds AVX2 GEMM dispatch. |
src/layer/x86/gelu_x86_avx2.cpp |
Adds AVX2 BF16 GELU wrapper. |
src/layer/x86/gelu_bf16s.h |
Adds AVX2 GELU dispatch. |
src/layer/x86/erf_x86_avx2.cpp |
Adds AVX2 BF16 erf wrapper. |
src/layer/x86/erf_bf16s.h |
Adds AVX2 erf dispatch. |
src/layer/x86/elu_x86_avx2.cpp |
Adds AVX2 BF16 ELU wrapper. |
src/layer/x86/elu_bf16s.h |
Adds AVX2 ELU dispatch. |
src/layer/x86/eltwise_x86_avx2.cpp |
Adds AVX2 BF16 eltwise wrapper. |
src/layer/x86/eltwise_bf16s.h |
Adds AVX2 eltwise dispatch. |
src/layer/x86/dropout_x86_avx2.cpp |
Adds AVX2 BF16 dropout wrapper. |
src/layer/x86/dropout_bf16s.h |
Adds AVX2 dropout dispatch. |
src/layer/x86/dequantize_x86_avx2.cpp |
Adds AVX2 BF16 dequantization wrapper. |
src/layer/x86/dequantize_bf16s.h |
Adds AVX2 dequantization dispatch. |
src/layer/x86/deconvolution_x86_avx2.cpp |
Adds AVX2 BF16 deconvolution wrappers. |
src/layer/x86/deconvolution_packed_bf16s.h |
Adds AVX2 deconvolution dispatch. |
src/layer/x86/convolution1d_x86_avx2.cpp |
Adds AVX2 BF16 convolution1d wrappers. |
src/layer/x86/convolution1d_packed_bf16s.h |
Adds AVX2 convolution1d dispatch. |
src/layer/x86/convolution_x86_avx2.cpp |
Adds AVX2 BF16 convolution wrappers. |
src/layer/x86/convolution_packed_bf16s.h |
Adds packed convolution dispatch. |
src/layer/x86/convolution_im2col_gemm_bf16s.h |
Adds im2col GEMM dispatch. |
src/layer/x86/convolution_3x3_winograd_bf16s.h |
Adds Winograd AVX2 dispatch. |
src/layer/x86/clip_x86_avx2.cpp |
Adds AVX2 BF16 clip wrapper. |
src/layer/x86/clip_bf16s.h |
Adds AVX2 clip dispatch. |
src/layer/x86/bnll_x86_avx2.cpp |
Adds AVX2 BF16 BNLL wrapper. |
src/layer/x86/bnll_bf16s.h |
Adds AVX2 BNLL dispatch. |
src/layer/x86/binaryop_x86.cpp |
Dispatches BF16 binary operations. |
src/layer/x86/binaryop_x86_avx2.cpp |
Adds AVX2 binary operation wrapper. |
src/layer/x86/batchnorm_x86_avx2.cpp |
Adds AVX2 BatchNorm wrappers. |
src/layer/x86/batchnorm_bf16s.h |
Adds AVX2 BatchNorm dispatch. |
src/layer/vulkan/sdpa_vulkan.cpp |
Rejects batched Vulkan KV cache. |
src/layer/vulkan/multiheadattention_vulkan.cpp |
Rejects batched Vulkan MHA cache. |
src/layer/sdpa.cpp |
Revises generic cache allocation behavior. |
src/layer/riscv/multiheadattention_riscv.cpp |
Rejects batched RISC-V MHA cache. |
src/layer/multiheadattention.cpp |
Revises generic MHA cache handling. |
src/layer/mips/sdpa_mips.cpp |
Rejects batched MIPS SDPA cache. |
src/layer/mips/multiheadattention_mips.cpp |
Rejects batched MIPS MHA cache. |
src/layer/loongarch/sdpa_loongarch.cpp |
Rejects batched LoongArch SDPA cache. |
src/layer/loongarch/multiheadattention_loongarch.cpp |
Rejects batched LoongArch MHA cache. |
src/layer/arm/sdpa_arm.cpp |
Rejects batched ARM SDPA cache. |
src/layer/arm/multiheadattention_arm.cpp |
Rejects batched ARM MHA cache. |
docs/developer-guide/operators.md |
Documents finite attention masks. |
docs/developer-guide/kvcache.md |
Documents private cache layout behavior. |
benchmark/models/llm/llama3.2_1b_decoder.ncnn.param |
Removes explicit GQA expansion layers. |
benchmark/benchncnn_llm.cpp |
Uses finite causal-mask values. |
Review details
- Files reviewed: 108/112 changed files
- Comments generated: 15
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| #if NCNN_RUNTIME_CPU && NCNN_AVX2 && __AVX__ && !__AVX2__ && !__AVX512BF16__ | ||
| if (ncnn::cpu_support_x86_avx2()) | ||
| { | ||
| dequantize_forward_bf16s_avx2(bottom_blob, top_blob, scale_data, scale_data_size, bias_data, bias_data_size, opt); | ||
| return 0; |
No description provided.