Migrate cpu sdpa from NS flash attn - #2092

Open

jijiaz wants to merge 64 commits into

jijiaz:copilot/migrate-cpu-flash-attention

jijiaz commented Jul 27, 2026 •

edited

Loading

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.
Dispatch Routes – 5 internal routes: Tier 0 scalar fallback (always active), Tier 1 mixed fp16 (AVX2) / mixed bf16 (AVX512F or AMX-BF16) via BestLA, Tier 2 homogeneous fp16 (AVX512-FP16) / homogeneous bf16 (AMX-BF16, no-GQA only). Routes 1/2 are promoted to default.
Packed KV Cache – Transparent per-forward KV reorder cache for decode workloads in the public sdpa() path. Explicit lifecycle helpers (packed_kv_alloc, update_packed_kv, bestla_sdpa_packed) under auto_round_kernel.internal.cpu for power users.
Internal Features – BestLA-only extensions (padding-right, ALiBi slopes, TANH30 activation, prefer_fp32) exposed through auto_round_kernel.internal.cpu for opt-in use; rejected by the public sdpa() contract.
Tests – 5 new Python test/benchmark files covering public API semantics, mixed-dtype correctness, internal feature parity, packed KV lifecycle, and performance benchmarks. ISA-aware skip guards for AVX2/AVX512F/AMX-BF16/AVX512-FP16.
C++ Unit Tests – TestReorderKV, TestPersistentPackedKV, TestPackedForwardSetup, TestHomogeneousForwardSetup, TestMixedPaddingRight, TestMixedAlibiTanh, TestMixedNumericalFeatures.
CI – New GitHub Actions workflow (ark_cpu_sdpa.yml) with AVX2 (ubuntu-latest), AVX512F, and AMX-BF16 (self-hosted) runner tiers.

TODO / Follow-up

Debug route cleanup – remove debug_resolve_sdpa_route and debug_cpu_sdpa_route once route stability is confirmed in CI; tests should assert numerical parity rather than route integers.

Type of Change

Feature

Checklist Before Submitting

My code has been tested locally.
Documentation has been updated as needed.
New or updated tests are included where applicable.
The CUDA CI has passed. You can trigger it by commenting /azp run Unit-Test-CUDA-AutoRound.

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)

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

jijiaz and others added 28 commits

August 5, 2026 18:52


          added avx512 stable branch & amx-bf16 branch to mixed fp16 route

577a81e

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


          added workflow tests

10ae00c

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


          added workflow tests

110eb2c

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


          fixed internal feature router

f28f32d

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


          fixed internal feature router on mixed fp16

12bf444

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


          fixed internal feature avoiding wrong router

22b065a

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


          WIP: added CI runner debuger

771da18

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


          blocked unsupported internal feature kernel

715c4ce

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


          blocked unsupported internal feature kernel

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


          cleaned up out-of-scope codes

183fa37

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


          removed excessive redefinition of TENSOR_LAYOUT

5c27cd2

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


          fix: skip cpu-related tests when cpuinfo is missing

040ce05

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


          fix: skip cpu-related tests when cpuinfo is missing

a54ccd6

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


          fix: skip cpu-related tests when cpuinfo is missing

c32ad82

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


          fix: harden CPU sdpa test imports and document scalar reference

52bb3e2

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


          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

View reviewed changes

auto_round_extension/ark/test/test_ark_cpu_mixed_bestla_sdpa.py

+              @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

a32543254 Aug 7, 2026

Contributor

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

auto_round_extension/ark/test/test_ark_cpu_mixed_bestla_sdpa.py

+                  "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)}

a32543254 Aug 7, 2026

Contributor

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

auto_round_extension/ark/.gitignore

-              *.pyc

                
                    No newline at end of file
+              *.pyc
+              *.csv.venv/
+              auto_round_extension/ark/auto_round_kernel/build_*/

a32543254 Aug 7, 2026

Contributor

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

I believe the above already included what you add

auto_round_extension/ark/README.md

               | [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.

a32543254 Aug 7, 2026

Contributor

“ 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 ？

.gitignore

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

a32543254 Aug 7, 2026

Contributor

same as above

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet