ark: add BF16/FP16 Sparge variants - #2145
Draft
yiliu30 wants to merge 47 commits into
Draft
Conversation
Add ark-xbuild*, ark-xbuild-builtin*, and xbuild_bf16* to the ark .gitignore. These are local CMake build outputs (hundreds of MB each) and must not be tracked.
- sdpa_sparse_bf16.cpp: route BF16 sparse prefill through the sparse
SAGE launchers (launch_sparse_sage_prefill_kernel_{128,64}).
- sycl_tla_common.hpp: declare the BF16 sparse sdpa entry points
(d64, row_linear, qtile256_row64k).
- ark.cpp: expose the low-level BF16 sparse binding via pybind.
- __init__.py: export sage_sparse_bf16 and the BF16 e2e entry point.
- CMakeLists.txt: add the BF16 sparse kernel sources to the build.
Two device-level bugs made the native BF16 sparse kernel unusable: - SparseSageConfig hardcoded the Q*K MMA to an int8 DPAS (XE_DPAS_TT<..., int32_t, int8_t>) regardless of ElementQ. Feeding bf16 fragments into the int8 DPAS hung the device; INT8 worked only because its fragments genuinely are int8. The MMA is now ElementQ-aware: int8 keeps the int32/int8 DPAS, while bf16/fp16 use a float/<ElementQ> DPAS. - The mainloop applied the softmax scale only inside the scale_block_size != 0 branches. The native BF16 path uses scale_block_size 0 (no quantization), so the 1/sqrt(d) scale was never applied and the softmax ran on unscaled scores. Added the else branch that multiplies by params.scale.
- _block_map_lut_torch: replace the boolean-mask assignment (filled_matrix[~block_map] = 10_000_000) with torch.where. On the XPU backend that assignment silently misses entries, leaving 0s that become -1 LUT entries and drive _lut_to_block_map's scatter_ out of bounds (ScatterGatherKernels assert). Replace the same unsafe pattern (final_tile_map[~mask] = True) with |= in sparse_attention and the triton preprocess. - _lut_to_block_map: scatter True at the recovered block positions instead of scattering the valid mask. Masked-out tail indices all point at slot 0, so their trailing False writes were silently dropping a selected block at position 0; the round-trip is now lossless (empty rows cleared explicitly). - sage_sparse_bf16: route to the native BF16 sparse kernel when available, keeping the dense-mask fallback for unsupported configs. The fallback materializes a [B,Hq,Sq,Skv] mask and OOMs at long sequences. - _fill_block_map_torch: vectorize the O(k_blocks) python loop (thousands of sequential kernel launches at long sequences) into a gather/cumsum/scatter. Preprocess at seq 75600 drops from ~1170 ms to ~450 ms.
- bench_sparse_topk.py: add --dtype bf16, the sparse_bf16_kernel_only / sparse_bf16_e2e modes, and baseline/effective TFLOPS columns. The binding loader also accepts the ark-xbuild artifact. - tools/repro_sparse_bf16_sdpa_bench.sh: one-command repro that sources oneAPI 2026.1, pins a device, and runs the BF16 sparse benchmark against the rebuilt ark-xbuild extension.
- run_lut_roundtrip_case: the LUT must be non-negative and the block_map -> lut -> block_map round-trip lossless (regression for the XPU boolean-mask assignment bug that produced -1 LUT entries). - run_fill_block_map_case: the vectorized _fill_block_map_torch must match the reference loop implementation (regression for the preprocess speedup). - test_sage_sparse_prefill_e2e: include ark-xbuild in the binding search roots.
sage_sparse_bf16 now uses the native XPU sparse kernel exclusively. The dense-mask fallback (rebuild the full block map, then run per-head dense sagev1) is removed: it materialized a [B, Hq, Sq, Skv] mask and OOM'd at long sequences, and the native kernel now covers every supported config. Unsupported configs raise instead of silently degrading to the slow path.
tools/run_bf16_sparse_bench.sh documents and runs the exact BF16 sparse SDPA benchmark command: sources oneAPI 2026.1, pins a GPU via ZE_AFFINITY_MASK, uses the rebuilt ark-xbuild extension, and saves a timestamped CSV + log under benchmarks/results/. Includes env overrides and prerequisites in the header.
triton 3.7.x enables has_predicated_io for non-LTS Intel GPU drivers and emits the SPV_INTEL_predicated_io SPIR-V extension. Some level-zero loaders (e.g. 1.13.35563) reject that extension at kernel load time, aborting every triton kernel launch with 'InvalidModule: ... unknown extension'. Since it is not a catchable Python exception, the torch preprocess fallback never ran and the whole benchmark died. Add auto_round_kernel/_xpu_triton_compat.py which patches the intel triton backend to force has_predicated_io off (standard predicated loads instead), and apply it when the triton preprocess module loads. The triton-xpu preprocess backend now works and is faster than torch (seq 32768/50% e2e: 327ms -> 206ms).
tools/run_sparse_sagev1_bench.sh documents and runs the sparse SAGE v1 (fp16 / INT8) benchmark on main: sources oneAPI, pins a GPU, uses the rebuilt extension, and saves timestamped CSV + log under benchmarks/results/. Includes env overrides and a note on the triton preprocess backend.
triton 3.7.x enables has_predicated_io for non-LTS Intel GPU drivers and emits the SPV_INTEL_predicated_io SPIR-V extension. Some level-zero loaders (e.g. 1.13.35563) reject that extension at kernel load time, aborting every triton kernel launch with 'InvalidModule: ... unknown extension'. Since it is not a catchable Python exception, the torch preprocess fallback never ran and the whole benchmark died. Add auto_round_kernel/_xpu_triton_compat.py which patches the intel triton backend to force has_predicated_io off (standard predicated loads instead), and apply it when the triton preprocess module loads. The triton-xpu preprocess backend now works and is faster than torch (seq 32768/50% e2e: 327ms -> 206ms).
run_bf16_sparse_bench.sh now uses the triton_xpu preprocess backend by default (override via SPARGE_PREPROCESS_BACKEND=torch), enabled by the SPV_INTEL_predicated_io workaround in _xpu_triton_compat. This lifts BF16 sparse e2e speedup vs dense torch SDPA to 1.1-1.3x at 50% selection and ~4x at 12.5%. Update the speedup chart with the triton-backend results.
Wire up the dormant SPARSESDPAFwdMainloop into a new SparseSDPAConfig with ElementQ-aware (bf16/fp16) MMA, plus launchers, a wrapper, declarations, and a sage_sparse_sdpa pybind. Native-precision path applies the softmax scale directly and carries no INT8 dequant baggage, unlike the shared sparse-SAGE path.
… modes Add sage_sparse_sdpa and sparge_sage2_attn_meansim_topk_xpu_sdpa (bf16+fp16), route sage_sparse_bf16 through the sparse-SDPA path, and add sparse_sdpa_* modes to the benchmark.
Add run_case_sdpa/run_case_sdpa_full (bf16+fp16 vs dense ref + sparse-SAGE parity + topk=1.0 dense gate). Compute the dense reference on CPU because XPU torch.softmax is numerically broken on large masked tensors (sums to <1), which also un-breaks the existing bf16 qtile256 test case.
… path The BF16 native sparse attention path previously ran through the INT8-centric SAGE kernel structures (sage_sparse_bf16 -> SparseSageConfig -> SPARSESAGEV1FwdMainloop). Remove that coupling so each dtype has one clean path: - INT8: sage_sparse / sparge_sage2_attn_meansim_topk_xpu - BF16/FP16: sage_sparse_sdpa / sparge_sage2_attn_meansim_topk_xpu_sdpa Changes: - delete sdpa_sparse_bf16.cpp (superseded by sdpa_sparse_sdpa.cpp) - drop sage_sparse_bf16 pybind + Python alias and sparge_sage2_attn_meansim_topk_xpu_bf16, plus __init__ exports - remove sdpa_impl_bf16_sparse_* declarations from sycl_tla_common.hpp - hardcode SparseSageConfig MMA to int32/int8 DPAS; static_assert ElementQ==int8 - remove the native softmax-scale branch from the INT8 SAGE mainloop - guard the scale_block_size division in XeSparseSageFwdKernel (0 on native path) - align _query_tile_tokens_for_head_dim(64) default to 64 with the e2e wrappers - drop bf16-only cases from the prefill e2e test and bf16-only bench modes - switch the flux example to the native sparse-SDPA entry points - gitignore generated benchmarks/results/ Co-Authored-By: Claude <noreply@anthropic.com>
Bring the handoff in sync with the refactored state: one native sage_sparse_sdpa path per precision, the deleted sdpa_sparse_bf16.cpp, the re-configure requirement for the GLOB source list, the torch preprocess backend caveat, and the confirmed full-sweep benchmark numbers. Co-Authored-By: Claude <noreply@anthropic.com>
512x512 render (50 steps) with the native bf16 sparse attention patch at topk 0.5 (avg sparsity 50%, all 2850 attention calls through the sparse path). Reference for the FLUX sparse-attention investigation. Co-Authored-By: Claude <noreply@anthropic.com>
for more information, see https://pre-commit.ci
The bf16 sparse feature is source-only for the formal PR: remove the dev handoff note (no _CN counterpart) and the generated benchmark/render PNGs. Co-Authored-By: Claude <noreply@anthropic.com>
…l/auto-round into feat/sparse-bf16-prefill-v2
for more information, see https://pre-commit.ci
Signed-off-by: yiliu30 <yi4.liu@intel.com>
# Conflicts: # auto_round_extension/ark/.gitignore
for more information, see https://pre-commit.ci
…o feat/sparse-bf16-prefill-v2
…l/auto-round into feat/sparse-bf16-prefill-v2
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
1 participant
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
Please briefly describe your main changes, the motivation.
Type of Change
Bug fix
Related Issues
Fixes or relates to #
Checklist Before Submitting
/azp run Unit-Test-CUDA-AutoRound.Sparse BF16 Benchmark — 75000 tokens (bf16, head_dim 128, 40 heads, qtile256)
-B60 (K32)
BF16 Sparse SDPA
FP16 Sparse SDPA
Source:
benchmarks/results/bench_sparse_bf16_20260812_145411.csv