Migrate cpu sdpa from NS flash attn by jijiaz · Pull Request #2092 · intel/auto-round

jijiaz · 2026-07-27T08:01:29Z

Description

Migrate CPU SDPA backend from Neural Speed to auto_round_kernel, providing a drop-in CPU
implementation with the same public API as the existing XPU sdpa(). The implementation
leverages BestLA kernels for mixed-precision (f32 Q + f16/bf16 KV) and homogeneous
(fp16/bf16) dispatch, with a scalar fallback for unsupported configurations.

Public API — auto_round_kernel.sdpa() now accepts device="cpu" tensors. Supports
HND/NHD layouts, causal masking, additive masks, GQA, and multi-tile K/V sequences.
Mixed-dtype (f32 Q + f16/bf16 KV) is transparently accelerated via a hidden packed KV cache.
Tests — test_ark_cpu_sdpa.py (public API semantics), test_ark_cpu_mixed_bestla_sdpa.py
(mixed-dtype accuracy with real LLM shapes), bench_ark_cpu_sdpa.py (correctness benchmark),
benchmark_sdpa.py (perf micro-benchmark). ISA-aware skip guards for AVX2/AVX512F/
AMX-BF16/AVX512-FP16.
Numerical Tolerances — Aligned with PyTorch CPU SDPA and in-repo XPU attention:
fp16 → 1–2e-2, bf16 → 3–5e-2 (empirically measured; softmax amplifies BF16's 7-bit
mantissa error ~3–5× vs FP16's 10-bit).

Performance

ARK CPU SDPA delivers 1.3–1.6× speedup over Torch FP32 SDPA for decode workloads across
real-world LLM attention shapes on 32-thread Xeon with AMX-BF16. Representative results
(geomean across all non-scalar routes: 1.28×):

model	workload	B	Hq/Hkv	Skv	dtype	ARK (ms)	Torch FP32 (ms)	speedup
gemma2-27b	decode	4	32/16	8K	bf16	2.49	3.91	1.57×
qwen2.5-7b	decode	4	28/4	8K	bf16	1.35	2.10	1.55×
deepseek-V3	decode	4	128/8	8K	bf16	6.64	10.00	1.51×
llama3.1-70b	decode	4	64/8	8K	bf16	3.41	5.05	1.48×
llama3.1-8b	prefill	1	32/8	1K	bf16	2.91	3.99	1.37×
llama3.1-70b	prefill	1	64/8	512	bf16	2.40	2.76	1.15×
llama3.1-70b	prefill	1	64/8	512	fp16	3.34	2.79	0.84×

Notes:

FP16 decode speedups (1.0–1.3×) trail BF16 (1.3–1.6×) because the current build
uses gcc, which lacks AVX512-FP16 intrinsic support — the native fp16 kernel path
is compiled out, leaving only the AVX2+F16C conversion path (fp16→fp32 per element).
Building with icx would restore parity.

Prefill underperforms (0.84–1.37×) because the mixed-datatype matmul incurs
per-element dtype conversion overhead that Torch's pure-FP32 MKL SGEMM avoids.

Homogeneous routes (all-FP16 or all-BF16 Q/K/V) are excluded from the table because
hom-FP16 requires AVX512-FP16 compiler support (not available under gcc), and
hom-BF16 is restricted to non-GQA models (all benchmark shapes use GQA).

Signed-off-by: jijiaz <jijia.zhou@intel.com>

…dpa benchmark Signed-off-by: jijiaz <jijia.zhou@intel.com>

Signed-off-by: jijiaz <jijia.zhou@intel.com>

…(phase 2 step 1) Signed-off-by: jijiaz <jijia.zhou@intel.com>

…phase 2 step 2) Signed-off-by: jijiaz <jijia.zhou@intel.com>

…2 step 3) Signed-off-by: jijiaz <jijia.zhou@intel.com>

Signed-off-by: jijiaz <jijia.zhou@intel.com>

…hase 4 step 1) Signed-off-by: jijiaz <jijia.zhou@intel.com>

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

a32543254 · 2026-08-07T08:49:16Z

+@pytest.mark.parametrize("layout", ["HND", "NHD"])
+def test_bestla_mixed_sdpa_matches_torch(kv_dtype, is_causal, layout):
+    torch.manual_seed(4003)
+    batch, heads_q, heads_kv, head_dim, seq = 1, 8, 2, 64, 64


better also have medium and large size test for acc and performance
maybe take some shape from really LLMs

Added LLM shape tests and benchmarks

a32543254 · 2026-08-07T08:50:55Z

+    "auto_round_kernel", reason="compiled ARK extension not built in this environment"
+)
+
+_TOL = {torch.float16: (3e-2, 3e-2), torch.bfloat16: (8e-2, 8e-2)}


the TOL seem a little bit too high for float
maybe you could try to set a lower threshold

a32543254 · 2026-08-07T08:53:16Z

-*.pyc
+*.pyc
+*.csv.venv/
+auto_round_extension/ark/auto_round_kernel/build_*/


I don‘t think we need add more here
build
xbuild
*.csv
*.so
*.pyc

I believe the above already included what you add

Removed in new commit.

a32543254 · 2026-08-07T08:54:27Z

 | [test_matmul.py](test/test_matmul.py) | Low-level matmul |
 | [test_packq.py](test/test_packq.py) | Weight packing utilities |
+  Notes:
+  * The patch only routes calls to ARK on XPU when the inputs match ARK kernel constraints; otherwise it falls back to the original torch SDPA.


“ The patch only routes calls to ARK on XPU when the inputs match ARK kernel constraints; otherwise it falls back to the original torch SDPA.”

why here mention xpu ？ all the code should only focus on cpu right ？

Merge remains, now removed in the new commits.

a32543254 · 2026-08-07T08:54:42Z

 docs/plan/
+.venv/
+auto_round_extension/ark/auto_round_kernel/build_*/
+auto_round_extension/ark/build-*/


same as above

Removed in new commit.

Signed-off-by: jijiaz <jijia.zhou@intel.com> (cherry picked from commit 37f00a8)

Signed-off-by: jijiaz <jijia.zhou@intel.com>

Defect-1 (workspace cache) + Defect-2 (fp16-bf16-AMX prefill) already in tree; this snapshot captures the JIT-accelerated BF16 K/V update paths that were in the working tree before the GCC12 investigation. Signed-off-by: jijiaz <jijia.zhou@intel.com>

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

Copilot

Pull request overview

Adds CPU support to ARK’s SDPA API using scalar and BestLA-accelerated routes, including packed KV caching.

Changes:

Implements CPU SDPA dispatch, packed caches, and mixed/homogeneous precision routes.
Adds CPU tests, benchmarks, validation tooling, and CI.
Updates build configuration, BestLA kernels, documentation, and dependencies.

Reviewed changes

Copilot reviewed 18 out of 20 changed files in this pull request and generated 14 comments.

Show a summary per file

File	Description
`.gitignore`	Ignores ARK build artifacts.
`.github/workflows/ark_cpu_sdpa.yml`	Adds CPU SDPA CI.
`auto_round_extension/ark/.gitignore`	Normalizes ignored Python artifacts.
`auto_round_extension/ark/README.md`	Documents SDPA evaluation usage.
`auto_round_extension/ark/requirements.txt`	Adds CPU feature detection dependency.
`auto_round_extension/ark/test/validate_non_int8_cpu_sdpa.py`	Adds route readiness runbook.
`auto_round_extension/ark/test/test_ark_cpu_sdpa.py`	Tests public CPU SDPA behavior.
`auto_round_extension/ark/test/test_ark_cpu_mixed_bestla_sdpa.py`	Tests mixed-precision BestLA routes.
`auto_round_extension/ark/test/bench_ark_cpu_sdpa.py`	Benchmarks CPU SDPA routes.
`auto_round_extension/ark/auto_round_kernel/CMakeLists.txt`	Builds CPU SDPA sources and feature flags.
`auto_round_extension/ark/auto_round_kernel/__init__.py`	Exposes CPU SDPA and packed-cache APIs.
`auto_round_extension/ark/auto_round_kernel/ark.cpp`	Adds native CPU dispatch and bindings.
`auto_round_extension/ark/auto_round_kernel/ark/cpu/mha_dense.cpp`	Implements scalar CPU attention.
`auto_round_extension/ark/auto_round_kernel/ark/cpu/mha_dense.h`	Defines CPU attention contracts.
`auto_round_extension/ark/auto_round_kernel/ark/cpu/sdpa.h`	Declares BestLA SDPA routes and cache helpers.
`auto_round_extension/ark/auto_round_kernel/bestla/bestla/bestla_gemm.h`	Adds masked tail stores.
`auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp`	Removes obsolete SDPA documentation.
`auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp`	Updates unsupported-type assertion.

Suppressed comments (1)

auto_round_extension/ark/auto_round_kernel/init.py:835

This repeats the nonexistent ARK_UNSAFE_BESTLA_MIXED_SDPA gate for n_padding. Document the actual ARK_ENABLE_INTERNAL_SDPA_FEATURES=ON build requirement so users do not set an ineffective environment variable.

    - n_padding: Number of valid (non-padding) K/V positions when the K/V
      sequence is right-padded. Must be in (0, seq_kv] and mutually exclusive
      with is_causal. Only supported on the BestLA mixed-precision CPU path
      with ARK_UNSAFE_BESTLA_MIXED_SDPA=1. [CPU-only]

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Signed-off-by: jijiaz <jijia.zhou@intel.com>

for more information, see https://pre-commit.ci

Signed-off-by: jijiaz <jijia.zhou@intel.com>

jijiaz force-pushed the copilot/migrate-cpu-flash-attention branch 14 times, most recently from 41e938e to 32356f9 Compare August 3, 2026 03:37

jijiaz force-pushed the copilot/migrate-cpu-flash-attention branch 2 times, most recently from f90eab3 to d4a3493 Compare August 5, 2026 02:57

Copilot AI added 14 commits August 5, 2026 18:52

feat: add ark cpu sdpa implementation

2639e8e

Signed-off-by: jijiaz <jijia.zhou@intel.com>

test: wrap ark cpu kv update call

51ce63d

Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: tiled online softmax (flash attention) for ark cpu sdpa

442db5f

Signed-off-by: jijiaz <jijia.zhou@intel.com>

perf: cache fp32 query row in ark cpu flash-attention; add cpu-only s…

cb2612b

…dpa benchmark Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: migrate neural-speed-style cpu attention base types (phase 1)

de98d00

Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: migrate neural-speed cpu attention softmax/epilogue components …

831a65f

…(phase 2 step 1) Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: migrate neural-speed cpu attention gemm dispatch/packer layer (…

1c71889

…phase 2 step 2) Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: migrate mha_stable_interface_t stable-softmax attention (phase …

5a435b9

…2 step 3) Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: migrate bestla_fusion_attn_forward dtype dispatch (phase 2 step 4)

478449c

Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: wire bestla_fusion_attn_forward into CPU sdpa (phase 3 step 1)

6ce7575

Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: route CPU sdpa to BestLA mixed-precision path (phase 3 step 2)

e87979a

Signed-off-by: jijiaz <jijia.zhou@intel.com>

fix: make phase 3 bestla mixed sdpa route safe and non-misleading

400160d

Signed-off-by: jijiaz <jijia.zhou@intel.com>

feat: bridge raw K/V into NTILE packed cache for bestla mixed sdpa (p…

40142c2

…hase 4 step 1) Signed-off-by: jijiaz <jijia.zhou@intel.com>

test: validate and harden packed K/V reorder bridge (phase 4 step 2)

9632014

Signed-off-by: jijiaz <jijia.zhou@intel.com>

jijiaz and others added 13 commits August 5, 2026 18:52

fix: defer .so import to end of __init__.py to prevent circular import

8e6617e

Signed-off-by: jijiaz <jijia.zhou@intel.com>

rebased and merged upstream main

8a4e8e5

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

b6f366a

for more information, see https://pre-commit.ci

fixed cpu sdpa regressions

8708f9f

Signed-off-by: jijiaz <jijia.zhou@intel.com>

fixed cpu sdpa regression

57007cd

Signed-off-by: jijiaz <jijia.zhou@intel.com>

Removed out-of-scope xpu alignment

d4a3493

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

4f4c98f

for more information, see https://pre-commit.ci

Fixed import.so route bug & removed out-of-scope tests

4d94d75

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

62096e0

for more information, see https://pre-commit.ci

Improved BF16 packed cache

4f55c50

Signed-off-by: jijiaz <jijia.zhou@intel.com>

speed up KV reorder in mixed route

9f407ae

Signed-off-by: jijiaz <jijia.zhou@intel.com>

fix prefill shape corruption

b5db39f

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

b09eea6

for more information, see https://pre-commit.ci

a32543254 self-requested a review August 7, 2026 08:43

a32543254 reviewed Aug 7, 2026

View reviewed changes

jijiaz force-pushed the copilot/migrate-cpu-flash-attention branch from 37f00a8 to 3c2f775 Compare August 8, 2026 10:59

jijiaz and others added 7 commits August 8, 2026 11:03

fixed excessive workspace

3c2f775

Signed-off-by: jijiaz <jijia.zhou@intel.com> (cherry picked from commit 37f00a8)

removed excessive tests and improved prefill routes

312a19e

Signed-off-by: jijiaz <jijia.zhou@intel.com>

Added LLM shape benchmark and adjusted TOL

7a5a09d

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

88a418e

for more information, see https://pre-commit.ci

Merge branch 'main' into copilot/migrate-cpu-flash-attention

3021be3

Merge branch 'main' into copilot/migrate-cpu-flash-attention

b599218

luoyu-intel requested a balanced review from Copilot August 10, 2026 02:55

Copilot started reviewing on behalf of luoyu-intel August 10, 2026 02:56 View session

Copilot AI reviewed Aug 10, 2026

View reviewed changes

jijiaz and others added 4 commits August 10, 2026 14:57

Merge branch 'main' into copilot/migrate-cpu-flash-attention

13de405

Resolved issues by Copilot comments

cd00fab

Signed-off-by: jijiaz <jijia.zhou@intel.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

4af90b4

for more information, see https://pre-commit.ci

Resolved issues by Copilot comments

60ece37

Signed-off-by: jijiaz <jijia.zhou@intel.com>

Uh oh!

Conversation

jijiaz commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Performance

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jijiaz commented Jul 27, 2026 •

edited

Loading