Migrate cpu sdpa from NS flash attn - #2092
Open
jijiaz wants to merge 54 commits into
Open
Conversation
jijiaz
force-pushed
the
copilot/migrate-cpu-flash-attention
branch
14 times, most recently
from
August 3, 2026 03:37
41e938e to
32356f9
Compare
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
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>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
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>
…hase 4 step 3) Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
…e 5 Step 1) Signed-off-by: jijiaz <jijia.zhou@intel.com>
…Phase 5 Step 2) Signed-off-by: jijiaz <jijia.zhou@intel.com>
…U (Phase 5) Signed-off-by: jijiaz <jijia.zhou@intel.com>
… attention routes Signed-off-by: jijiaz <jijia.zhou@intel.com>
- Extend ark.cpp CPU sdpa() signature with use_alibi, use_tanh, prefer_fp32_flag, n_padding_arg after is_causal - Build attn_flags from individual flags in the mixed_bestla block; remove stale n_padding=0 comment - Reject alibi/tanh/n_padding on scalar Tier-0 path with clear error - Extend Python sdpa() with use_alibi, use_tanh, prefer_fp32, n_padding kwargs; reject all four on XPU; split lib.sdpa() call by device type - Add Python tests: prefer_fp32 smoke, padding-right, alibi, tanh (all skip gracefully when extension or ISA unavailable) - Update Phase 6 TIER 1 comment in sdpa.cpp to mark barrier (b) closed Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
- validate_non_int8_cpu_sdpa.py: add CI/readiness matrix, promotion decision (routes 1/2 remain gated; explicit blockers B1–B5 and follow-up items F1–F7), and final delivery summary - bench_ark_cpu_sdpa.py: add --mode raw|packed|both for raw-vs-packed comparison and run_case_packed() for packed KV cache path benchmarking - .github/workflows/non_int8_cpu_sdpa.yml: ISA-matrix CI workflow (avx2 on ubuntu-latest; avx512f/amx-bf16/avx512-fp16 as self-hosted manual-dispatch stubs until hardware is available) Signed-off-by: jijiaz <jijia.zhou@intel.com>
…gap-fix
Module A (gating/entry surface): already complete; gating test and all kwarg
plumbing tests existed. Added Module C homogeneous test to test_ark_cpu_sdpa.py.
Module B (mixed-route numerical parity):
- Add test_bestla_packed_sdpa_numerical_parity (packed path vs torch ref,
parametrised over kv_dtype × is_causal) to test_ark_cpu_mixed_bestla_sdpa.py.
- Add test_bestla_raw_vs_packed_output_consistency (raw path == packed path on
same inputs) to test_ark_cpu_mixed_bestla_sdpa.py.
- Add _packed_sdpa() helper that gates/restores ARK_UNSAFE_BESTLA_MIXED_SDPA
and calls ark_cpu_packed_kv_alloc / ark_cpu_update_packed_kv /
ark_cpu_bestla_sdpa_packed with the correct Python API signatures.
Module C (homogeneous route classification):
- Add test_homogeneous_half_uses_tier0_not_internal_routes (parametrised over
fp16/bf16) to test_ark_cpu_sdpa.py — asserts that homogeneous Q/K/V inputs
are handled by Tier 0 scalar with or without the env gate, and produce exact
bitwise-identical output in both cases.
- Add `import os` to test_ark_cpu_sdpa.py (needed by the new test).
Module D (benchmark):
- Fix three API call bugs in run_case_packed in bench_ark_cpu_sdpa.py:
* ark_cpu_packed_kv_alloc: dtype is keyword-only; add dtype= keyword arg.
* ark_cpu_update_packed_kv: 6th positional arg is capacity (int), not dtype.
* ark_cpu_bestla_sdpa_packed: positional args are (query, cache_k, cache_v,
seq_len_kv, capacity, num_heads_kv); remove wrong batch/heads_q/head_dim
positional args and the non-existent dtype= kwarg.
- Add NotImplementedError to except clauses (raised by the packed-path Python
wrappers when the C extension is not built).
Bugfix (auto_round_kernel/__init__.py):
- Remove three _get_cpu_lib() calls in ark_cpu_packed_kv_alloc,
ark_cpu_update_packed_kv, and ark_cpu_bestla_sdpa_packed. _get_cpu_lib is not
defined anywhere in the module; these functions should use the module-level
cpu_lib directly, consistent with all other CPU path functions in the file.
The NameError caused the packed-path tests to fail unconditionally instead of
raising NotImplementedError and being skipped.
Module E (runbook/workflow):
- validate_non_int8_cpu_sdpa.py: update TEST_COVERAGE section to document the
two new packed-path tests and the homogeneous classification test; add a
"Tier 1 packed path" entry to COMMANDS for --run mode.
- non_int8_cpu_sdpa.yml: add "Tier 1 packed KV path" step to avx2 job (F16,
continue-on-error); split avx512f and amx-bf16 jobs into separate steps for
raw mixed tests (-k "not packed") and packed path (-k "packed").
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
Signed-off-by: jijiaz <jijia.zhou@intel.com>
for more information, see https://pre-commit.ci
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 XPUsdpa(). 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.auto_round_kernel.sdpa()now acceptsdevice="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.sdpa()path. Explicit lifecycle helpers (packed_kv_alloc,update_packed_kv,bestla_sdpa_packed) underauto_round_kernel.internal.cpufor power users.prefer_fp32) exposed throughauto_round_kernel.internal.cpufor opt-in use; rejected by the publicsdpa()contract.TestReorderKV,TestPersistentPackedKV,TestPackedForwardSetup,TestHomogeneousForwardSetup,TestMixedPaddingRight,TestMixedAlibiTanh,TestMixedNumericalFeatures.ark_cpu_sdpa.yml) with AVX2 (ubuntu-latest), AVX512F, and AMX-BF16 (self-hosted) runner tiers.TODO / Follow-up
debug_resolve_sdpa_routeanddebug_cpu_sdpa_routeonce route stability is confirmed in CI; tests should assert numerical parity rather than route integers.Type of Change
Feature
Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.