diff --git a/auto_round_extension/ark/BKC_KERNEL_BENCH.md b/auto_round_extension/ark/BKC_KERNEL_BENCH.md new file mode 100644 index 000000000..9e4e894ad --- /dev/null +++ b/auto_round_extension/ark/BKC_KERNEL_BENCH.md @@ -0,0 +1,321 @@ +# Kernel Build / Benchmark BKC + +This is the shortest known-good flow to compile the XPU sparse prefill kernel and run the kernel benchmark on this branch. + +## Environment + +```bash +cd /home/yiliu4/workspace/auto-round-prefill-clean-sparse-pr/auto_round_extension/ark +source /opt/intel/oneapi/setvars.sh +source /home/yiliu4/workspace/auto-round-py/.venv/bin/activate +``` + +## Build + +Recommended build for the current sparse prefill path: + +```bash +python -m cmake -S auto_round_kernel -B auto_round_kernel/xbuild \ + -DARK_XPU=ON \ + -DARK_SYCL_TLA=ON + +python -m cmake --build auto_round_kernel/xbuild \ + --target auto_round_kernel_xpu \ + -j 4 +``` + +## Sanity + +```bash +python auto_round_kernel/wrapper/test/test_sage_sparse_prefill_e2e.py +python auto_round_kernel/wrapper/test/test_sparge_preprocess_topk_e2e.py +python -m pytest -q test/test_sparge_preprocess_helpers.py +``` + +## Benchmark + +Recommended benchmark for the supported `q_tile=256` kernel: + +```bash +ZE_AFFINITY_MASK=1 \ +python test/bench_sparse_topk.py \ + --batch 1 \ + --num-heads-q 40 \ + --num-heads-kv 40 \ + --seq-len 75600 \ + --head-dim 128 \ + --tensor-layout NHD \ + --topk 0.5 \ + --q-tile-override 256 \ + --sparse-q-block-tokens 256 \ + --sparse-k-block-tokens 64 \ + --warmup 2 \ + --iters 3 \ + --output-csv bench_sparse_topk_bkc_qtile256_k64.csv +``` + +The script prints a summary table and writes the CSV to: + +```bash +bench_sparse_topk_bkc_qtile256_k64.csv +``` + +## Pick A Free GPU + +Before running the benchmark, check that the target XPU is idle: + +```bash +xpu-smi ps +xpu-smi dump -d -1 -m 18 +``` + +What to look for: + +- no `bench_sparse_topk.py` or other user compute process attached in `xpu-smi ps` +- flat, low memory on all devices in `xpu-smi dump` + +On this node, `xpu-smi` compute-util counters were not available without extra privilege, so the practical check was: + +- no active XPU process attached +- steady low memory footprint across devices + +The validated rerun used `ZE_AFFINITY_MASK=6`. + +## Historical Sparse K-Prefetch A/B + +The current branch keeps sparse `K` prefetch enabled by default, so there is no +longer a build-time toggle in `CMakeLists.txt`. + +The `prefetch=ON/OFF` comparisons below are historical A/B data captured before +that cleanup, when both paths were still build-selectable. + +Run `NHD` and `HND` on the chosen device: + +```bash +ZE_AFFINITY_MASK=6 \ +python test/bench_sparse_topk.py \ + --batch 1 \ + --num-heads-q 40 \ + --num-heads-kv 40 \ + --seq-len 75600 \ + --head-dim 128 \ + --tensor-layout NHD \ + --topk 0.5 0.3 0.1 \ + --q-tile-override 256 \ + --sparse-q-block-tokens 256 \ + --sparse-k-block-tokens 64 \ + --warmup 2 \ + --iters 3 \ + --output-csv bench_sparse_topk_prefetch_on_qtile256_k64_nhd_gpu6.csv + +ZE_AFFINITY_MASK=6 \ +python test/bench_sparse_topk.py \ + --batch 1 \ + --num-heads-q 40 \ + --num-heads-kv 40 \ + --seq-len 75600 \ + --head-dim 128 \ + --tensor-layout HND \ + --topk 0.5 0.3 0.1 \ + --q-tile-override 256 \ + --sparse-q-block-tokens 256 \ + --sparse-k-block-tokens 64 \ + --warmup 2 \ + --iters 3 \ + --output-csv bench_sparse_topk_prefetch_on_qtile256_k64_hnd_gpu6.csv +``` + +For the old `prefetch=OFF` path, rerun the same `NHD` and `HND` commands with: + +```bash +--output-csv bench_sparse_topk_prefetch_off_qtile256_k64_nhd_gpu6.csv +--output-csv bench_sparse_topk_prefetch_off_qtile256_k64_hnd_gpu6.csv +``` + +## Validated Results + +GPU6 rerun, same node, same shape, same `q_tile=256` / `k_block=64` kernel. + +`NHD`, prefetch `ON` vs `OFF`: + +- `topk=0.5`: kernel `796.465 -> 482.584 ms` (`39.4%` faster), e2e `916.359 -> 577.503 ms` (`37.0%` faster) +- `topk=0.3`: kernel `478.918 -> 291.787 ms` (`39.1%` faster), e2e `599.025 -> 387.981 ms` (`35.2%` faster) +- `topk=0.1`: kernel `161.504 -> 98.049 ms` (`39.3%` faster), e2e `282.359 -> 213.140 ms` (`24.5%` faster) + +`HND`, prefetch `ON` vs `OFF`: + +- `topk=0.5`: kernel `694.479 -> 409.535 ms` (`41.0%` faster), e2e `823.376 -> 517.995 ms` (`37.1%` faster) +- `topk=0.3`: kernel `417.675 -> 256.283 ms` (`38.6%` faster), e2e `544.377 -> 367.859 ms` (`32.4%` faster) +- `topk=0.1`: kernel `144.716 -> 89.758 ms` (`38.0%` faster), e2e `272.264 -> 216.839 ms` (`20.4%` faster) + +Cross-device consistency against the earlier GPU4 run: + +- `NHD`: GPU4 vs GPU6 delta stayed within `0.0-2.7%` for prefetch `ON`, and within `0.2-0.6%` for prefetch `OFF` +- `HND`: GPU4 vs GPU6 delta stayed within `1.0-4.2%` for prefetch `ON`, and within `0.0-0.3%` for prefetch `OFF` + +Saved CSVs: + +- `bench_sparse_topk_prefetch_on_qtile256_k64_nhd_gpu6.csv` +- `bench_sparse_topk_prefetch_on_qtile256_k64_hnd_gpu6.csv` +- `bench_sparse_topk_prefetch_off_qtile256_k64_nhd_gpu6.csv` +- `bench_sparse_topk_prefetch_off_qtile256_k64_hnd_gpu6.csv` + +## B70 Profiling Snapshot + +This node is `0xe223` (`B70` / `bmg_g31`). The following `unitrace` A/B was run +for the same sparse case in `HND` layout: + +- `batch=1` +- `num_heads_q=40` +- `num_heads_kv=40` +- `seq_len=75600` +- `head_dim=128` +- `topk=0.5` +- `q_tile=256` +- `sparse_q_block_tokens=256` +- `sparse_k_block_tokens=64` +- `ZE_AFFINITY_MASK=0` + +Historical build toggle used for this A/B: + +- `ARK_SPARSE_SAGE_ENABLE_K_PREFETCH=OFF` +- `ARK_SPARSE_SAGE_ENABLE_K_PREFETCH=ON` + +Machine-specific profiling command: + +```bash +sg render -c 'bash -lc " +cd /home/yiliu4/workspace/auto-round-prefill-clean-sparse-pr/auto_round_extension/ark +export LD_LIBRARY_PATH=/opt/intel/oneapi/compiler/2025.3/lib:/opt/intel/oneapi/2025.3/lib:${LD_LIBRARY_PATH} +ZE_AFFINITY_MASK=0 \ +/home/yiliu4/workspace/pti-gpu/tools/unitrace/install_local/bin/unitrace \ + -d -s --chrome-kernel-logging --demangle \ + --devices-to-sample 0 \ + --output sparse_vecstall_sampling_prefetch_ \ + /home/yiliu4/workspace/auto-round-py/.venv/bin/python \ + test/bench_sparse_topk.py \ + --batch 1 \ + --num-heads-q 40 \ + --num-heads-kv 40 \ + --seq-len 75600 \ + --head-dim 128 \ + --tensor-layout HND \ + --topk 0.5 \ + --q-tile-override 256 \ + --sparse-q-block-tokens 256 \ + --sparse-k-block-tokens 64 \ + --warmup 0 \ + --iters 1 +"' +``` + +Kernel properties from `unitrace`: + +| Setting | Sparse Private Mem / Thread | Sparse Spill / Thread | Dense Spill / Thread | Notes | +|---|---:|---:|---:|---| +| `prefetch=OFF` | `1536 B` | `3392 B` | `128 B` | large sparse spill remains | +| `prefetch=ON` | `2048 B` | `640 B` | `128 B` | spill drops a lot, private mem rises | + +Sparse kernel timing from the same profiled runs: + +| Setting | XeSparseSageFwdKernel Calls | Total Time (ns) | Avg Time (ns) | +|---|---:|---:|---:| +| `prefetch=OFF` | `2` | `1377784270` | `688892135` | +| `prefetch=ON` | `2` | `868062916` | `434031458` | + +Profiled benchmark output from the same runs: + +| Setting | dense_torch_sdpa | dense_sagev1 | sparse kernel_only | sparse e2e | +|---|---:|---:|---:|---:| +| `prefetch=OFF` | `2160.233 ms` | `8896.800 ms` | `40102.633 ms` | `837.385 ms` | +| `prefetch=ON` | `2014.259 ms` | `873.501 ms` | `42590.231 ms` | `620.531 ms` | + +Interpretation: + +- On this `B70` node, enabling sparse `K` prefetch increases sparse private + memory but reduces sparse spill substantially. +- The sparse kernel also ran faster under `unitrace` with prefetch enabled. +- Treat the benchmark latencies in this section as profiler-distorted. Use them + only for `prefetch ON/OFF` comparison inside `unitrace`, not as normal + benchmark numbers. + +Saved profiling artifacts: + +- `sparse_vecstall_sampling_prefetch_off.3716760` +- `python.3716760.json` +- `sparse_vecstall_sampling_prefetch_on.3756979` +- `python.3756979.json` + +## B70 `sycl_tla` Tag A/B + +On the same `B70` node, with sparse `K` prefetch kept enabled, a separate build +was created with: + +- baseline tree: `auto_round_kernel/xbuild` +- baseline tag: `SYCL_TLA_GIT_TAG=260409` +- comparison tree: `auto_round_kernel/xbuild_tla260630_prefetch_on` +- comparison tag: `SYCL_TLA_GIT_TAG=260630` + +The benchmark script only loads the extension from `xbuild/`, so the rerun used +the `xbuild` binary first, then temporarily swapped in the `260630` `.so`, ran +the same benchmark, and restored the original `xbuild` binary afterward. + +Rerun conditions: + +- free device chosen from `xpu-smi ps`: `GPU7` +- `ZE_AFFINITY_MASK=7` +- `batch=1` +- `num_heads_q=40` +- `num_heads_kv=40` +- `seq_len=75600` +- `head_dim=128` +- `topk=0.5` +- `q_tile=256` +- `sparse_q_block_tokens=256` +- `sparse_k_block_tokens=64` +- `warmup=2` +- `iters=3` + +Sparse benchmark comparison: + +| Layout | `260409` kernel-only | `260630` kernel-only | Delta | +|---|---:|---:|---:| +| `HND` | `409.834 ms` | `399.086 ms` | `2.6%` faster | +| `NHD` | `481.936 ms` | `465.911 ms` | `3.3%` faster | + +| Layout | `260409` e2e | `260630` e2e | Delta | +|---|---:|---:|---:| +| `HND` | `525.080 ms` | `515.407 ms` | `1.8%` faster | +| `NHD` | `582.289 ms` | `564.575 ms` | `3.0%` faster | + +Interpretation: + +- On this rerun, `SYCL_TLA_GIT_TAG=260630` was consistently but modestly faster + than `260409`. +- The improvement was small on `HND` and slightly larger on `NHD`. +- This comparison was done with sparse `K` prefetch enabled in both builds. + +Saved CSVs: + +- `bench_sparse_topk_prefetch_on_tla260409_hnd_gpu7.csv` +- `bench_sparse_topk_prefetch_on_tla260409_nhd_gpu7.csv` +- `bench_sparse_topk_prefetch_on_tla260630_hnd_gpu7.csv` +- `bench_sparse_topk_prefetch_on_tla260630_nhd_gpu7.csv` + +## Optional `q_tile=64` Reference + +```bash +ZE_AFFINITY_MASK=4 \ +python test/bench_sparse_topk.py \ + --batch 1 \ + --num-heads-q 40 \ + --num-heads-kv 40 \ + --seq-len 75600 \ + --head-dim 128 \ + --tensor-layout NHD \ + --topk 0.5 0.3 0.1 \ + --q-tile-override 64 \ + --warmup 2 \ + --iters 3 \ + --output-csv bench_sparse_topk_bkc_qtile64.csv +``` diff --git a/auto_round_extension/ark/REPRODUCE.md b/auto_round_extension/ark/REPRODUCE.md new file mode 100644 index 000000000..6ded18a76 --- /dev/null +++ b/auto_round_extension/ark/REPRODUCE.md @@ -0,0 +1,137 @@ +# Sparse Prefill Clean Branch + +This branch keeps the sparse prefill runtime on top of `main` with two supported query-tile modes: + +- `q_tile=64` via the generic sparse prefill path +- `q_tile=256` with `sparse_q_block_tokens=256` and `sparse_k_block_tokens=64` + +## Build + +```bash +cd auto_round_extension/ark +source /opt/intel/oneapi/setvars.sh +python -m cmake -S auto_round_kernel -B auto_round_kernel/xbuild -DARK_XPU=ON -DARK_SYCL_TLA=ON +python -m cmake --build auto_round_kernel/xbuild --target auto_round_kernel_xpu -j 4 +``` + +## Tests + +```bash +cd auto_round_extension/ark +python \ + auto_round_kernel/wrapper/test/test_sage_sparse_prefill_e2e.py +python \ + auto_round_kernel/wrapper/test/test_sparge_preprocess_topk_e2e.py +python -m pytest -q \ + test/test_sparge_preprocess_helpers.py +``` + +## Benchmark + +`q_tile=64`: + +```bash +ZE_AFFINITY_MASK=4 \ +python \ + test/bench_sparse_topk.py \ + --batch 1 --num-heads-q 40 --num-heads-kv 40 --seq-len 75600 --head-dim 128 \ + --tensor-layout NHD --topk 0.5 0.3 --q-tile-override 64 +``` + +`q_tile=256`, decoupled sparse rows: + +```bash +ZE_AFFINITY_MASK=4 \ +python \ + test/bench_sparse_topk.py \ + --batch 1 --num-heads-q 40 --num-heads-kv 40 --seq-len 75600 --head-dim 128 \ + --tensor-layout NHD --topk 0.5 0.3 \ + --q-tile-override 256 \ + --sparse-q-block-tokens 256 \ + --sparse-k-block-tokens 64 +``` + +## Wan Example + +The WAN example defaults now follow the `Wan2.2-T2V-A14B` diffusers model-card +example: `1280x720`, `81` frames, `40` steps, `guidance_scale=4.0`, +`guidance_scale_2=3.0`, `fps=16`. + +`q_tile=64`: + +```bash +cd auto_round_extension/ark/examples +WAN_USE_SPARSE=1 \ +WAN_SPARSE_TOPK=0.5 \ +WAN_SPARSE_Q_TILE_OVERRIDE=64 \ +python run_wan.py +``` + +`q_tile=256`: + +```bash +cd auto_round_extension/ark/examples +WAN_USE_SPARSE=1 \ +WAN_SPARSE_TOPK=0.5 \ +WAN_SPARSE_Q_TILE_OVERRIDE=256 \ +WAN_SPARSE_Q_BLOCK_TOKENS=256 \ +WAN_SPARSE_K_BLOCK_TOKENS=64 \ +python run_wan.py +``` + +## Flux Example + +`q_tile=64`: + +```bash +cd auto_round_extension/ark/examples +FLUX_USE_SPARSE=1 \ +FLUX_SPARSE_TOPK=0.5 \ +FLUX_SPARSE_Q_TILE_OVERRIDE=64 \ +python run_flux.py +``` + +`q_tile=256`: + +```bash +cd auto_round_extension/ark/examples +FLUX_USE_SPARSE=1 \ +FLUX_SPARSE_TOPK=0.5 \ +FLUX_SPARSE_Q_TILE_OVERRIDE=256 \ +FLUX_SPARSE_Q_BLOCK_TOKENS=256 \ +FLUX_SPARSE_K_BLOCK_TOKENS=64 \ +python run_flux.py +``` + +## Flux Sweep + +Dense baseline plus sparse `topk=1.0..0.1`, using 8 devices and the `q_tile=256`, +`sparse_q_block_tokens=256`, `sparse_k_block_tokens=64` kernel: + +```bash +cd auto_round_extension/ark/examples +FLUX_SWEEP_PYTHON=/home/yiliu4/workspace/auto-round-py/.venv/bin/python \ +FLUX_MODEL=/home/yiliu4/workspace/models/black-forest-labs/FLUX.1-dev \ +FLUX_SWEEP_DEVICES=0,1,2,3,4,5,6,7 \ +bash run_flux_sweep.sh +``` + +## WAN Sweep + +Dense baseline plus sparse `topk=1.0..0.1`, using the `Wan2.2-T2V-A14B` model-card +default settings and fixed sparse kernel settings +(`q_tile=256`, `sparse_q_block_tokens=256`, `sparse_k_block_tokens=64`): + +```bash +cd auto_round_extension/ark/examples +WAN_SWEEP_PYTHON=/home/yiliu4/workspace/auto-round-py/.venv/bin/python \ +WAN_MODEL=/home/yiliu4/workspace/models/Wan-AI/Wan2.2-T2V-A14B-Diffusers \ +WAN_SWEEP_DEVICE_POOLS='0,1,2,3;4,5,6,7' \ +bash run_wan_sweep.sh +``` + +Execution order per config: + +- first try `WAN_DEVICE_MAP=balanced` on a 4-XPU pool +- then retry `WAN_CPU_OFFLOAD_MODE=model` on the pool's primary device +- finally retry `WAN_CPU_OFFLOAD_MODE=sequential` if needed diff --git a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt index 5d40df7a4..10fced65c 100755 --- a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt +++ b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt @@ -72,7 +72,11 @@ else() endif() list(APPEND libs bestla) -file(GLOB SRCS ${CMAKE_CURRENT_SOURCE_DIR}/*.cpp) +set(SRCS + ${CMAKE_CURRENT_SOURCE_DIR}/ark.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sdpa.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/sdpa_sparse.cpp +) set(SDPA_GENERATED_SRCS) set(SDPA_KERNEL_DECLARATIONS) @@ -128,6 +132,7 @@ if(ARK_XPU AND ARK_SYCL_TLA) ${CUTLASS_DIR}/examples/06_bmg_flash_attention ${CUTLASS_DIR}/examples/12_xe20_moe_gemm_cute_interface ) + set(_ark_compat_include_dir ${CMAKE_CURRENT_SOURCE_DIR}/compat/include) include(${CMAKE_CURRENT_LIST_DIR}/sdpa_generation.cmake) endif() @@ -141,6 +146,7 @@ if(ARK_XPU AND ARK_SYCL_TLA) target_compile_definitions(${PY_NAME} PRIVATE ARK_SYCL_TLA=1 CUTLASS_ENABLE_SYCL=1 SYCL_INTEL_TARGET=1) target_include_directories(${PY_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/wrapper/include) target_include_directories(${PY_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/generated/sdpa) + target_include_directories(${PY_NAME} SYSTEM PRIVATE ${_ark_compat_include_dir}) # Use SYSTEM include directories for sycl-tla to suppress warnings/errors from third-party headers # (e.g., std::common_type specialization issues in traits.hpp) foreach(_inc_dir IN LISTS _sycl_tla_include_dirs) @@ -167,6 +173,7 @@ if(ARK_UT) if(ARK_XPU AND ARK_SYCL_TLA) target_compile_definitions(${TEST_NAME} PRIVATE ARK_SYCL_TLA=1 CUTLASS_ENABLE_SYCL=1 SYCL_INTEL_TARGET=1) target_include_directories(${TEST_NAME} PRIVATE ${CMAKE_CURRENT_BINARY_DIR}/generated/sdpa) + target_include_directories(${TEST_NAME} SYSTEM PRIVATE ${_ark_compat_include_dir}) foreach(_inc_dir IN LISTS _sycl_tla_include_dirs) if(EXISTS "${_inc_dir}") target_include_directories(${TEST_NAME} SYSTEM PRIVATE "${_inc_dir}") diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index 512fab6a5..79ca62aac 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -13,6 +13,9 @@ # limitations under the License. import os +import importlib.util +import sys +from pathlib import Path from typing import Optional import torch @@ -207,6 +210,37 @@ def _empty_attention_output( cpu_lib = None xpu_lib = None + +def _load_local_xpu_lib(): + module_dir = Path(__file__).resolve().parent + candidates = sorted((module_dir / "xbuild").glob("auto_round_kernel_xpu*.so")) + if not candidates: + return None + ext_path = candidates[-1] + module_name = "auto_round_kernel._local.auto_round_kernel_xpu" + spec = importlib.util.spec_from_file_location(module_name, ext_path) + if spec is None or spec.loader is None: + return None + sys.modules.pop(module_name, None) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +def _prefer_local_xpu_lib(lib): + module_dir = Path(__file__).resolve().parent + lib_path = getattr(lib, "__file__", None) + if lib_path is not None: + try: + if Path(lib_path).resolve().is_relative_to(module_dir): + return lib + except Exception: + pass + local_lib = _load_local_xpu_lib() + return local_lib if local_lib is not None else lib + + try: from . import auto_round_kernel_cpu as _cpu_lib_mod @@ -219,10 +253,11 @@ def _empty_attention_output( try: from . import auto_round_kernel_xpu as _xpu_lib_mod - xpu_lib = _xpu_lib_mod + xpu_lib = _prefer_local_xpu_lib(_xpu_lib_mod) except ImportError as _e: - print(f"ARK is unable to load XPU lib: {_e}") - xpu_lib = None + xpu_lib = _load_local_xpu_lib() + if xpu_lib is None: + print(f"ARK is unable to load XPU lib: {_e}") def get_lib(A: torch.Tensor): @@ -1280,6 +1315,23 @@ def sageattn( ) +from .sparse_attention import ( + _block_map_lut_torch, + _build_block_causal_mask, + _build_sparge_preprocess_context, + _finalize_sparge_preprocess_outputs, + _from_hnd, + _sparge_preprocess_topk_torch, + _sequence_mean_native_layout, + _slice_sequence_native_layout, + _to_hnd, + sage_sparse, + sparge_block_map_to_mask, + sparge_preprocess_topk, + sparge_sage2_attn_meansim_topk_xpu, +) + + def sageattn_varlen( q: torch.Tensor, k: torch.Tensor, diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index 279083b26..94b71ab3d 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -404,6 +404,109 @@ static void sage_pvi8(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, t is_causal, (BTLA_DTYPE)o_dtype, (float*)lse); } +static void sage_sparse(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, torch_ptr O, torch_ptr mask, + int scale_block_size, torch_ptr qscale, torch_ptr kscale, torch_ptr lut, + torch_ptr valid_block_num, int num_q_blocks, int num_k_blocks, int q_tile_override, int q_stride_s, int q_stride_d, + int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, int k_stride_h, int k_stride_b, + int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, int o_stride_s, + int o_stride_d, int o_stride_h, int o_stride_b, int q_dtype, int k_dtype, int o_dtype, + int batch, int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, int head_dim, + float softmax_scale, bool is_causal) { + if (mask && is_causal) { + throw std::invalid_argument("ark::sage_sparse: mask and is_causal cannot both be set"); + } + if (!lut || !valid_block_num) { + throw std::invalid_argument("ark::sage_sparse: lut and valid_block_num must be provided"); + } + if (head_dim == 64) { + ark::sdpa_impl_qks8_sparse_d64_pvhalf( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, scale_block_size, (void*)qscale, + (void*)kscale, (void*)lut, (void*)valid_block_num, num_q_blocks, num_k_blocks, q_tile_override, q_stride_s, + q_stride_d, q_stride_h, q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, + v_stride_h, v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal, (BTLA_DTYPE)o_dtype); + return; + } + ark::sdpa_impl_qks8_sparse_row_linear_pvhalf( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, scale_block_size, (void*)qscale, + (void*)kscale, (void*)lut, (void*)valid_block_num, num_q_blocks, num_k_blocks, q_tile_override, q_stride_s, + q_stride_d, q_stride_h, q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, + v_stride_h, v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal, (BTLA_DTYPE)o_dtype); +} + +static void sage_sparse_qtile256_row64k( + torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, torch_ptr O, torch_ptr mask, int scale_block_size, + torch_ptr qscale, torch_ptr kscale, torch_ptr lut, torch_ptr valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int q_dtype, int k_dtype, int o_dtype, int batch, + int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, + bool is_causal) { + if (mask && is_causal) { + throw std::invalid_argument("ark::sage_sparse_qtile256_row64k: mask and is_causal cannot both be set"); + } + if (!lut || !valid_block_num) { + throw std::invalid_argument("ark::sage_sparse_qtile256_row64k: lut and valid_block_num must be provided"); + } + ark::sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, scale_block_size, (void*)qscale, + (void*)kscale, (void*)lut, (void*)valid_block_num, num_q_blocks, num_k_blocks, q_tile_override, q_stride_s, + q_stride_d, q_stride_h, q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, + v_stride_h, v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal, (BTLA_DTYPE)o_dtype); +} + +static void sage_sparse_row_linear(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, torch_ptr O, torch_ptr mask, + int scale_block_size, torch_ptr qscale, torch_ptr kscale, torch_ptr lut, + torch_ptr valid_block_num, int num_q_blocks, int num_k_blocks, int q_tile_override, + int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, + int v_stride_h, int v_stride_b, int o_stride_s, int o_stride_d, int o_stride_h, + int o_stride_b, int q_dtype, int k_dtype, int o_dtype, int batch, int num_heads_q, + int num_heads_kv, int seq_len_q, int seq_len_kv, int head_dim, + float softmax_scale, bool is_causal) { + if (mask && is_causal) { + throw std::invalid_argument("ark::sage_sparse_row_linear: mask and is_causal cannot both be set"); + } + if (!lut || !valid_block_num) { + throw std::invalid_argument("ark::sage_sparse_row_linear: lut and valid_block_num must be provided"); + } + ark::sdpa_impl_qks8_sparse_row_linear_pvhalf( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, scale_block_size, (void*)qscale, + (void*)kscale, (void*)lut, (void*)valid_block_num, num_q_blocks, num_k_blocks, q_tile_override, q_stride_s, + q_stride_d, q_stride_h, q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, + v_stride_h, v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal, (BTLA_DTYPE)o_dtype); +} + +static void sage_sparse_decode(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, torch_ptr K_cache, + torch_ptr V_cache, torch_ptr O, torch_ptr mask, int scale_block_size, torch_ptr qscale, + torch_ptr kscale, torch_ptr lut, torch_ptr valid_block_num, int num_q_blocks, + int num_k_blocks, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, + int k_stride_s, int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, + int v_stride_s, int v_stride_h, int v_stride_b, int o_stride_s, int o_stride_d, + int o_stride_h, int o_stride_b, int q_dtype, int k_dtype, int o_dtype, int batch, + int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, int seq_len_kv_cache, + int head_dim, float softmax_scale, bool is_causal) { + if (mask && is_causal) { + throw std::invalid_argument("ark::sage_sparse_decode: mask and is_causal cannot both be set"); + } + if (!lut || !valid_block_num) { + throw std::invalid_argument("ark::sage_sparse_decode: lut and valid_block_num must be provided"); + } + if (!K_cache || !V_cache) { + throw std::invalid_argument("ark::sage_sparse_decode: K_cache and V_cache must be provided"); + } + ark::sdpa_impl_qks8_sparse_decode_pvhalf( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)K_cache, (void*)V_cache, (void*)O, (void*)mask, + scale_block_size, (void*)qscale, (void*)kscale, (void*)lut, (void*)valid_block_num, num_q_blocks, num_k_blocks, + q_stride_s, q_stride_d, q_stride_h, q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, + v_stride_s, v_stride_h, v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, + num_heads_kv, seq_len_q, seq_len_kv, seq_len_kv_cache, head_dim, softmax_scale, is_causal, + (BTLA_DTYPE)o_dtype); +} + static void moe_gemm_wrapper(torch_ptr stream, torch_ptr activations, torch_ptr weights, torch_ptr scales, torch_ptr outputs, int dtype, int N, int K, torch_ptr num_tokens_per_expert, int num_experts) { @@ -709,6 +812,10 @@ PYBIND11_MODULE(PY_NAME, m) { pybind11::arg("seq_len_q"), pybind11::arg("seq_len_kv"), pybind11::arg("head_dim"), pybind11::arg("softmax_scale"), pybind11::arg("is_causal"), pybind11::arg("tensor_layout"), pybind11::arg("lse") = 0); + m.def("sage_sparse", &ark::sage_sparse); + m.def("sage_sparse_qtile256_row64k", &ark::sage_sparse_qtile256_row64k); + m.def("sage_sparse_row_linear", &ark::sage_sparse_row_linear); + m.def("sage_sparse_decode", &ark::sage_sparse_decode); // Low-level SAGE PVi8 API: input Q/K/V are pre-quantized int8 with qscale/kscale/vscale. m.def("sage_pvi8", &ark::sage_pvi8, pybind11::arg("stream"), pybind11::arg("Q"), pybind11::arg("K"), pybind11::arg("V"), pybind11::arg("O"), pybind11::arg("mask"), @@ -730,4 +837,4 @@ PYBIND11_MODULE(PY_NAME, m) { m.def("moe_gemm_prefill_int_dpas", &ark::moe_gemm_prefill_int_dpas_wrapper); m.def("matmul_sycl_tla", &ark::matmul_sycl_tla); #endif -} \ No newline at end of file +} diff --git a/auto_round_extension/ark/auto_round_kernel/sdpa_sparse.cpp b/auto_round_extension/ark/auto_round_kernel/sdpa_sparse.cpp new file mode 100644 index 000000000..0046bc366 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/sdpa_sparse.cpp @@ -0,0 +1,362 @@ +// Copyright (c) 2026 Intel Corporation +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +#include +#include +#include +#include "bestla/bestla.h" + +#include +#include "sycl_tla_sdpa_sparse.hpp" + +namespace ark { + +namespace detail = sparse_detail; + +namespace { + +using KernelLauncher = int (*)(detail::Options const& options); + +int launch_prefill_kernel_f16_128_sparse_sage(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128(options); +} + +int launch_prefill_kernel_f16_128_sparse_sage_qtile128(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128_qtile128(options); +} + +int launch_prefill_kernel_f16_128_sparse_sage_qtile64(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128_qtile64(options); +} + +int launch_prefill_kernel_f16_64_sparse_sage(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_64(options); +} + +int launch_prefill_kernel_bf16_128_sparse_sage(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128(options); +} + +int launch_prefill_kernel_bf16_128_sparse_sage_qtile128(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128_qtile128(options); +} + +int launch_prefill_kernel_bf16_128_sparse_sage_qtile64(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_128_qtile64(options); +} + +int launch_prefill_kernel_bf16_64_sparse_sage(detail::Options const& options) { + return detail::launch_sparse_sage_prefill_kernel_64(options); +} + +KernelLauncher select_sparse_sage_prefill_launcher(BTLA_DTYPE pv_dtype, int head_dim, int q_tile_override) { + switch (head_dim) { + case 128: + if (q_tile_override == 64) { + return pv_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_128_sparse_sage_qtile64 + : launch_prefill_kernel_f16_128_sparse_sage_qtile64; + } + if (q_tile_override == 128) { + return pv_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_128_sparse_sage_qtile128 + : launch_prefill_kernel_f16_128_sparse_sage_qtile128; + } + if (q_tile_override != 0 && q_tile_override != 256) return nullptr; + return pv_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_128_sparse_sage + : launch_prefill_kernel_f16_128_sparse_sage; + case 64: + if (q_tile_override != 0 && q_tile_override != 128) return nullptr; + return pv_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_64_sparse_sage + : launch_prefill_kernel_f16_64_sparse_sage; + default: + return nullptr; + } +} + +detail::Options make_common_options(void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int q_stride_s, + int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, + int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, + int v_stride_b, int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, + int batch, int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, + int head_dim, float softmax_scale, bool is_causal) { + if (q_stride_d != 1 || k_stride_d != 1 || v_stride_d != 1 || o_stride_d != 1) { + throw std::invalid_argument("make_common_options: head-dim stride must be 1 for Q/K/V/O"); + } + detail::Options options; + options.q = Q_ptr; + options.k = K_ptr; + options.v = V_ptr; + options.mask = mask; + options.o = O_ptr; + options.use_tensor_strides = true; + options.q_stride_s = q_stride_s; + options.q_stride_d = q_stride_d; + options.q_stride_h = q_stride_h; + options.q_stride_b = q_stride_b; + options.k_stride_s = k_stride_s; + options.k_stride_d = k_stride_d; + options.k_stride_h = k_stride_h; + options.k_stride_b = k_stride_b; + options.v_stride_d = v_stride_d; + options.v_stride_s = v_stride_s; + options.v_stride_h = v_stride_h; + options.v_stride_b = v_stride_b; + options.o_stride_s = o_stride_s; + options.o_stride_d = o_stride_d; + options.o_stride_h = o_stride_h; + options.o_stride_b = o_stride_b; + options.batch = batch; + options.num_heads_q = num_heads_q; + options.num_heads_kv = num_heads_kv; + options.seq_len_qo = seq_len_q; + options.seq_len_kv = seq_len_kv; + options.head_size_qk = head_dim; + options.head_size_vo = head_dim; + options.softmax_scale = softmax_scale; + options.is_causal = is_causal; + return options; +} + +} // namespace + +void sparse_sage_prefill(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, + int scale_block_size, void* qscale, void* kscale, void* lut, void* valid_block_num, + int num_q_blocks, int num_k_blocks, int q_tile_override, BTLA_DTYPE pv_dtype, int q_stride_s, + int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, + int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, + int v_stride_b, int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, + int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, int head_dim, + float softmax_scale, bool is_causal, int sparse_q_block_size = 0) { + const int effective_q_tile_override = (head_dim == 128 && q_tile_override == 0) ? 64 : q_tile_override; + detail::Options options = + make_common_options(Q_ptr, K_ptr, V_ptr, O_ptr, mask, q_stride_s, q_stride_d, q_stride_h, q_stride_b, + k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, + v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, + num_heads_kv, seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal); + options.scale_block_size = scale_block_size; + options.sparse_q_block_size = sparse_q_block_size; + options.q_tile_override = effective_q_tile_override; + options.qscale = qscale; + options.kscale = kscale; + options.lut = static_cast(lut); + options.valid_block_num = static_cast(valid_block_num); + options.num_q_blocks = num_q_blocks; + options.num_k_blocks = num_k_blocks; + compat::set_default_queue(*q); + + KernelLauncher launcher = select_sparse_sage_prefill_launcher(pv_dtype, head_dim, effective_q_tile_override); + if (launcher == nullptr) { + throw std::runtime_error("Unsupported sparse_sage_prefill config"); + } + + launcher(options); +} + +void sparse_sage_decode(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* K_cache_ptr, void* V_cache_ptr, + void* O_ptr, void* mask, int scale_block_size, void* qscale, void* kscale, void* lut, + void* valid_block_num, int num_q_blocks, int num_k_blocks, BTLA_DTYPE pv_dtype, int q_stride_s, + int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, + int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, + int v_stride_b, int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, + int num_heads_q, int num_heads_kv, int seq_len_q, int seq_len_kv, int seq_len_kv_cache, + int head_dim, float softmax_scale, bool is_causal) { + const int effective_q_tile_override = head_dim == 128 ? 64 : 0; + detail::Options options = + make_common_options(Q_ptr, K_ptr, V_ptr, O_ptr, mask, q_stride_s, q_stride_d, q_stride_h, q_stride_b, + k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, + v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, + num_heads_kv, seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal); + options.scale_block_size = scale_block_size; + options.q_tile_override = effective_q_tile_override; + options.qscale = qscale; + options.kscale = kscale; + options.lut = static_cast(lut); + options.valid_block_num = static_cast(valid_block_num); + options.num_q_blocks = num_q_blocks; + options.num_k_blocks = num_k_blocks; + options.block_K = K_cache_ptr; + options.block_V = V_cache_ptr; + options.seq_len_kv_cache = seq_len_kv_cache; + compat::set_default_queue(*q); + + KernelLauncher launcher = select_sparse_sage_prefill_launcher(pv_dtype, head_dim, effective_q_tile_override); + if (launcher == nullptr) { + throw std::runtime_error( + "Unsupported dtype or head dimension for sparse_sage_decode (only F16/BF16 PV and 64/128 are supported)"); + } + + launcher(options); +} + +void sdpa_impl_qks8_sparse_d64_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, BTLA_DTYPE pv_dtype) { + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: seq_len_q and seq_len_kv must be greater than 0"); + } + if (pv_dtype != BTLA_DTYPE::F16 && pv_dtype != BTLA_DTYPE::BF16) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: only F16 and BF16 are supported for V/O dtype"); + } + if (qscale == nullptr || kscale == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: qscale and kscale must be provided"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_d64_pvhalf: num_q_blocks and num_k_blocks must be greater than 0"); + } + if (scale_block_size <= 0) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: scale_block_size must be greater than 0"); + } + if (head_dim != 64) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_d64_pvhalf: head_dim must be 64"); + } + + sparse_sage_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, scale_block_size, qscale, kscale, lut, valid_block_num, + num_q_blocks, num_k_blocks, q_tile_override, pv_dtype, q_stride_s, q_stride_d, q_stride_h, + q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, + v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal); +} + +void sdpa_impl_qks8_sparse_row_linear_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, BTLA_DTYPE pv_dtype) { + if (q_tile_override != 0 && q_tile_override != 64) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_row_linear_pvhalf: q_tile_override must be 0 or 64 for the row-linear backend"); + } + if (scale_block_size != 64) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_row_linear_pvhalf: scale_block_size must be 64 so one sparse row maps to one workgroup"); + } + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_row_linear_pvhalf: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_row_linear_pvhalf: seq_len_q and seq_len_kv must be greater than 0"); + } + if (pv_dtype != BTLA_DTYPE::F16 && pv_dtype != BTLA_DTYPE::BF16) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_row_linear_pvhalf: only F16 and BF16 are supported for V/O dtype"); + } + if (qscale == nullptr || kscale == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_row_linear_pvhalf: qscale and kscale must be provided"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_row_linear_pvhalf: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_row_linear_pvhalf: num_q_blocks and num_k_blocks must be greater than 0"); + } + + sparse_sage_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, scale_block_size, qscale, kscale, lut, valid_block_num, + num_q_blocks, num_k_blocks, 64, pv_dtype, q_stride_s, q_stride_d, q_stride_h, q_stride_b, + k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, v_stride_b, + o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, seq_len_q, + seq_len_kv, head_dim, softmax_scale, is_causal); +} + +void sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, BTLA_DTYPE pv_dtype) { + if (q_tile_override != 256) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf: q_tile_override must be 256 for the decoupled backend"); + } + if (scale_block_size != 64) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf: scale_block_size must be 64 for the decoupled backend"); + } + if (head_dim != 128) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf: head_dim must be 128 for the decoupled backend"); + } + sparse_sage_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, scale_block_size, qscale, kscale, lut, valid_block_num, + num_q_blocks, num_k_blocks, q_tile_override, pv_dtype, q_stride_s, q_stride_d, q_stride_h, + q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, + v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, head_dim, softmax_scale, is_causal, /*sparse_q_block_size=*/256); +} + +void sdpa_impl_qks8_sparse_decode_pvhalf(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* K_cache_ptr, + void* V_cache_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, + int num_q_blocks, int num_k_blocks, int q_stride_s, int q_stride_d, + int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, + int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, + int v_stride_h, int v_stride_b, int o_stride_s, int o_stride_d, + int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int seq_len_kv_cache, int head_dim, + float softmax_scale, bool is_causal, BTLA_DTYPE pv_dtype) { + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: mask and is_causal cannot both be set"); + } + if (seq_len_q != 1) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: only seq_len_q == 1 is supported in v1"); + } + if (seq_len_kv <= 0 || seq_len_kv_cache <= 0) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_decode_pvhalf: seq_len_kv and seq_len_kv_cache must be greater than 0"); + } + if (pv_dtype != BTLA_DTYPE::F16 && pv_dtype != BTLA_DTYPE::BF16) { + throw std::invalid_argument( + "sdpa_impl_qks8_sparse_decode_pvhalf: only F16 and BF16 are supported for V/O dtype"); + } + if (qscale == nullptr || kscale == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: qscale and kscale must be provided"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: lut and valid_block_num must be provided"); + } + if (K_cache_ptr == nullptr || V_cache_ptr == nullptr) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: K_cache and V_cache must be provided"); + } + if (num_q_blocks != 1 || num_k_blocks <= 0) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: only one Q block row is supported in v1"); + } + if (scale_block_size <= 0) { + throw std::invalid_argument("sdpa_impl_qks8_sparse_decode_pvhalf: scale_block_size must be greater than 0"); + } + + sparse_sage_decode(q, Q_ptr, K_ptr, V_ptr, K_cache_ptr, V_cache_ptr, O_ptr, mask, scale_block_size, qscale, kscale, + lut, valid_block_num, num_q_blocks, num_k_blocks, pv_dtype, q_stride_s, q_stride_d, q_stride_h, + q_stride_b, k_stride_s, k_stride_d, k_stride_h, k_stride_b, v_stride_d, v_stride_s, v_stride_h, + v_stride_b, o_stride_s, o_stride_d, o_stride_h, o_stride_b, batch, num_heads_q, num_heads_kv, + seq_len_q, seq_len_kv, seq_len_kv_cache, head_dim, softmax_scale, is_causal); +} + +} // namespace ark + +#endif // ARK_XPU && ARK_SYCL_TLA diff --git a/auto_round_extension/ark/auto_round_kernel/sparge_preprocess_triton.py b/auto_round_extension/ark/auto_round_kernel/sparge_preprocess_triton.py new file mode 100644 index 000000000..5306cf4e8 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/sparge_preprocess_triton.py @@ -0,0 +1,449 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import logging +from typing import Any, Callable + +import torch +import triton +import triton.language as tl + +logger = logging.getLogger(__name__) + +_TRITON_FALLBACK_WARNING_LOGGED = False + + +def _log_fallback_warning_once(error: Exception) -> None: + global _TRITON_FALLBACK_WARNING_LOGGED + + if _TRITON_FALLBACK_WARNING_LOGGED: + return + + logger.warning( + "ARK SPARGE preprocess Triton-XPU backend failed and fell back to torch. Subsequent failures will be suppressed. Error: %s", + error, + ) + _TRITON_FALLBACK_WARNING_LOGGED = True + + +def _normalize_backend_preference(backend_preference: str) -> str: + backend = backend_preference.lower() + if backend not in {"auto", "torch", "triton_xpu"}: + raise ValueError(f"Unsupported SPARGE preprocess backend: {backend_preference}") + return backend + + +def _ensure_triton_xpu_available(query: torch.Tensor, head_dim: int) -> None: + if query.device.type != "xpu": + raise RuntimeError("Triton-XPU preprocess requires XPU tensors") + if head_dim not in (64, 128): + raise ValueError(f"Unsupported head_dim={head_dim} for Triton-XPU preprocess") + + +def _sequence_mean_native_layout(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: + layout = tensor_layout.upper() + if layout == "HND": + return tensor.mean(dim=2).contiguous() + return tensor.mean(dim=1).contiguous() + + +def _slice_sequence_native_layout(tensor: torch.Tensor, tensor_layout: str, start: int, end: int) -> torch.Tensor: + layout = tensor_layout.upper() + if layout == "HND": + return tensor[:, :, start:end, :] + return tensor[:, start:end, :, :] + + +@triton.jit +def _triton_bmm_pool_sim_simmean_fuse_quant_xpu( + x_ptr, + xm_ptr, + pool_ptr, + sim_ptr, + x_quant_ptr, + scale_ptr, + simthreshd1_ptr, + N, + x_stride_b, + x_stride_s, + x_stride_h, + x_stride_d, + xq_stride_b, + xq_stride_s, + xq_stride_h, + xq_stride_d, + D: tl.constexpr, + BS: tl.constexpr, + FUSE_MEAN: tl.constexpr, +): + b = tl.program_id(0) + h = tl.program_id(1) + nb = tl.program_id(2) + H = tl.num_programs(1) + + row_ids = nb * BS + tl.arange(0, BS) + xmask = row_ids[:, None] < N + x_ptrs = ( + x_ptr + b * x_stride_b + row_ids[:, None] * x_stride_s + h * x_stride_h + tl.arange(0, D)[None, :] * x_stride_d + ) + x = tl.load(x_ptrs, mask=xmask, other=0.0) + valid_rows = N - nb * BS + bs_eff = tl.minimum(valid_rows, BS) + x_fp32 = x.to(tl.float32) + + if FUSE_MEAN: + xm_ptrs = xm_ptr + (b * H * D) + (h * D) + tl.arange(0, D) + x_mean = tl.load(xm_ptrs).to(tl.float32) + x_fp32 = x_fp32 - x_mean[None, :] + x_fp32 = tl.where(xmask, x_fp32, 0.0) + + cur_h1 = tl.load(simthreshd1_ptr + h) + pool = tl.sum(x_fp32, axis=0) / bs_eff + x_sq = x_fp32 * x_fp32 + x_norm = tl.sqrt(tl.sum(x_sq, axis=1, keep_dims=True)) + x_norm = tl.where(x_norm > 0, x_norm, 1.0) + x_normed = (x_fp32 / x_norm).to(tl.float16) + grams = tl.dot(x_normed, tl.trans(x_normed)) + sum_value = tl.sum(grams).to(tl.float32) + cur_sim = (sum_value / (bs_eff * bs_eff)) > cur_h1 + # Match the CUDA reference (SpargeAttn spas_sage_attn/utils.py): there the + # self-similarity divides zero-padded rows by a zero norm (0/0 -> NaN), so a + # partial tail block compares NaN > thr == False and is forced dense. Our + # guarded norm would instead mark the tail block similar/prunable, sparsifying + # the sequence tail (= last video frames). Force partial blocks to False. + cur_sim = cur_sim & (bs_eff >= BS) + + num_blocks = tl.cdiv(N, BS) + pool_block_offset = (b * H * num_blocks * D) + (h * num_blocks * D) + (nb * D) + tl.store(pool_ptr + pool_block_offset + tl.arange(0, D), pool) + sim_offset = (b * H * num_blocks) + (h * num_blocks) + nb + tl.store(sim_ptr + sim_offset, cur_sim) + + scale = tl.max(tl.abs(x_fp32)) / 127.0 + scale = scale + 1.0e-7 + x_int8 = x_fp32 / scale + x_int8 = x_int8 + 0.5 * tl.where(x_int8 >= 0, 1, -1) + x_int8 = x_int8.to(tl.int8) + x_quant_ptrs = ( + x_quant_ptr + + b * xq_stride_b + + row_ids[:, None] * xq_stride_s + + h * xq_stride_h + + tl.arange(0, D)[None, :] * xq_stride_d + ) + scale_ptrs = scale_ptr + (b * H * num_blocks) + (h * num_blocks) + nb + tl.store(x_quant_ptrs, x_int8, mask=xmask) + tl.store(scale_ptrs, scale) + + +def _safe_softmax(scores: torch.Tensor) -> torch.Tensor: + finite = torch.isfinite(scores) + safe_scores = torch.where(finite, scores, torch.full_like(scores, -1.0e9)) + probs = torch.softmax(safe_scores, dim=-1) + probs = torch.where(finite, probs, torch.zeros_like(probs)) + denom = probs.sum(dim=-1, keepdim=True) + return torch.where(denom > 0, probs / denom, torch.zeros_like(probs)) + + +def _build_block_causal_mask( + num_q_tiles: int, + num_k_blocks: int, + q_route_block_tokens: int, + k_route_block_tokens: int, + device: torch.device, +) -> torch.Tensor: + q_idx = torch.arange(num_q_tiles, device=device, dtype=torch.int64).view(-1, 1) + k_idx = torch.arange(num_k_blocks, device=device, dtype=torch.int64).view(1, -1) + valid_k_per_q = ((q_idx + 1) * q_route_block_tokens + k_route_block_tokens - 1) // k_route_block_tokens + return k_idx < valid_k_per_q + + +@triton.jit +def _triton_fill_block_map_kernel(final_map_ptr, num_to_select_ptr, sorted_indices_ptr, NK: tl.constexpr): + b = tl.program_id(0) + h = tl.program_id(1) + q = tl.program_id(2) + H = tl.num_programs(1) + Q = tl.num_programs(2) + + cur_num_to_select = tl.load(num_to_select_ptr + b * H * Q + h * Q + q) + sorted_row_ptr = sorted_indices_ptr + b * H * Q * NK + h * Q * NK + q * NK + final_row_ptr = final_map_ptr + b * H * Q * NK + h * Q * NK + q * NK + cur_num_to_select = tl.where(cur_num_to_select == 0, 1, cur_num_to_select) + added = 0 + for i in range(NK): + if added < cur_num_to_select: + cur_idx = tl.load(sorted_row_ptr + i) + cur_val = tl.load(final_row_ptr + cur_idx) + if cur_val == 0: + tl.store(final_row_ptr + cur_idx, 1) + added += 1 + + +def _fill_block_map_triton( + final_map: torch.Tensor, num_to_select: torch.Tensor, sorted_indices: torch.Tensor +) -> torch.Tensor: + final_map_u8 = final_map.contiguous().to(torch.uint8) + num_to_select = num_to_select.contiguous() + sorted_indices = sorted_indices.contiguous() + bsz, num_heads, num_q, num_k = final_map.shape + grid = (bsz, num_heads, num_q) + _triton_fill_block_map_kernel[grid](final_map_u8, num_to_select, sorted_indices, NK=num_k) + return final_map_u8.to(torch.bool) + + +@triton.jit +def _triton_block_map_to_lut_kernel(map_ptr, lut_ptr, valid_block_num_ptr, NK: tl.constexpr): + b = tl.program_id(0) + h = tl.program_id(1) + q = tl.program_id(2) + H = tl.num_programs(1) + Q = tl.num_programs(2) + + row_map_ptr = map_ptr + b * H * Q * NK + h * Q * NK + q * NK + row_lut_ptr = lut_ptr + b * H * Q * NK + h * Q * NK + q * NK + row_valid_ptr = valid_block_num_ptr + b * H * Q + h * Q + q + + valid_block_num = 0 + prev_block = 0 + for i in range(NK): + cur_block = tl.load(row_map_ptr + i) + if cur_block: + tl.store(row_lut_ptr + valid_block_num, i - prev_block) + valid_block_num += 1 + prev_block = i + tl.store(row_valid_ptr, valid_block_num) + + +def _block_map_lut_triton(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + block_map_u8 = block_map.contiguous().to(torch.uint8) + bsz, num_heads, num_q, num_k = block_map.shape + lut = torch.zeros((bsz, num_heads, num_q, num_k), dtype=torch.int32, device=block_map.device) + valid_block_num = torch.zeros((bsz, num_heads, num_q), dtype=torch.int32, device=block_map.device) + grid = (bsz, num_heads, num_q) + _triton_block_map_to_lut_kernel[grid](block_map_u8, lut, valid_block_num, NK=num_k) + return lut, valid_block_num + + +def _fill_block_map_torch( + final_map: torch.Tensor, num_to_select: torch.Tensor, sorted_indices: torch.Tensor +) -> torch.Tensor: + k_blocks = final_map.shape[-1] + filled = final_map.clone() + column_ids = torch.arange(k_blocks, device=final_map.device).view(1, 1, 1, k_blocks) + target_new = torch.maximum(num_to_select, torch.ones_like(num_to_select)) + added = torch.zeros_like(num_to_select) + for rank in range(k_blocks): + idx_match = column_ids == sorted_indices[..., rank : rank + 1] + is_new = idx_match & ~filled + should_add = (added < target_new).unsqueeze(-1) + newly_selected = should_add & is_new + filled |= newly_selected + added = added + newly_selected.any(dim=-1).to(added.dtype) + return filled + + +def _get_pool_sim_triton_simmean_fuse_quant( + x: torch.Tensor, + x_mean: torch.Tensor | None, + block_size: int, + simthreshd1: torch.Tensor, + tensor_layout: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + x = x.contiguous() + layout = tensor_layout.upper() + if layout == "HND": + bsz, num_heads, seq_len, head_dim = x.shape + x_stride_b, x_stride_h, x_stride_s, x_stride_d = x.stride() + else: + bsz, seq_len, num_heads, head_dim = x.shape + x_stride_b, x_stride_s, x_stride_h, x_stride_d = x.stride() + num_blocks = (seq_len + block_size - 1) // block_size + pool = torch.empty((bsz, num_heads, num_blocks, head_dim), device=x.device, dtype=x.dtype) + sim_u8 = torch.empty((bsz, num_heads, num_blocks), device=x.device, dtype=torch.uint8) + x_quant = torch.empty_like(x, dtype=torch.int8) + x_scale = torch.empty((bsz, num_heads, num_blocks), device=x.device, dtype=torch.float32) + if layout == "HND": + xq_stride_b, xq_stride_h, xq_stride_s, xq_stride_d = x_quant.stride() + else: + xq_stride_b, xq_stride_s, xq_stride_h, xq_stride_d = x_quant.stride() + mean = None if x_mean is None else x_mean.contiguous().squeeze(-2) + grid = (bsz, num_heads, num_blocks) + num_warps = 4 if block_size <= 64 else 8 + _triton_bmm_pool_sim_simmean_fuse_quant_xpu[grid]( + x, + mean, + pool, + sim_u8, + x_quant, + x_scale, + simthreshd1.contiguous(), + seq_len, + x_stride_b, + x_stride_s, + x_stride_h, + x_stride_d, + xq_stride_b, + xq_stride_s, + xq_stride_h, + xq_stride_d, + D=head_dim, + BS=block_size, + FUSE_MEAN=mean is not None, + num_warps=num_warps, + ) + return pool, sim_u8.to(torch.bool), x_quant, x_scale.unsqueeze(-1).contiguous() + + +def _run_triton_xpu_preprocess(ctx: Any) -> dict[str, Any]: + key_mean = _sequence_mean_native_layout(ctx.key, ctx.tensor_layout) if ctx.smooth_k else None + pooled_q, sim_qblocks, q_int8_hnd, q_scale = _get_pool_sim_triton_simmean_fuse_quant( + ctx.query, + None, + ctx.quant_block_size, + ctx.simthreshd1, + ctx.tensor_layout, + ) + pooled_k, sim_kblocks, k_int8_hnd, k_scale = _get_pool_sim_triton_simmean_fuse_quant( + ctx.key, + key_mean, + ctx.quant_block_size, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + ) + k_quant_gran = getattr(ctx, "k_quant_granularity", ctx.quant_block_size) + if k_quant_gran > ctx.quant_block_size: + _, _, k_int8_hnd, k_scale = _get_pool_sim_triton_simmean_fuse_quant( + ctx.key, + key_mean, + k_quant_gran, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + ) + ratio = k_quant_gran // ctx.quant_block_size + num_k_blocks_fine = (ctx.seq_len_kv + ctx.quant_block_size - 1) // ctx.quant_block_size + k_scale = k_scale.repeat_interleave(ratio, dim=2)[:, :, :num_k_blocks_fine, :] + if ctx.q_blocks_per_tile > 1: + pooled_q_for_routing, sim_q_for_routing, _, _ = _get_pool_sim_triton_simmean_fuse_quant( + ctx.query, + None, + ctx.query_tile_tokens, + ctx.simthreshd1, + ctx.tensor_layout, + ) + else: + pooled_q_for_routing = pooled_q + sim_q_for_routing = sim_qblocks + + if ctx.k_blocks_per_tile > 1: + pooled_k_for_routing, sim_k_for_routing, _, _ = _get_pool_sim_triton_simmean_fuse_quant( + ctx.key, + key_mean, + ctx.k_route_block_tokens, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + ) + else: + pooled_k_for_routing = pooled_k + sim_k_for_routing = sim_kblocks + + kv_head_index = torch.arange(ctx.num_heads_q, device=ctx.query.device, dtype=torch.int64) // ( + ctx.num_heads_q // ctx.num_heads_kv + ) + pooled_k_for_q = pooled_k_for_routing[:, kv_head_index] + sim_k_for_q = sim_k_for_routing[:, kv_head_index] + sim_k_expand = sim_k_for_q.unsqueeze(-2).expand(-1, -1, ctx.num_q_tiles, -1) + sim_q_expand = sim_q_for_routing.unsqueeze(-1).expand(-1, -1, -1, pooled_k_for_routing.size(2)) + + pooled_score = torch.matmul( + pooled_q_for_routing.to(torch.float32), + pooled_k_for_q.transpose(-1, -2).to(torch.float32), + ) + pooled_score *= ctx.head_dim**-0.5 + pooled_score = pooled_score.masked_fill(~sim_k_expand, -torch.inf) + if ctx.is_causal: + causal_mask = _build_block_causal_mask( + ctx.num_q_tiles, + pooled_k_for_routing.size(2), + ctx.q_route_block_tokens, + ctx.k_route_block_tokens, + ctx.query.device, + ) + pooled_score = pooled_score.masked_fill( + ~causal_mask.view(1, 1, ctx.num_q_tiles, pooled_k_for_routing.size(2)), -torch.inf + ) + else: + causal_mask = None + + pooled_prob = _safe_softmax(pooled_score) + sorted_prob = torch.sort(pooled_prob, dim=-1, descending=True) + _, _, _, num_k_route_blocks = pooled_prob.shape + num_to_select = ( + (ctx.topk.view(1, ctx.num_heads_q, 1) * num_k_route_blocks) + .to(torch.int64) + .expand(ctx.batch, -1, ctx.num_q_tiles) + .contiguous() + ) + final_tile_map = torch.zeros_like(pooled_prob, dtype=torch.bool) + final_tile_map[~sim_k_expand] = True + final_tile_map[~sim_q_expand] = True + final_tile_map = _fill_block_map_triton(final_tile_map, num_to_select, sorted_prob.indices) + if causal_mask is not None: + final_tile_map &= causal_mask.view(1, 1, ctx.num_q_tiles, num_k_route_blocks) + if ctx.attention_sink: + final_tile_map[..., 0] = True + + q_block_to_tile = ( + torch.arange(ctx.num_sparse_q_blocks, device=ctx.query.device, dtype=torch.int64) + // ctx.q_sparse_blocks_per_tile + ) + q_block_to_tile = q_block_to_tile.clamp_max(ctx.num_q_tiles - 1) + k_block_to_tile = ( + torch.arange(ctx.num_sparse_k_blocks, device=ctx.query.device, dtype=torch.int64) + // ctx.k_sparse_blocks_per_tile + ) + k_block_to_tile = k_block_to_tile.clamp_max(ctx.num_k_tiles - 1) + raw_block_map = final_tile_map.index_select(2, q_block_to_tile).index_select(3, k_block_to_tile).contiguous() + lut, valid_block_num = _block_map_lut_triton(raw_block_map) + + return { + "query_i8": q_int8_hnd, + "key_i8": k_int8_hnd, + "qscale": q_scale, + "kscale": k_scale, + "lut": lut, + "valid_block_num": valid_block_num, + "raw_block_map": raw_block_map, + "tile_block_map": final_tile_map.contiguous(), + "sim_qblocks": sim_q_for_routing.contiguous(), + "sim_kblocks": sim_k_for_routing.contiguous(), + "backend": "triton_xpu", + } + + +def dispatch_sparge_preprocess_backend( + *, + ctx: Any, + torch_backend: Callable[[], dict[str, Any]], + backend_preference: str = "auto", +) -> dict[str, Any]: + backend = _normalize_backend_preference(backend_preference) + if backend == "torch": + result = torch_backend() + result["backend"] = "torch" + return result + + try: + _ensure_triton_xpu_available(ctx.query, ctx.head_dim) + result = _run_triton_xpu_preprocess(ctx) + result["backend"] = "triton_xpu" + return result + except (NotImplementedError, RuntimeError, ValueError, triton.runtime.errors.TritonError) as error: + if backend == "triton_xpu": + raise + _log_fallback_warning_once(error) + result = torch_backend() + result["backend"] = "torch" + return result diff --git a/auto_round_extension/ark/auto_round_kernel/sparse_attention.py b/auto_round_extension/ark/auto_round_kernel/sparse_attention.py new file mode 100644 index 000000000..427353d46 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/sparse_attention.py @@ -0,0 +1,1621 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +import warnings +from dataclasses import dataclass +from typing import Any, Optional + +import torch + +from . import ( + _attention_strides_qko, + _attention_strides_v, + _empty_attention_output, + _normalize_tensor_layout, + _validate_attention_tensor, + cvt_dtype, + get_lib, + get_stream, +) +from . import sagev1 as _dense_sagev1 +from . import sagev1_pvi8 as _dense_sagev1_pvi8 +from . import ( + sdpa, +) + + +def _get_xpu_sparse_kernel_backend() -> str: + backend = os.getenv("SAGE_ATTN_XPU_SPARSE_KERNEL_BACKEND", "sycl").strip().lower() + if backend != "sycl": + raise ValueError("Unsupported SAGE_ATTN_XPU_SPARSE_KERNEL_BACKEND=" f"{backend!r}; expected: sycl") + return backend + + +def _get_sparse_preprocess_backend_preference() -> str: + return os.getenv("SAGE_ATTN_SPARSE_PREPROCESS_BACKEND", "auto").strip().lower() + + +def sage_sparse( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + lut: torch.Tensor, + valid_block_num: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + quant_block_size: int = 64, + qscale: torch.Tensor = None, + kscale: torch.Tensor = None, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, + tensor_layout: str = "HND", +) -> torch.Tensor: + """Low-level sparse SAGE attention with pre-quantized INT8 Q/K and LUT-driven block selection.""" + if query.device.type != "xpu": + raise NotImplementedError("sage_sparse is only supported on XPU") + if query.dtype != torch.int8 or key.dtype != torch.int8: + raise ValueError(f"Q/K must be int8, got Q={query.dtype}, K={key.dtype}") + if value.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"V must be float16 or bfloat16, got {value.dtype}") + if qscale is None or kscale is None: + raise ValueError("qscale and kscale must be provided for sage_sparse") + if q_tile_override not in (0, 64, 128, 256): + raise ValueError(f"Unsupported q_tile_override={q_tile_override}; supported values: 0, 64, 128, 256") + if lut.dtype != torch.int32 or valid_block_num.dtype != torch.int32: + raise ValueError("lut and valid_block_num must be int32 tensors") + if lut.device != query.device or valid_block_num.device != query.device: + raise ValueError("lut and valid_block_num must be on the same XPU device as Q/K/V") + if qscale.device != query.device or kscale.device != query.device: + raise ValueError("qscale and kscale must be on the same XPU device as Q/K/V") + if quant_block_size <= 0: + raise ValueError(f"quant_block_size must be positive, got {quant_block_size}") + + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout) + Bv, Hkv2, Skv2, Dv = _validate_attention_tensor(value, "V", tensor_layout) + + if Bk != B or Bv != B: + raise ValueError("Batch size mismatch between Q/K/V") + if Hkv2 != Hkv or Skv2 != Skv or Dv != Dk: + raise ValueError("K/V shape mismatch") + if Dk != D: + raise ValueError("Head dim mismatch between Q and K/V") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + _validate_sparse_q_tile_override_for_head_dim(D, q_tile_override) + + effective_sparse_q_block_tokens = quant_block_size if sparse_q_block_tokens is None else int(sparse_q_block_tokens) + effective_sparse_k_block_tokens = quant_block_size if sparse_k_block_tokens is None else int(sparse_k_block_tokens) + if effective_sparse_q_block_tokens <= 0 or effective_sparse_k_block_tokens <= 0: + raise ValueError( + "sparse_q_block_tokens and sparse_k_block_tokens must be positive when provided; " + f"got {effective_sparse_q_block_tokens} and {effective_sparse_k_block_tokens}" + ) + + q_scale_blocks = (Sq + quant_block_size - 1) // quant_block_size + kv_scale_blocks = (Skv + quant_block_size - 1) // quant_block_size + q_sparse_blocks = (Sq + effective_sparse_q_block_tokens - 1) // effective_sparse_q_block_tokens + kv_sparse_blocks = (Skv + effective_sparse_k_block_tokens - 1) // effective_sparse_k_block_tokens + decoupled_sparse_rows = ( + effective_sparse_q_block_tokens != quant_block_size or effective_sparse_k_block_tokens != quant_block_size + ) + if tuple(lut.shape) != (B, Hq, q_sparse_blocks, kv_sparse_blocks): + raise ValueError(f"lut must have shape {(B, Hq, q_sparse_blocks, kv_sparse_blocks)}, got {tuple(lut.shape)}") + if tuple(valid_block_num.shape) != (B, Hq, q_sparse_blocks): + raise ValueError( + f"valid_block_num must have shape {(B, Hq, q_sparse_blocks)}, got {tuple(valid_block_num.shape)}" + ) + if qscale.numel() != B * Hq * q_scale_blocks: + raise ValueError( + f"qscale must have {B * Hq * q_scale_blocks} elements for shape [B, Hq, ceil(Sq/block), 1], got {qscale.numel()}" + ) + if kscale.numel() != B * Hkv * kv_scale_blocks: + raise ValueError( + f"kscale must have {B * Hkv * kv_scale_blocks} elements for shape [B, Hkv, ceil(Skv/block), 1], got {kscale.numel()}" + ) + if torch.any(valid_block_num < 0).item(): + raise ValueError("valid_block_num entries must be non-negative") + if torch.any(valid_block_num > kv_sparse_blocks).item(): + raise ValueError(f"valid_block_num entries must be <= {kv_sparse_blocks}") + if decoupled_sparse_rows and not _is_sparse_qtile256_row64k_config( + head_dim=D, + quant_block_size=quant_block_size, + q_tile_override=q_tile_override, + sparse_q_block_tokens=effective_sparse_q_block_tokens, + sparse_k_block_tokens=effective_sparse_k_block_tokens, + tensor_layout=tensor_layout, + ): + raise ValueError( + "Only the decoupled sparse config " + "(head_dim=128, quant_block_size=64, q_tile_override=256, " + "sparse_q_block_tokens=256, sparse_k_block_tokens=64, tensor_layout in {'HND', 'NHD'}) is supported" + ) + + lib = get_lib(query) + stream = get_stream(query) + O = _empty_attention_output( + B, + Hq, + Sq, + D, + dtype=value.dtype, + device=query.device, + tensor_layout=tensor_layout, + ) + q_strides = _attention_strides_qko(query, tensor_layout) + k_strides = _attention_strides_qko(key, tensor_layout) + v_strides = _attention_strides_v(value, tensor_layout) + o_strides = _attention_strides_qko(O, tensor_layout) + sparse_fn_name = "sage_sparse_qtile256_row64k" if decoupled_sparse_rows else "sage_sparse" + if not hasattr(lib, sparse_fn_name): + raise RuntimeError(f"Loaded XPU extension does not expose {sparse_fn_name}") + sparse_fn = getattr(lib, sparse_fn_name) + sparse_fn( + stream, + query.data_ptr(), + key.data_ptr(), + value.data_ptr(), + O.data_ptr(), + attn_mask.data_ptr() if attn_mask is not None else 0, + quant_block_size, + qscale.data_ptr(), + kscale.data_ptr(), + lut.data_ptr(), + valid_block_num.data_ptr(), + q_sparse_blocks, + kv_sparse_blocks, + q_tile_override, + *q_strides, + *k_strides, + *v_strides, + *o_strides, + cvt_dtype(query.dtype), + cvt_dtype(key.dtype), + cvt_dtype(O.dtype), + B, + Hq, + Hkv, + Sq, + Skv, + D, + float(scale) if scale is not None else 1.0 / (D**0.5), + bool(is_causal), + ) + return O + + +def sage_sparse_row_linear( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + lut: torch.Tensor, + valid_block_num: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + quant_block_size: int = 64, + qscale: torch.Tensor = None, + kscale: torch.Tensor = None, + q_tile_override: int = 64, + tensor_layout: str = "HND", +) -> torch.Tensor: + """Sparse SAGE prefill with one sparse row per workgroup via the q_tile=64 launcher.""" + if quant_block_size != 64: + raise ValueError(f"sage_sparse_row_linear currently requires quant_block_size == 64, got {quant_block_size}") + if q_tile_override not in (0, 64): + raise ValueError(f"Unsupported q_tile_override={q_tile_override}; row-linear backend supports only 0 or 64") + if query.device.type != "xpu": + raise NotImplementedError("sage_sparse_row_linear is only supported on XPU") + if query.dtype != torch.int8 or key.dtype != torch.int8: + raise ValueError(f"Q/K must be int8, got Q={query.dtype}, K={key.dtype}") + if value.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"V must be float16 or bfloat16, got {value.dtype}") + if qscale is None or kscale is None: + raise ValueError("qscale and kscale must be provided for sage_sparse_row_linear") + if lut.dtype != torch.int32 or valid_block_num.dtype != torch.int32: + raise ValueError("lut and valid_block_num must be int32 tensors") + if lut.device != query.device or valid_block_num.device != query.device: + raise ValueError("lut and valid_block_num must be on the same XPU device as Q/K/V") + if qscale.device != query.device or kscale.device != query.device: + raise ValueError("qscale and kscale must be on the same XPU device as Q/K/V") + + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout) + Bv, Hkv2, Skv2, Dv = _validate_attention_tensor(value, "V", tensor_layout) + + if Bk != B or Bv != B: + raise ValueError("Batch size mismatch between Q/K/V") + if Hkv2 != Hkv or Skv2 != Skv or Dv != Dk: + raise ValueError("K/V shape mismatch") + if Dk != D: + raise ValueError("Head dim mismatch between Q and K/V") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + + q_blocks = (Sq + quant_block_size - 1) // quant_block_size + kv_blocks = (Skv + quant_block_size - 1) // quant_block_size + if tuple(lut.shape) != (B, Hq, q_blocks, kv_blocks): + raise ValueError(f"lut must have shape {(B, Hq, q_blocks, kv_blocks)}, got {tuple(lut.shape)}") + if tuple(valid_block_num.shape) != (B, Hq, q_blocks): + raise ValueError(f"valid_block_num must have shape {(B, Hq, q_blocks)}, got {tuple(valid_block_num.shape)}") + if qscale.numel() != B * Hq * q_blocks: + raise ValueError( + f"qscale must have {B * Hq * q_blocks} elements for shape [B, Hq, ceil(Sq/block), 1], got {qscale.numel()}" + ) + if kscale.numel() != B * Hkv * kv_blocks: + raise ValueError( + f"kscale must have {B * Hkv * kv_blocks} elements for shape [B, Hkv, ceil(Skv/block), 1], got {kscale.numel()}" + ) + if torch.any(valid_block_num < 0).item(): + raise ValueError("valid_block_num entries must be non-negative") + if torch.any(valid_block_num > kv_blocks).item(): + raise ValueError(f"valid_block_num entries must be <= {kv_blocks}") + + lib = get_lib(query) + stream = get_stream(query) + O = _empty_attention_output( + B, + Hq, + Sq, + D, + dtype=value.dtype, + device=query.device, + tensor_layout=tensor_layout, + ) + q_strides = _attention_strides_qko(query, tensor_layout) + k_strides = _attention_strides_qko(key, tensor_layout) + v_strides = _attention_strides_v(value, tensor_layout) + o_strides = _attention_strides_qko(O, tensor_layout) + lib.sage_sparse_row_linear( + stream, + query.data_ptr(), + key.data_ptr(), + value.data_ptr(), + O.data_ptr(), + attn_mask.data_ptr() if attn_mask is not None else 0, + quant_block_size, + qscale.data_ptr(), + kscale.data_ptr(), + lut.data_ptr(), + valid_block_num.data_ptr(), + q_blocks, + kv_blocks, + q_tile_override, + *q_strides, + *k_strides, + *v_strides, + *o_strides, + cvt_dtype(query.dtype), + cvt_dtype(key.dtype), + cvt_dtype(O.dtype), + B, + Hq, + Hkv, + Sq, + Skv, + D, + float(scale) if scale is not None else 1.0 / (D**0.5), + bool(is_causal), + ) + return O + + +def sage_sparse_decode( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + lut: torch.Tensor, + valid_block_num: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + quant_block_size: int = 64, + qscale: torch.Tensor | None = None, + kscale: torch.Tensor | None = None, + kscale_cache: torch.Tensor | None = None, + tensor_layout: str = "HND", +) -> torch.Tensor: + """Low-level sparse SAGE decode with pre-quantized INT8 Q/K and explicit KV cache.""" + if query.device.type != "xpu": + raise NotImplementedError("sage_sparse_decode is only supported on XPU") + if dropout_p != 0.0: + raise NotImplementedError("dropout_p must be 0.0 for sage_sparse_decode") + if enable_gqa: + raise NotImplementedError("enable_gqa is not used by sage_sparse_decode; provide explicit Hq/Hkv tensors") + if query.dtype != torch.int8 or key.dtype != torch.int8 or key_cache.dtype != torch.int8: + raise ValueError(f"Q/K/K_cache must be int8, got Q={query.dtype}, K={key.dtype}, K_cache={key_cache.dtype}") + if value.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"V must be float16 or bfloat16, got {value.dtype}") + if value_cache.dtype != value.dtype: + raise ValueError(f"V_cache dtype must match V dtype, got {value_cache.dtype} vs {value.dtype}") + if qscale is None or kscale is None or kscale_cache is None: + raise ValueError("qscale, kscale, and kscale_cache must be provided for sage_sparse_decode") + if lut.dtype != torch.int32 or valid_block_num.dtype != torch.int32: + raise ValueError("lut and valid_block_num must be int32 tensors") + if quant_block_size <= 0: + raise ValueError(f"quant_block_size must be positive, got {quant_block_size}") + + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout) + Bkc, Hkvc, Skvc, Dkc = _validate_attention_tensor(key_cache, "K_cache", tensor_layout) + Bv, Hkv2, Skv2, Dv = _validate_attention_tensor(value, "V", tensor_layout) + Bvc, Hkvc2, Skvc2, Dvc = _validate_attention_tensor(value_cache, "V_cache", tensor_layout) + if Bk != B or Bkc != B or Bv != B or Bvc != B: + raise ValueError("Batch size mismatch between Q/K/V/cache tensors") + if Hkv2 != Hkv or Hkvc != Hkv or Hkvc2 != Hkv: + raise ValueError("KV head-count mismatch") + if Skv2 != Skv or Skvc2 != Skvc: + raise ValueError("K/V sequence length mismatch") + if Dk != D or Dkc != D or Dv != D or Dvc != D: + raise ValueError("Head dim mismatch between Q/K/V/cache tensors") + if Sq != 1: + raise ValueError(f"sage_sparse_decode currently supports only seq_len_q == 1, got {Sq}") + if Skv <= 0 or Skvc <= 0: + raise ValueError("sage_sparse_decode requires both current-step KV and KV cache to be non-empty") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + + q_blocks = (Sq + quant_block_size - 1) // quant_block_size + cur_blocks = (Skv + quant_block_size - 1) // quant_block_size + cache_blocks = (Skvc + quant_block_size - 1) // quant_block_size + total_blocks = cur_blocks + cache_blocks + if tuple(lut.shape) != (B, Hq, q_blocks, total_blocks): + raise ValueError(f"lut must have shape {(B, Hq, q_blocks, total_blocks)}, got {tuple(lut.shape)}") + if tuple(valid_block_num.shape) != (B, Hq, q_blocks): + raise ValueError(f"valid_block_num must have shape {(B, Hq, q_blocks)}, got {tuple(valid_block_num.shape)}") + if qscale.numel() != B * Hq * q_blocks: + raise ValueError(f"qscale must have {B * Hq * q_blocks} elements, got {qscale.numel()}") + if kscale.numel() != B * Hkv * cur_blocks: + raise ValueError(f"kscale must have {B * Hkv * cur_blocks} elements, got {kscale.numel()}") + if kscale_cache.numel() != B * Hkv * cache_blocks: + raise ValueError(f"kscale_cache must have {B * Hkv * cache_blocks} elements, got {kscale_cache.numel()}") + if torch.any(valid_block_num < 0).item(): + raise ValueError("valid_block_num entries must be non-negative") + if torch.any(valid_block_num > total_blocks).item(): + raise ValueError(f"valid_block_num entries must be <= {total_blocks}") + + q_hnd = _to_hnd(query, tensor_layout) + k_hnd = _to_hnd(key, tensor_layout) + v_hnd = _to_hnd(value, tensor_layout) + k_cache_hnd = _to_hnd(key_cache, tensor_layout) + v_cache_hnd = _to_hnd(value_cache, tensor_layout) + O_hnd = torch.empty((B, Hq, Sq, D), dtype=value.dtype, device=query.device) + q_strides = _attention_strides_qko(q_hnd, "HND") + k_strides = _attention_strides_qko(k_hnd, "HND") + v_strides = _attention_strides_v(v_hnd, "HND") + o_strides = _attention_strides_qko(O_hnd, "HND") + kscale_total = torch.cat( + [kscale_cache.reshape(B, Hkv, cache_blocks, 1), kscale.reshape(B, Hkv, cur_blocks, 1)], dim=2 + ) + lib = get_lib(query) + stream = get_stream(query) + lib.sage_sparse_decode( + stream, + q_hnd.data_ptr(), + k_hnd.data_ptr(), + v_hnd.data_ptr(), + k_cache_hnd.data_ptr(), + v_cache_hnd.data_ptr(), + O_hnd.data_ptr(), + attn_mask.data_ptr() if attn_mask is not None else 0, + quant_block_size, + qscale.data_ptr(), + kscale_total.data_ptr(), + lut.data_ptr(), + valid_block_num.data_ptr(), + q_blocks, + total_blocks, + *q_strides, + *k_strides, + *v_strides, + *o_strides, + cvt_dtype(q_hnd.dtype), + cvt_dtype(k_hnd.dtype), + cvt_dtype(O_hnd.dtype), + B, + Hq, + Hkv, + Sq, + Skv, + Skvc, + D, + float(scale) if scale is not None else 1.0 / (D**0.5), + bool(is_causal), + ) + return _from_hnd(O_hnd, tensor_layout) + + +def sagev1( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + quant_block_size: int = 64, + tensor_layout: str = "HND", +) -> torch.Tensor: + return _dense_sagev1( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + ) + + +def sagev1_pvi8( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, + quant_block_size: int = 64, + tensor_layout: str = "HND", +) -> torch.Tensor: + return _dense_sagev1_pvi8( + query=query, + key=key, + value=value, + attn_mask=attn_mask, + dropout_p=dropout_p, + is_causal=is_causal, + scale=scale, + enable_gqa=enable_gqa, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + ) + + +def sageattn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + tensor_layout: str = "HND", + is_causal: bool = False, + sm_scale: Optional[float] = None, + return_lse: bool = False, + kernel: str = "v1_pvhalf", + **kwargs, +) -> torch.Tensor: + """SAGE attention dispatcher. + + Signature mirrors ``sageattention.sageattn``. + + Args: + - q, k, v: Query/Key/Value tensors. Layout selected by ``tensor_layout``. + - tensor_layout: "HND" or "NHD". + - is_causal: Whether to apply causal mask. + - sm_scale: Softmax scale. Uses ``1 / sqrt(head_dim)`` when None. + - return_lse: Not supported; must be False. + - kernel: Which SAGE variant to dispatch to. + - "v1_pvhalf" (default): PV in half precision (calls ``sagev1``). + - "v1_pvi8": PV in INT8 precision (calls ``sagev1_pvi8``). + - kwargs: Forwarded to the underlying kernel (e.g. ``attn_mask``, + ``dropout_p``, ``enable_gqa``, ``quant_block_size``). + + Returns: + - O: same layout as the input tensors. + """ + if return_lse: + raise NotImplementedError("return_lse is not supported in ARK sageattn") + + if kernel == "v1_pvhalf": + impl = sagev1 + elif kernel == "v1_pvi8": + impl = sagev1_pvi8 + else: + raise ValueError(f"Unsupported sageattn kernel={kernel!r}; supported: 'v1_pvhalf', 'v1_pvi8'") + + return impl( + query=q, + key=k, + value=v, + is_causal=is_causal, + scale=sm_scale, + tensor_layout=tensor_layout, + **kwargs, + ) + + +def _normalize_sparse_mask( + attn_mask: torch.Tensor | None, + batch: int, + seq_q: int, + seq_kv: int, + device: torch.device, +) -> torch.Tensor | None: + if attn_mask is None: + return None + if attn_mask.dtype == torch.bool: + raise ValueError("Boolean attention masks are not supported") + + mask = attn_mask + if mask.ndim == 2: + if mask.shape != (seq_q, seq_kv): + raise ValueError(f"Unsupported 2D attention mask shape {tuple(mask.shape)}") + mask = mask.view(1, 1, seq_q, seq_kv) + elif mask.ndim == 3: + if mask.shape != (batch, seq_q, seq_kv): + raise ValueError(f"Unsupported 3D attention mask shape {tuple(mask.shape)}") + mask = mask.unsqueeze(1) + elif mask.ndim == 4: + if mask.shape[-2:] != (seq_q, seq_kv): + raise ValueError(f"Unsupported 4D attention mask shape {tuple(mask.shape)}") + if mask.shape[1] != 1: + raise ValueError("Only attention masks with head dimension 1 are supported") + if mask.shape[0] == 1 and batch != 1: + mask = mask.expand(batch, -1, -1, -1) + elif mask.shape[0] != batch: + raise ValueError(f"Unsupported attention mask batch dimension {mask.shape[0]} != {batch}") + else: + raise ValueError(f"Unsupported attention mask rank {mask.ndim}") + + return mask.contiguous().to(device=device, dtype=torch.float32) + + +def _normalize_per_head_hparam( + value: float | int | torch.Tensor, + num_heads: int, + device: torch.device, + name: str, +) -> torch.Tensor: + if torch.is_tensor(value): + out = value.to(device=device, dtype=torch.float32).flatten() + if out.numel() == 1: + out = out.expand(num_heads) + elif out.numel() != num_heads: + raise ValueError(f"{name} must have 1 or {num_heads} elements, got {out.numel()}") + return out.contiguous() + return torch.full((num_heads,), float(value), device=device, dtype=torch.float32) + + +def _query_tile_tokens_for_head_dim(head_dim: int) -> int: + if head_dim == 64: + return 128 + if head_dim == 128: + return 64 + raise ValueError(f"Unsupported head_dim={head_dim}; supported: 64, 128") + + +def _normalize_query_tile_tokens( + query_tile_tokens: int | None, + *, + head_dim: int, + quant_block_size: int, +) -> int: + if query_tile_tokens is None: + return _query_tile_tokens_for_head_dim(head_dim) + + tokens = int(query_tile_tokens) + if tokens not in (64, 128, 256): + raise ValueError(f"query_tile_tokens={tokens} is not supported; supported values: 64, 128, 256") + if tokens <= 0 or tokens % quant_block_size != 0: + raise ValueError( + f"query_tile_tokens={tokens} must be a positive multiple of quant_block_size={quant_block_size}" + ) + return tokens + + +def _validate_sparse_q_tile_override_for_head_dim(head_dim: int, q_tile_override: int) -> None: + if head_dim == 128: + supported = (0, 64, 128, 256) + elif head_dim == 64: + supported = (0, 128) + else: + raise ValueError(f"Unsupported head_dim={head_dim}; supported: 64, 128") + + if q_tile_override not in supported: + raise ValueError( + f"q_tile_override={q_tile_override} is not supported for head_dim={head_dim}; " + f"supported values: {', '.join(str(v) for v in supported)}" + ) + + +def _routing_k_block_tokens_for_head_dim(head_dim: int, quant_block_size: int) -> int: + if head_dim == 64: + return quant_block_size + if head_dim == 128: + return 128 + raise ValueError(f"Unsupported head_dim={head_dim}; supported: 64, 128") + + +def _normalize_sparse_q_block_tokens( + sparse_q_block_tokens: int | None, + *, + quant_block_size: int, + q_route_block_tokens: int, +) -> int: + tokens = quant_block_size if sparse_q_block_tokens is None else int(sparse_q_block_tokens) + if tokens not in (64, 128, 256): + raise ValueError(f"sparse_q_block_tokens={tokens} is not supported; supported values: 64, 128, 256") + if tokens <= 0 or tokens % quant_block_size != 0: + raise ValueError( + f"sparse_q_block_tokens={tokens} must be a positive multiple of quant_block_size={quant_block_size}" + ) + if q_route_block_tokens % tokens != 0: + raise ValueError( + f"q_route_block_tokens={q_route_block_tokens} must be divisible by sparse_q_block_tokens={tokens}" + ) + return tokens + + +def _normalize_sparse_k_block_tokens( + sparse_k_block_tokens: int | None, + *, + quant_block_size: int, + k_route_block_tokens: int, +) -> int: + tokens = quant_block_size if sparse_k_block_tokens is None else int(sparse_k_block_tokens) + if tokens not in (64, 128): + raise ValueError(f"sparse_k_block_tokens={tokens} is not supported; supported values: 64, 128") + if tokens <= 0 or tokens % quant_block_size != 0: + raise ValueError( + f"sparse_k_block_tokens={tokens} must be a positive multiple of quant_block_size={quant_block_size}" + ) + if k_route_block_tokens % tokens != 0: + raise ValueError( + f"k_route_block_tokens={k_route_block_tokens} must be divisible by sparse_k_block_tokens={tokens}" + ) + return tokens + + +def _is_sparse_qtile256_row64k_config( + *, + head_dim: int, + quant_block_size: int, + q_tile_override: int, + sparse_q_block_tokens: int, + sparse_k_block_tokens: int, + tensor_layout: str, +) -> bool: + return ( + head_dim == 128 + and quant_block_size == 64 + and q_tile_override == 256 + and sparse_q_block_tokens == 256 + and sparse_k_block_tokens == 64 + and _normalize_tensor_layout(tensor_layout) in {"HND", "NHD"} + ) + + +def _sequence_mean_native_layout(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: + layout = _normalize_tensor_layout(tensor_layout) + seq_dim = 2 if layout == "HND" else 1 + return tensor.mean(dim=seq_dim).contiguous() + + +def _slice_sequence_native_layout(tensor: torch.Tensor, tensor_layout: str, start: int, end: int) -> torch.Tensor: + layout = _normalize_tensor_layout(tensor_layout) + if layout == "HND": + return tensor[:, :, start:end, :] + return tensor[:, start:end, :, :] + + +def _to_hnd(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: + layout = _normalize_tensor_layout(tensor_layout) + if layout == "HND": + return tensor.contiguous() + return tensor.permute(0, 2, 1, 3).contiguous() + + +def _from_hnd(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: + layout = _normalize_tensor_layout(tensor_layout) + if layout == "HND": + return tensor.contiguous() + return tensor.permute(0, 2, 1, 3).contiguous() + + +def _safe_softmax(scores: torch.Tensor) -> torch.Tensor: + finite = torch.isfinite(scores) + safe_scores = torch.where(finite, scores, torch.full_like(scores, -1.0e9)) + probs = torch.softmax(safe_scores, dim=-1) + probs = torch.where(finite, probs, torch.zeros_like(probs)) + denom = probs.sum(dim=-1, keepdim=True) + return torch.where(denom > 0, probs / denom, torch.zeros_like(probs)) + + +def _build_block_causal_mask( + num_q_tiles: int, + num_k_blocks: int, + q_route_block_tokens: int, + k_route_block_tokens: int, + device: torch.device, +) -> torch.Tensor: + q_idx = torch.arange(num_q_tiles, device=device, dtype=torch.int64).view(-1, 1) + k_idx = torch.arange(num_k_blocks, device=device, dtype=torch.int64).view(1, -1) + valid_k_per_q = ((q_idx + 1) * q_route_block_tokens + k_route_block_tokens - 1) // k_route_block_tokens + return k_idx < valid_k_per_q + + +def _fill_block_map_torch( + final_map: torch.Tensor, num_to_select: torch.Tensor, sorted_indices: torch.Tensor +) -> torch.Tensor: + k_blocks = final_map.shape[-1] + filled = final_map.clone() + column_ids = torch.arange(k_blocks, device=final_map.device).view(1, 1, 1, k_blocks) + target_new = torch.maximum(num_to_select, torch.ones_like(num_to_select)) + added = torch.zeros_like(num_to_select) + for rank in range(k_blocks): + idx_match = column_ids == sorted_indices[..., rank : rank + 1] + is_new = idx_match & ~filled + should_add = (added < target_new).unsqueeze(-1) + newly_selected = should_add & is_new + filled |= newly_selected + added = added + newly_selected.any(dim=-1).to(added.dtype) + return filled + + +def _block_map_lut_torch(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + valid_block_num = block_map.to(torch.int32).sum(dim=-1) + _, _, _, num_k_blocks = block_map.shape + one_matrix = torch.ones(block_map.shape, dtype=torch.int32, device=block_map.device) + cum_matrix = torch.cumsum(one_matrix, dim=-1) + masked_cum_matrix = cum_matrix * block_map.to(torch.int32) + filled_matrix = masked_cum_matrix.clone() + filled_matrix[~block_map] = 10_000_000 + lut = torch.sort(filled_matrix, dim=-1)[0] - 1 + lut[..., 1:] = lut[..., 1:] - lut[..., :-1] + invalid_mask = torch.arange(num_k_blocks, device=block_map.device).view( + 1, 1, 1, num_k_blocks + ) >= valid_block_num.unsqueeze(-1) + lut = torch.where(invalid_mask, torch.zeros_like(lut), lut) + return lut.to(torch.int32).contiguous(), valid_block_num.to(torch.int32).contiguous() + + +def _pool_sim_and_quant_torch( + x: torch.Tensor, + block_size: int, + sim_threshold: torch.Tensor, + tensor_layout: str, + mean_subtract: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + layout = _normalize_tensor_layout(tensor_layout) + if layout == "HND": + bsz, num_heads, seq_len, head_dim = x.shape + else: + bsz, seq_len, num_heads, head_dim = x.shape + num_blocks = (seq_len + block_size - 1) // block_size + pad_tokens = num_blocks * block_size - seq_len + if layout == "HND": + x_pad = torch.nn.functional.pad(x, (0, 0, 0, pad_tokens)) + x_blocks = x_pad.view(bsz, num_heads, num_blocks, block_size, head_dim).to(torch.float32) + else: + if pad_tokens: + pad = torch.zeros((bsz, pad_tokens, num_heads, head_dim), dtype=x.dtype, device=x.device) + x_pad = torch.cat([x, pad], dim=1) + else: + x_pad = x + x_blocks = x_pad.view(bsz, num_blocks, block_size, num_heads, head_dim).permute(0, 3, 1, 2, 4).to(torch.float32) + + valid_tokens = torch.ones((seq_len,), device=x.device, dtype=torch.float32) + if pad_tokens: + valid_tokens = torch.nn.functional.pad(valid_tokens, (0, pad_tokens)) + valid_mask = valid_tokens.view(1, 1, num_blocks, block_size, 1) + counts = valid_mask.sum(dim=-2).clamp_min_(1.0) + + if mean_subtract is not None: + mean_values = mean_subtract.squeeze(-2) if mean_subtract.ndim == 4 else mean_subtract + x_blocks = x_blocks - mean_values.to(torch.float32).unsqueeze(2).unsqueeze(2) + x_blocks = x_blocks * valid_mask + + pooled = x_blocks.sum(dim=-2) / counts + + norms = torch.linalg.vector_norm(x_blocks, dim=-1, keepdim=True) + normalized = torch.where(norms > 0, x_blocks / norms, torch.zeros_like(x_blocks)) + grams = torch.matmul(normalized, normalized.transpose(-1, -2)) + mean_sim = grams.sum(dim=(-1, -2)) / counts.squeeze(-1).squeeze(-1).pow(2) + sim_blocks = mean_sim > sim_threshold.view(1, num_heads, 1) + # Match the CUDA reference (SpargeAttn spas_sage_attn/utils.py): a partial + # tail block (pad_tokens > 0 only ever affects the last block) divides padded + # rows by a zero norm there (0/0 -> NaN -> NaN > thr == False), forcing it + # dense. Our guarded norm would mark it prunable and sparsify the sequence + # tail (= last video frames), so force the partial last block to False. + if pad_tokens: + sim_blocks[:, :, -1] = False + + max_abs = x_blocks.abs().amax(dim=(-1, -2)).clamp_min_(0.0) + scales = (max_abs / 127.0) + 1.0e-7 + q = x_blocks / scales.unsqueeze(-1).unsqueeze(-1) + q = torch.where(q >= 0, torch.floor(q + 0.5), torch.ceil(q - 0.5)) + q = q.clamp_(-127, 127).to(torch.int8) + if layout == "HND": + q_native = q.view(bsz, num_heads, num_blocks * block_size, head_dim)[:, :, :seq_len, :].contiguous() + else: + q_native = ( + q.permute(0, 2, 3, 1, 4) + .reshape(bsz, num_blocks * block_size, num_heads, head_dim)[:, :seq_len, :, :] + .contiguous() + ) + + return pooled.to(x.dtype), sim_blocks.contiguous(), q_native, scales.unsqueeze(-1).to(torch.float32).contiguous() + + +def _get_sparse_block_sparsity_stats( + valid_block_num: torch.Tensor, + lut: torch.Tensor, + *, + is_causal: bool = False, +) -> tuple[int, int, float, float, float]: + total_selected = int(valid_block_num.sum().item()) + del is_causal + total_candidates = lut.size(3) * lut.size(2) * lut.size(0) * lut.size(1) + selected_ratio = float(total_selected / total_candidates) if total_candidates > 0 else 0.0 + sparsity_ratio = 1.0 - selected_ratio + total_rows = valid_block_num.numel() + selected_blocks_per_row = float(total_selected / total_rows) if total_rows > 0 else 0.0 + return total_selected, total_candidates, selected_ratio, sparsity_ratio, selected_blocks_per_row + + +@dataclass(frozen=True) +class _SpargePreprocessContext: + query: torch.Tensor + key: torch.Tensor + batch: int + num_heads_q: int + num_heads_kv: int + seq_len_q: int + seq_len_kv: int + head_dim: int + is_causal: bool + smooth_k: bool + simthreshd1: torch.Tensor + topk: torch.Tensor + attention_sink: bool + quant_block_size: int + tensor_layout: str + q_route_block_tokens: int + k_route_block_tokens: int + sparse_q_block_tokens: int + sparse_k_block_tokens: int + query_tile_tokens: int + q_blocks_per_tile: int + k_blocks_per_tile: int + q_sparse_blocks_per_tile: int + k_sparse_blocks_per_tile: int + num_q_blocks: int + num_k_blocks: int + num_sparse_q_blocks: int + num_sparse_k_blocks: int + num_q_tiles: int + num_k_tiles: int + k_quant_granularity: int + + +def _build_sparge_preprocess_context( + query: torch.Tensor, + key: torch.Tensor, + *, + is_causal: bool, + smooth_k: bool, + simthreshd1: float | torch.Tensor, + topk: float | torch.Tensor, + attention_sink: bool, + quant_block_size: int, + tensor_layout: str, + k_quant_granularity: int = 64, + query_tile_tokens: int | None = None, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +) -> _SpargePreprocessContext: + if query.device.type != "xpu": + raise NotImplementedError("sparge_preprocess_topk is only supported on XPU") + if query.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"Q must be float16 or bfloat16, got {query.dtype}") + if key.dtype != query.dtype: + raise ValueError(f"K dtype must match Q dtype, got K={key.dtype}, Q={query.dtype}") + if quant_block_size != 64: + raise ValueError( + f"quant_block_size={quant_block_size} is not supported in sparge_preprocess_topk; only 64 is supported" + ) + if k_quant_granularity not in (64, 128): + raise ValueError(f"k_quant_granularity={k_quant_granularity} is not supported; only 64 or 128") + if k_quant_granularity % quant_block_size != 0: + raise ValueError( + f"k_quant_granularity={k_quant_granularity} must be a multiple of quant_block_size={quant_block_size}" + ) + + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout, expected_dtype=query.dtype) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout, expected_dtype=query.dtype) + if Bk != B: + raise ValueError("Batch size mismatch between Q and K") + if Dk != D: + raise ValueError("Head dim mismatch between Q and K") + if Hq % Hkv != 0: + raise ValueError("num_heads_q must be divisible by num_heads_kv") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + + q_route_block_tokens = _normalize_query_tile_tokens( + query_tile_tokens, + head_dim=D, + quant_block_size=quant_block_size, + ) + requested_sparse_q_block_tokens = None if sparse_q_block_tokens is None else int(sparse_q_block_tokens) + requested_sparse_k_block_tokens = None if sparse_k_block_tokens is None else int(sparse_k_block_tokens) + k_route_block_tokens = _routing_k_block_tokens_for_head_dim(D, quant_block_size) + if ( + D == 128 + and q_route_block_tokens == 256 + and requested_sparse_q_block_tokens == 256 + and requested_sparse_k_block_tokens in (None, 64) + ): + k_route_block_tokens = 64 + sparse_q_block_tokens = _normalize_sparse_q_block_tokens( + sparse_q_block_tokens, + quant_block_size=quant_block_size, + q_route_block_tokens=q_route_block_tokens, + ) + sparse_k_block_tokens = _normalize_sparse_k_block_tokens( + sparse_k_block_tokens, + quant_block_size=quant_block_size, + k_route_block_tokens=k_route_block_tokens, + ) + q_blocks_per_tile = q_route_block_tokens // quant_block_size + k_blocks_per_tile = k_route_block_tokens // quant_block_size + q_sparse_blocks_per_tile = q_route_block_tokens // sparse_q_block_tokens + k_sparse_blocks_per_tile = k_route_block_tokens // sparse_k_block_tokens + num_q_blocks = (Sq + quant_block_size - 1) // quant_block_size + num_k_blocks = (Skv + quant_block_size - 1) // quant_block_size + num_sparse_q_blocks = (Sq + sparse_q_block_tokens - 1) // sparse_q_block_tokens + num_sparse_k_blocks = (Skv + sparse_k_block_tokens - 1) // sparse_k_block_tokens + num_q_tiles = (Sq + q_route_block_tokens - 1) // q_route_block_tokens + num_k_tiles = (Skv + k_route_block_tokens - 1) // k_route_block_tokens + + return _SpargePreprocessContext( + query=query, + key=key, + batch=B, + num_heads_q=Hq, + num_heads_kv=Hkv, + seq_len_q=Sq, + seq_len_kv=Skv, + head_dim=D, + is_causal=is_causal, + smooth_k=smooth_k, + simthreshd1=_normalize_per_head_hparam(simthreshd1, Hq, query.device, "simthreshd1"), + topk=_normalize_per_head_hparam(topk, Hq, query.device, "topk").clamp_(0.0, 1.0), + attention_sink=attention_sink, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + q_route_block_tokens=q_route_block_tokens, + k_route_block_tokens=k_route_block_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + query_tile_tokens=q_route_block_tokens, + q_blocks_per_tile=q_blocks_per_tile, + k_blocks_per_tile=k_blocks_per_tile, + q_sparse_blocks_per_tile=q_sparse_blocks_per_tile, + k_sparse_blocks_per_tile=k_sparse_blocks_per_tile, + num_q_blocks=num_q_blocks, + num_k_blocks=num_k_blocks, + num_sparse_q_blocks=num_sparse_q_blocks, + num_sparse_k_blocks=num_sparse_k_blocks, + num_q_tiles=num_q_tiles, + num_k_tiles=num_k_tiles, + k_quant_granularity=k_quant_granularity, + ) + + +def _sparge_preprocess_topk_torch_impl(ctx: _SpargePreprocessContext) -> dict[str, Any]: + key_mean = _sequence_mean_native_layout(ctx.key, ctx.tensor_layout) if ctx.smooth_k else None + pooled_q, sim_qblocks, q_int8_hnd, q_scale = _pool_sim_and_quant_torch( + ctx.query, + ctx.quant_block_size, + ctx.simthreshd1, + ctx.tensor_layout, + ) + pooled_k, sim_kblocks, k_int8_hnd, k_scale = _pool_sim_and_quant_torch( + ctx.key, + ctx.quant_block_size, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + key_mean, + ) + + if ctx.q_blocks_per_tile > 1: + tile_pooled_q = [] + tile_sim_q = [] + for qtile in range(ctx.num_q_tiles): + qblk_start = qtile * ctx.q_blocks_per_tile + qblk_end = min(qblk_start + ctx.q_blocks_per_tile, pooled_q.size(2)) + tile_tokens = _slice_sequence_native_layout( + ctx.query, + ctx.tensor_layout, + qblk_start * ctx.quant_block_size, + min((qblk_end * ctx.quant_block_size), ctx.seq_len_q), + ) + pooled_tile, sim_tile, _, _ = _pool_sim_and_quant_torch( + tile_tokens, + ctx.query_tile_tokens, + ctx.simthreshd1, + ctx.tensor_layout, + ) + tile_pooled_q.append(pooled_tile[:, :, 0, :]) + tile_sim_q.append(sim_tile[:, :, 0]) + pooled_q_for_routing = torch.stack(tile_pooled_q, dim=2) + sim_q_for_routing = torch.stack(tile_sim_q, dim=2) + else: + pooled_q_for_routing = pooled_q + sim_q_for_routing = sim_qblocks + + if ctx.k_blocks_per_tile > 1: + pooled_k_for_routing, sim_k_for_routing, _, _ = _pool_sim_and_quant_torch( + ctx.key, + ctx.k_route_block_tokens, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + key_mean, + ) + else: + pooled_k_for_routing = pooled_k + sim_k_for_routing = sim_kblocks + + kv_head_index = torch.arange(ctx.num_heads_q, device=ctx.query.device, dtype=torch.int64) // ( + ctx.num_heads_q // ctx.num_heads_kv + ) + pooled_k_for_q = pooled_k_for_routing[:, kv_head_index] + sim_k_for_q = sim_k_for_routing[:, kv_head_index] + sim_k_expand = sim_k_for_q.unsqueeze(-2).expand(-1, -1, ctx.num_q_tiles, -1) + sim_q_expand = sim_q_for_routing.unsqueeze(-1).expand(-1, -1, -1, pooled_k_for_routing.size(2)) + + pooled_score = torch.matmul( + pooled_q_for_routing.to(torch.float32), pooled_k_for_q.transpose(-1, -2).to(torch.float32) + ) + pooled_score *= ctx.head_dim**-0.5 + pooled_score = pooled_score.masked_fill(~sim_k_expand, -torch.inf) + if ctx.is_causal: + causal_mask = _build_block_causal_mask( + ctx.num_q_tiles, + pooled_k_for_routing.size(2), + ctx.q_route_block_tokens, + ctx.k_route_block_tokens, + ctx.query.device, + ) + pooled_score = pooled_score.masked_fill( + ~causal_mask.view(1, 1, ctx.num_q_tiles, pooled_k_for_routing.size(2)), -torch.inf + ) + else: + causal_mask = None + + pooled_prob = _safe_softmax(pooled_score) + sorted_prob = torch.sort(pooled_prob, dim=-1, descending=True) + _, _, _, num_k_route_blocks = pooled_prob.shape + num_to_select = ( + (ctx.topk.view(1, ctx.num_heads_q, 1) * num_k_route_blocks) + .to(torch.int64) + .expand(ctx.batch, -1, ctx.num_q_tiles) + .contiguous() + ) + final_tile_map = torch.zeros_like(pooled_prob, dtype=torch.bool) + final_tile_map[~sim_k_expand] = True + final_tile_map[~sim_q_expand] = True + final_tile_map = _fill_block_map_torch(final_tile_map, num_to_select, sorted_prob.indices) + if causal_mask is not None: + final_tile_map &= causal_mask.view(1, 1, ctx.num_q_tiles, num_k_route_blocks) + if ctx.attention_sink: + final_tile_map[..., 0] = True + + q_block_to_tile = ( + torch.arange(ctx.num_sparse_q_blocks, device=ctx.query.device, dtype=torch.int64) + // ctx.q_sparse_blocks_per_tile + ) + q_block_to_tile = q_block_to_tile.clamp_max(ctx.num_q_tiles - 1) + k_block_to_tile = ( + torch.arange(ctx.num_sparse_k_blocks, device=ctx.query.device, dtype=torch.int64) + // ctx.k_sparse_blocks_per_tile + ) + k_block_to_tile = k_block_to_tile.clamp_max(ctx.num_k_tiles - 1) + raw_block_map = final_tile_map.index_select(2, q_block_to_tile).index_select(3, k_block_to_tile).contiguous() + + return { + "query_i8": q_int8_hnd, + "key_i8": k_int8_hnd, + "qscale": q_scale, + "kscale": k_scale, + "raw_block_map": raw_block_map, + "tile_block_map": final_tile_map.contiguous(), + "sim_qblocks": sim_q_for_routing.contiguous(), + "sim_kblocks": sim_k_for_routing.contiguous(), + "backend": "torch", + } + + +def _finalize_sparge_preprocess_outputs( + ctx: _SpargePreprocessContext, + backend_result: dict[str, Any], +) -> dict[str, Any]: + raw_block_map = backend_result["raw_block_map"].contiguous() + block_map = raw_block_map + lut = backend_result.get("lut") + valid_block_num = backend_result.get("valid_block_num") + if lut is None or valid_block_num is None: + lut, valid_block_num = _block_map_lut_torch(block_map) + total_selected, total_candidates, selected_ratio, sparsity_ratio, selected_blocks_per_row = ( + _get_sparse_block_sparsity_stats( + valid_block_num, + lut, + is_causal=ctx.is_causal, + ) + ) + + return { + "query_i8": backend_result["query_i8"].contiguous(), + "key_i8": backend_result["key_i8"].contiguous(), + "qscale": backend_result["qscale"].contiguous(), + "kscale": backend_result["kscale"].contiguous(), + "lut": lut, + "valid_block_num": valid_block_num, + "block_map": block_map, + "raw_block_map": raw_block_map, + "tile_block_map": backend_result["tile_block_map"].contiguous(), + "sim_qblocks": backend_result["sim_qblocks"].contiguous(), + "sim_kblocks": backend_result["sim_kblocks"].contiguous(), + "query_tile_tokens": ctx.query_tile_tokens, + "quant_block_size": ctx.quant_block_size, + "sparse_q_block_tokens": ctx.sparse_q_block_tokens, + "sparse_k_block_tokens": ctx.sparse_k_block_tokens, + "backend": backend_result["backend"], + "kernel_compatibility_added_blocks": 0, + "stats": { + "total_selected": total_selected, + "total_candidates": total_candidates, + "selected_ratio": selected_ratio, + "sparsity_ratio": sparsity_ratio, + "selected_blocks_per_row": selected_blocks_per_row, + }, + } + + +def _sparge_preprocess_topk_dispatch( + ctx: _SpargePreprocessContext, + *, + backend_preference: str = "auto", +) -> dict[str, Any]: + from .sparge_preprocess_triton import dispatch_sparge_preprocess_backend + + backend_result = dispatch_sparge_preprocess_backend( + ctx=ctx, + torch_backend=lambda: _sparge_preprocess_topk_torch_impl(ctx), + backend_preference=backend_preference, + ) + return _finalize_sparge_preprocess_outputs(ctx, backend_result) + + +def _sparge_preprocess_topk_torch( + query: torch.Tensor, + key: torch.Tensor, + *, + is_causal: bool = False, + smooth_k: bool = True, + simthreshd1: float | torch.Tensor = -0.1, + topk: float | torch.Tensor = 0.5, + attention_sink: bool = False, + quant_block_size: int = 64, + tensor_layout: str = "HND", + query_tile_tokens: int | None = None, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +) -> dict[str, Any]: + ctx = _build_sparge_preprocess_context( + query, + key, + is_causal=is_causal, + smooth_k=smooth_k, + simthreshd1=simthreshd1, + topk=topk, + attention_sink=attention_sink, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + query_tile_tokens=query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + return _finalize_sparge_preprocess_outputs(ctx, _sparge_preprocess_topk_torch_impl(ctx)) + + +def sparge_block_map_to_mask( + block_map: torch.Tensor, + *, + quant_block_size: int = 64, + q_block_tokens: int | None = None, + k_block_tokens: int | None = None, + seq_len_q: int | None = None, + seq_len_kv: int | None = None, + is_causal: bool = False, +) -> torch.Tensor: + if block_map.dtype != torch.bool or block_map.ndim != 4: + raise ValueError("block_map must be a 4D bool tensor") + if quant_block_size <= 0: + raise ValueError(f"quant_block_size must be positive, got {quant_block_size}") + q_tokens = quant_block_size if q_block_tokens is None else int(q_block_tokens) + k_tokens = quant_block_size if k_block_tokens is None else int(k_block_tokens) + if q_tokens <= 0 or k_tokens <= 0: + raise ValueError(f"q_block_tokens and k_block_tokens must be positive, got {q_tokens} and {k_tokens}") + batch, heads, q_blocks, kv_blocks = block_map.shape + full_q = q_blocks * q_tokens + full_k = kv_blocks * k_tokens + seq_q = full_q if seq_len_q is None else seq_len_q + seq_kv = full_k if seq_len_kv is None else seq_len_kv + expanded = block_map.repeat_interleave(q_tokens, dim=-2).repeat_interleave(k_tokens, dim=-1) + expanded = expanded[:, :, :seq_q, :seq_kv] + mask = torch.full(expanded.shape, -1.0e9, dtype=torch.float32, device=block_map.device) + mask = torch.where(expanded, torch.zeros_like(mask), mask) + if is_causal: + causal = torch.triu( + torch.full((seq_q, seq_kv), -1.0e9, dtype=torch.float32, device=block_map.device), diagonal=1 + ) + mask = torch.minimum(mask, causal.view(1, 1, seq_q, seq_kv)) + return mask.contiguous() + + +def sparge_preprocess_topk( + query: torch.Tensor, + key: torch.Tensor, + *, + is_causal: bool = False, + smooth_k: bool = True, + simthreshd1: float | torch.Tensor = -0.1, + topk: float | torch.Tensor = 0.5, + attention_sink: bool = False, + quant_block_size: int = 64, + tensor_layout: str = "HND", + k_quant_granularity: int = 64, + query_tile_tokens: int | None = None, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, + backend_preference: str | None = None, +) -> dict[str, Any]: + ctx = _build_sparge_preprocess_context( + query, + key, + is_causal=is_causal, + smooth_k=smooth_k, + simthreshd1=simthreshd1, + topk=topk, + attention_sink=attention_sink, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + k_quant_granularity=k_quant_granularity, + query_tile_tokens=query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + return _sparge_preprocess_topk_dispatch( + ctx, + backend_preference=backend_preference or _get_sparse_preprocess_backend_preference(), + ) + + +def sparge_preprocess_topk_decode( + query: torch.Tensor, + key: torch.Tensor, + key_cache: torch.Tensor, + *, + smooth_k: bool = True, + simthreshd1: float | torch.Tensor = -0.1, + topk: float | torch.Tensor = 0.5, + attention_sink: bool = False, + quant_block_size: int = 64, + tensor_layout: str = "HND", +) -> dict[str, Any]: + if query.device.type != "xpu": + raise NotImplementedError("sparge_preprocess_topk_decode is only supported on XPU") + if quant_block_size != 64: + raise ValueError("sparge_preprocess_topk_decode currently supports only quant_block_size=64") + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout, expected_dtype=query.dtype) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout, expected_dtype=key.dtype) + Bkc, Hkvc, Skvc, Dkc = _validate_attention_tensor( + key_cache, "K_cache", tensor_layout, expected_dtype=key_cache.dtype + ) + if Bk != B or Bkc != B: + raise ValueError("Batch size mismatch between Q/K/K_cache") + if Hkvc != Hkv: + raise ValueError("K and K_cache head-count mismatch") + if Dk != D or Dkc != D: + raise ValueError("Head dim mismatch between Q/K/K_cache") + if query.dtype != key.dtype or key_cache.dtype != query.dtype: + raise ValueError("Q/K/K_cache dtype must match") + if query.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"Q/K/K_cache must be float16 or bfloat16, got {query.dtype}") + if Sq != 1: + raise ValueError(f"sparge_preprocess_topk_decode currently supports only seq_len_q == 1, got {Sq}") + if Skv <= 0 or Skvc <= 0: + raise ValueError("sparge_preprocess_topk_decode requires non-empty current K and K_cache") + if Hq % Hkv != 0: + raise ValueError("num_heads_q must be divisible by num_heads_kv") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + + query_hnd = _to_hnd(query, tensor_layout) + key_hnd = _to_hnd(key, tensor_layout) + key_cache_hnd = _to_hnd(key_cache, tensor_layout) + key_total_hnd = torch.cat([key_cache_hnd, key_hnd], dim=2).contiguous() + total_seq = key_total_hnd.size(2) + num_k_blocks = (total_seq + quant_block_size - 1) // quant_block_size + simthreshd1_tensor = _normalize_per_head_hparam(simthreshd1, Hq, query.device, "simthreshd1") + topk_tensor = _normalize_per_head_hparam(topk, Hq, query.device, "topk").clamp_(0.0, 1.0) + + _, _, q_int8_hnd, q_scale = _pool_sim_and_quant_torch(query_hnd, quant_block_size, simthreshd1_tensor, "HND") + pooled_q_route, sim_q_route, _, _ = _pool_sim_and_quant_torch( + query_hnd, + _query_tile_tokens_for_head_dim(D), + simthreshd1_tensor, + "HND", + ) + key_mean = key_total_hnd.mean(dim=-2, keepdim=True) if smooth_k else None + pooled_k, sim_kblocks, key_i8_total_hnd, k_scale_total = _pool_sim_and_quant_torch( + key_total_hnd, + quant_block_size, + simthreshd1_tensor[:Hkv], + "HND", + key_mean, + ) + + kv_head_index = torch.arange(Hq, device=query.device, dtype=torch.int64) // (Hq // Hkv) + pooled_k_for_q = pooled_k[:, kv_head_index] + sim_k_for_q = sim_kblocks[:, kv_head_index] + pooled_score = torch.matmul( + pooled_q_route[:, :, :1, :].to(torch.float32), + pooled_k_for_q.transpose(-1, -2).to(torch.float32), + ) + pooled_score *= D**-0.5 + pooled_score = pooled_score.masked_fill(~sim_k_for_q.unsqueeze(-2), -torch.inf) + pooled_prob = _safe_softmax(pooled_score) + sorted_prob = torch.sort(pooled_prob, dim=-1, descending=True) + num_to_select = (topk_tensor.view(1, Hq, 1) * num_k_blocks).to(torch.int64).expand(B, -1, 1).contiguous() + final_tile_map = torch.zeros_like(pooled_prob, dtype=torch.bool) + final_tile_map[~sim_k_for_q.unsqueeze(-2)] = True + final_tile_map = _fill_block_map_torch(final_tile_map, num_to_select, sorted_prob.indices) + if attention_sink: + final_tile_map[..., 0] = True + + raw_block_map = final_tile_map.contiguous() + block_map = raw_block_map + lut, valid_block_num = _block_map_lut_torch(block_map) + total_selected, total_candidates, selected_ratio, sparsity_ratio, selected_blocks_per_row = ( + _get_sparse_block_sparsity_stats( + valid_block_num, + lut, + is_causal=False, + ) + ) + cache_blocks = (Skvc + quant_block_size - 1) // quant_block_size + cur_blocks = (Skv + quant_block_size - 1) // quant_block_size + + return { + "query_i8": _from_hnd(q_int8_hnd, tensor_layout), + "key_i8": _from_hnd(key_i8_total_hnd[:, :, Skvc:, :].contiguous(), tensor_layout), + "key_cache_i8": _from_hnd(key_i8_total_hnd[:, :, :Skvc, :].contiguous(), tensor_layout), + "qscale": q_scale.contiguous(), + "kscale": k_scale_total[:, :, cache_blocks:, :].contiguous(), + "kscale_cache": k_scale_total[:, :, :cache_blocks, :].contiguous(), + "lut": lut, + "valid_block_num": valid_block_num, + "block_map": block_map, + "raw_block_map": raw_block_map, + "tile_block_map": final_tile_map.contiguous(), + "sim_qblocks": sim_q_route[:, :, :1].contiguous(), + "sim_kblocks": sim_kblocks.contiguous(), + "query_tile_tokens": _query_tile_tokens_for_head_dim(D), + "quant_block_size": quant_block_size, + "backend": "torch", + "kernel_compatibility_added_blocks": 0, + "stats": { + "total_selected": total_selected, + "total_candidates": total_candidates, + "selected_ratio": selected_ratio, + "sparsity_ratio": sparsity_ratio, + "selected_blocks_per_row": selected_blocks_per_row, + }, + } + + +def sparge_sage2_decode_meansim_topk_xpu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = True, + scale: float | None = None, + smooth_k: bool = True, + simthreshd1: float | torch.Tensor = -0.1, + cdfthreshd: float | torch.Tensor | None = None, + topk: float | torch.Tensor = 0.5, + pvthreshd: float | torch.Tensor = 50, + attention_sink: bool = False, + tensor_layout: str = "HND", + output_dtype: torch.dtype | None = None, + return_sparsity: bool = False, + return_metadata: bool = False, +) -> torch.Tensor | tuple[Any, ...]: + if query.device.type != "xpu": + raise NotImplementedError("sparge_sage2_decode_meansim_topk_xpu is only supported on XPU") + if cdfthreshd is not None: + raise NotImplementedError("cdfthreshd routing is not implemented yet; use topk for the first slice") + if dropout_p != 0.0: + raise NotImplementedError("dropout_p must be 0.0 for sparge_sage2_decode_meansim_topk_xpu") + if attn_mask is not None and is_causal: + raise ValueError("attn_mask and is_causal cannot both be set") + if output_dtype is not None and output_dtype != value.dtype: + raise ValueError( + f"output_dtype must match value.dtype in the current implementation, got {output_dtype} vs {value.dtype}" + ) + if pvthreshd not in (None, 50): + warnings.warn("pvthreshd is not supported by the current ARK sparse kernel and is ignored", stacklevel=2) + + metadata = sparge_preprocess_topk_decode( + query, + key, + key_cache, + smooth_k=smooth_k, + simthreshd1=simthreshd1, + topk=topk, + attention_sink=attention_sink, + quant_block_size=64, + tensor_layout=tensor_layout, + ) + out = sage_sparse_decode( + metadata["query_i8"], + metadata["key_i8"], + value, + metadata["key_cache_i8"], + value_cache, + metadata["lut"], + metadata["valid_block_num"], + attn_mask=attn_mask, + is_causal=is_causal, + scale=scale, + quant_block_size=metadata["quant_block_size"], + qscale=metadata["qscale"], + kscale=metadata["kscale"], + kscale_cache=metadata["kscale_cache"], + tensor_layout=tensor_layout, + ) + sparsity_ratio = metadata["stats"]["sparsity_ratio"] + if return_metadata and return_sparsity: + return out, sparsity_ratio, metadata + if return_metadata: + return out, metadata + if return_sparsity: + return out, sparsity_ratio + return out + + +def sparge_sage2_attn_meansim_topk_xpu( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + smooth_k: bool = True, + simthreshd1: float | torch.Tensor = -0.1, + cdfthreshd: float | torch.Tensor | None = None, + topk: float | torch.Tensor = 0.5, + pvthreshd: float | torch.Tensor = 50, + attention_sink: bool = False, + tensor_layout: str = "HND", + output_dtype: torch.dtype | None = None, + return_sparsity: bool = False, + return_metadata: bool = False, + k_quant_granularity: int = 64, + query_tile_tokens: int | None = None, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +) -> torch.Tensor | tuple[Any, ...]: + if query.device.type != "xpu": + raise NotImplementedError("sparge_sage2_attn_meansim_topk_xpu is only supported on XPU") + if cdfthreshd is not None: + raise NotImplementedError("cdfthreshd routing is not implemented yet; use topk for the first slice") + if dropout_p != 0.0: + raise NotImplementedError("dropout_p must be 0.0 for sparge_sage2_attn_meansim_topk_xpu") + if attn_mask is not None and is_causal: + raise ValueError("attn_mask and is_causal cannot both be set") + if output_dtype is not None and output_dtype != value.dtype: + raise ValueError( + f"output_dtype must match value.dtype in the current implementation, got {output_dtype} vs {value.dtype}" + ) + if pvthreshd not in (None, 50): + warnings.warn("pvthreshd is not supported by the current ARK sparse kernel and is ignored", stacklevel=2) + + B, Hq, Sq, D = _validate_attention_tensor(query, "Q", tensor_layout, expected_dtype=query.dtype) + Bk, Hkv, Skv, Dk = _validate_attention_tensor(key, "K", tensor_layout, expected_dtype=key.dtype) + Bv, Hkv2, Skv2, Dv = _validate_attention_tensor(value, "V", tensor_layout, expected_dtype=value.dtype) + if Bk != B or Bv != B: + raise ValueError("Batch size mismatch between Q/K/V") + if Hkv2 != Hkv or Skv2 != Skv or Dv != Dk: + raise ValueError("K/V shape mismatch") + if Dk != D: + raise ValueError("Head dim mismatch between Q and K/V") + if query.dtype != key.dtype: + raise ValueError("Q and K dtype must match") + if query.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"Q/K must be float16 or bfloat16, got {query.dtype}") + if value.dtype not in (torch.float16, torch.bfloat16): + raise ValueError(f"V must be float16 or bfloat16, got {value.dtype}") + + effective_query_tile_tokens = query_tile_tokens + effective_q_tile_override = q_tile_override + if effective_query_tile_tokens is None and q_tile_override in (64, 128, 256): + effective_query_tile_tokens = q_tile_override + elif effective_query_tile_tokens is not None: + if q_tile_override == 0: + effective_q_tile_override = int(effective_query_tile_tokens) + elif q_tile_override != int(effective_query_tile_tokens): + raise ValueError( + "query_tile_tokens and q_tile_override must match when both are set; " + f"got query_tile_tokens={effective_query_tile_tokens}, q_tile_override={q_tile_override}" + ) + _validate_sparse_q_tile_override_for_head_dim(D, effective_q_tile_override) + + normalized_mask = _normalize_sparse_mask(attn_mask, B, Sq, Skv, query.device) + metadata = sparge_preprocess_topk( + query, + key, + is_causal=is_causal, + smooth_k=smooth_k, + simthreshd1=simthreshd1, + topk=topk, + attention_sink=attention_sink, + quant_block_size=64, + tensor_layout=tensor_layout, + k_quant_granularity=k_quant_granularity, + query_tile_tokens=effective_query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + _get_xpu_sparse_kernel_backend() + out = sage_sparse( + metadata["query_i8"], + metadata["key_i8"], + value, + metadata["lut"], + metadata["valid_block_num"], + attn_mask=normalized_mask, + is_causal=is_causal, + scale=scale, + quant_block_size=metadata["quant_block_size"], + qscale=metadata["qscale"], + kscale=metadata["kscale"], + q_tile_override=effective_q_tile_override, + sparse_q_block_tokens=metadata["sparse_q_block_tokens"], + sparse_k_block_tokens=metadata["sparse_k_block_tokens"], + tensor_layout=tensor_layout, + ) + sparsity_ratio = metadata["stats"]["sparsity_ratio"] + if return_metadata and return_sparsity: + return out, sparsity_ratio, metadata + if return_metadata: + return out, metadata + if return_sparsity: + return out, sparsity_ratio + return out diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_fmha_fwd_epilogue_compat.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_fmha_fwd_epilogue_compat.hpp new file mode 100644 index 000000000..adfd68f19 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_fmha_fwd_epilogue_compat.hpp @@ -0,0 +1,35 @@ +#pragma once + +namespace cutlass::fmha::kernel::detail { + +template +auto run_fmha_fwd_epilogue(Epilogue& epilogue, + TensorO2D const& O, + FragA& tArA, + FragARow& tA_max, + FragARow& tA_sum, + QVCoord blk_qv, + int thr_id, + int head_q, + int idx_b, + int) + -> decltype(epilogue(O, tArA, tA_max, tA_sum, blk_qv, thr_id, head_q, idx_b), void()) { + epilogue(O, tArA, tA_max, tA_sum, blk_qv, thr_id, head_q, idx_b); +} + +template +auto run_fmha_fwd_epilogue(Epilogue& epilogue, + TensorO2D const& O, + FragA& tArA, + FragARow& tA_max, + FragARow& tA_sum, + QVCoord blk_qv, + int thr_id, + int, + int, + long) + -> decltype(epilogue(O, tArA, tA_max, tA_sum, blk_qv, thr_id), void()) { + epilogue(O, tArA, tA_max, tA_sum, blk_qv, thr_id); +} + +} // namespace cutlass::fmha::kernel::detail diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sage_fwd_kernel.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sage_fwd_kernel.hpp index 4e773f49f..9037f6490 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sage_fwd_kernel.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sage_fwd_kernel.hpp @@ -353,4 +353,4 @@ class XeSageFwdKernel { } }; -} // namespace cutlass::fmha::kernel +} // namespace cutlass::fmha::kernel \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_fmha_fwd_epilogue.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_fmha_fwd_epilogue.hpp new file mode 100644 index 000000000..fc09bc7c4 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_fmha_fwd_epilogue.hpp @@ -0,0 +1,220 @@ +/*************************************************************************************************** + * Copyright (C) 2025 - 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + **************************************************************************************************/ + +#pragma once + +#include +#include "cutlass/cutlass.h" +#include "cutlass/detail/layout.hpp" +#include "cutlass/epilogue/collective/collective_epilogue.hpp" +#include "cutlass/epilogue/collective/detail.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" + +#include "cute/algorithm/subgroup_algorithms.hpp" +#include "cute/algorithm/tensor_algorithms.hpp" +#include "flash_attention_v2/collective/copy_block_slm.hpp" + +namespace cutlass::fmha::collective { + +using namespace cute; + +// Sparse kernels do not consume LSE today. Keep a local stateless epilogue so +// the sparse instantiations preserve the older no-LSE codegen shape. +template +class SparseFMHAFwdEpilogue { + public: + using TiledMMAPV = typename CollectiveMainloop::TiledMMAPV; + using TileShapePV = decltype(TiledMMAPV{}.tile_mnk()); + using TileShapeO = TileShapeO_; + using SGPerWG = decltype(product(take<1, 4>(shape(typename TiledMMAPV::ThrLayoutVMNK{})))); + + using TensorO = TensorO_; + using TensorO2D = decltype(TensorO_{}(append>(make_coord(_, _), 0))); + using ElementO = typename TensorO_::value_type; + + using FragA = typename CollectiveMainloop::FragA; + using FragARow = typename CollectiveMainloop::FragARow; + using ElementA = typename FragA::value_type; + + using ReduceK = decltype(size<3>(typename TiledMMAPV::ThrLayoutVMNK{})); + + static auto reduce_sg_v_helper() { + constexpr auto v_total_sg = get<1>(SGTileShapeA{}) / intel::_SGSize{}; + constexpr auto v_avail_sg = ReduceK{} / ReduceSGQ{}; + return Int<(v_total_sg > v_avail_sg) ? cute::gcd(v_total_sg, v_avail_sg) : v_total_sg>{}; + } + + using SGTileShapeA = decltype(atuple_coshape(FragA{}.tv_layout())); + using ReduceSGQ = decltype(cute::gcd(get<0>(SGTileShapeA{}), ReduceK{})); + using ReduceSGV = decltype(reduce_sg_v_helper()); + using ReduceSGLayout = decltype(make_identity_layout(Shape{})); + + using SGTileShapeO = decltype(shape_div(take<0, 2>(SGTileShapeA{}), shape(ReduceSGLayout{}))); + + using ReduceFragA = + decltype(make_subgroup_tensor(make_layout(select<1, 0>(SGTileShapeO{}), Stride, E<0>>{}))); + using ReduceFragARow = decltype(reduce<1>(ReduceFragA{}, sycl::plus{})); + + static auto default_tiled_copy_O_helper() { + if constexpr (ReduceK{} == _1{}) { + return make_block_2d_copy_D(TiledMMAPV{}, TensorO2D{}); + } else { + return make_block_2d_copy_D_subtiled(TiledMMAPV{}, ReduceFragA{}.tv_layout(), ReduceSGLayout{}, TensorO2D{}); + } + } + + using DefaultTiledCopyO = decltype(default_tiled_copy_O_helper()); + using TiledCopyO = conditional_t, DefaultTiledCopyO, TiledCopyO_>; + + struct Arguments {}; + struct Params {}; + + using AlignedSGTileA_Q = C<((size<0>(SGTileShapeA{}) + intel::sg_size - 1) / intel::sg_size) * intel::sg_size>; + + struct SharedStorageNone {}; + struct SharedStorageReduceK { + cute::array a_data; + cute::array a_sum_data, a_max_data; + }; + + using SharedStorage = conditional_t<(ReduceK{} > _1{}), SharedStorageReduceK, SharedStorageNone>; + + private: + SharedStorage& shared; + + public: + static constexpr Params to_underlying_arguments(Arguments const&, void*) { return {}; } + + CUTLASS_HOST_DEVICE static bool can_implement(Arguments const&) { return true; } + + CUTLASS_HOST_DEVICE + SparseFMHAFwdEpilogue(Params const&, SharedStorage& shared_) : shared(shared_) {} + + template + CUTLASS_DEVICE void operator()(TensorO2D const& O, FragA& tArA, FragARow& tA_max, FragARow& tA_sum, QVCoord blk_qv, + int thr_id) { + using namespace cute; + using ElementAcc = typename FragA::element_type; + + auto [rA, rA_sum, active] = reduce_A(tArA, tA_max, tA_sum, thr_id); + if (!active) { + return; + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < rA_sum.size(); i++) { + rA_sum(i) = ElementAcc(1) / rA_sum(i); + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < rA.size(); i++) { + rA(i) *= broadcast<0>(rA_sum, rA, i); + } + + Tensor cO = make_identity_tensor(O.shape()); + Tensor gO = local_tile(cO, TileShapeO{}, blk_qv); + + TiledCopyO copy_o{O}; + auto thr_copy_o = copy_o.get_slice(thr_id); + auto tOrO = thr_copy_o.partition_sg_fragment_S(gO); + auto tOgO = thr_copy_o.partition_D(gO); + + reorder(rA, tOrO); + copy(copy_o, tOrO, tOgO); + } + + template + CUTLASS_DEVICE decltype(auto) reduce_A(FragA_& tArA, FragARow_& tA_max, FragARow_& tA_sum, int thr_id) { + using namespace sycl::ext::oneapi::this_work_item; + + if constexpr (ReduceK{} == _1{}) { + return std::make_tuple(tArA, tA_sum, true); + } else { + auto thr_vak = group<1, 3>(TiledMMAPV{}.get_thr_layout_vmnk()).get_flat_coord(assert_uniform(thr_id)); + auto a_tile = get<1>(thr_vak); + auto k_blk = get<2>(thr_vak); + + auto shape_A = append(append(SGTileShapeA{}, ReduceK{}), SGPerWG{} / ReduceK{}); + auto shape_A_row = make_shape(get<0>(SGTileShapeO{}), shape(ReduceSGLayout{}), ReduceK{}, SGPerWG{} / ReduceK{}); + + auto sA_layout = group<2, 4>(flat_divide(make_ordered_layout(shape_A, Step<_1, _0, _2, _3>{}), SGTileShapeO{})); + auto sA_row_stride = make_stride(_1{}, make_stride(get<0>(shape_A_row), _0{}), AlignedSGTileA_Q{}, + AlignedSGTileA_Q{} * ReduceK{}); + auto sA_row_layout = make_layout(shape_A_row, sA_row_stride); + + auto basis2 = make_basis_like(SGTileShapeO{}); + auto sA_coords = make_layout(append(SGTileShapeO{}, shape(ReduceSGLayout{})), + append(basis2, product_each(zip(SGTileShapeO{}, basis2)))); + + auto sA = make_tensor(make_smem_ptr(&shared.a_data), sA_layout); + auto sA_max = make_tensor(make_smem_ptr(&shared.a_max_data), sA_row_layout); + auto sA_sum = make_tensor(make_smem_ptr(&shared.a_sum_data), sA_row_layout); + + copy_block_r2s(tA_max, sA_max(_, _, k_blk, a_tile)); + barrier_arrive(ScopeWorkgroup, SemanticsRelease | SemanticsWGMemory); + copy_block_r2s(tA_sum, sA_sum(_, _, k_blk, a_tile)); + copy_block_r2s(tArA, sA(_, _, _, k_blk, a_tile), sA_coords); + + bool active = (k_blk < size(ReduceSGLayout{})) || (ReduceK{} == size(ReduceSGLayout{})); + + barrier_wait(ScopeWorkgroup, SemanticsAcquire | SemanticsWGMemory); + barrier_arrive(ScopeWorkgroup, SemanticsRelease | SemanticsWGMemory); + + ReduceFragA rA; + ReduceFragARow rA_sum, rA_max, rA_kmax[ReduceK{}]; + + if (active) { + CUTLASS_PRAGMA_UNROLL + for (int kr = 0; kr < ReduceK{}; kr++) { + copy_block_s2r(sA_max(_, k_blk, kr, a_tile), rA_kmax[kr]); + } + + rA_max = rA_kmax[0]; + for (int kr = 1; kr < ReduceK{}; kr++) { + cute::transform(rA_max, rA_kmax[kr], rA_max, cute::max_fn{}); + } + + for (int kr = 0; kr < ReduceK{}; kr++) { + cute::transform(rA_max, rA_kmax[kr], rA_kmax[kr], + [](auto gmax, auto kmax) { return sycl::native::exp2(kmax - gmax); }); + } + } + + barrier_wait(ScopeWorkgroup, SemanticsAcquire | SemanticsWGMemory); + + if (active) { + clear(rA_sum); + + CUTLASS_PRAGMA_UNROLL + for (int kr = 0; kr < ReduceK{}; kr++) { + ReduceFragARow rA_sum_read; + copy_block_s2r(sA_sum(_, k_blk, kr, a_tile), rA_sum_read); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < rA_sum_read.size(); i++) { + rA_sum(i) += rA_sum_read(i) * rA_kmax[kr](i); + } + } + + clear(rA); + + CUTLASS_PRAGMA_UNROLL + for (int kr = 0; kr < ReduceK{}; kr++) { + ReduceFragA rA_read; + copy_block_s2r(sA(_, _, k_blk, kr, a_tile), sA_coords(_, _, 0), rA_read); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < rA_read.size(); i++) { + rA(i) += rA_read(i) * broadcast<0>(rA_kmax[kr], rA, i); + } + } + } + + return std::make_tuple(rA, rA_sum, active); + } + } +}; + +} // namespace cutlass::fmha::collective diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sage_fwd_kernel.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sage_fwd_kernel.hpp new file mode 100644 index 000000000..a5382ff2b --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sage_fwd_kernel.hpp @@ -0,0 +1,350 @@ +/*************************************************************************************************** + * Copyright (C) 2025 - 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/gemm.h" +#include "cutlass/kernel_hardware_info.hpp" + +#include "cute/util/type_traits.hpp" +#include "flash_attention_v2/collective/fmha_fusion.hpp" +#include "flash_attention_v2/collective/xe_fmha_fwd_epilogue.hpp" +#include "flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp" +#include "xe_fmha_fwd_epilogue_compat.hpp" +#include "xe_sparse_sagev1_fwd_mainloop.hpp" +#include "flash_attention_v2/kernel/xe_tile_scheduler.hpp" +#include + +namespace cutlass::fmha::kernel { + +using namespace cute; + +template +struct has_canonical_nhd_k : std::false_type {}; + +template +struct has_canonical_nhd_k().canonical_nhd_k)>> : std::true_type {}; + +template +struct has_sparse_q_block_size : std::false_type {}; + +template +struct has_sparse_q_block_size().sparse_q_block_size)>> : std::true_type {}; + +template +struct SparseSageProblemShape { + using SeqLenType = cute::conditional_t; + int batch; + int num_heads_q, num_heads_kv; + SeqLenType seq_len_qo, seq_len_kv, seq_len_kv_cache; + int head_size_qk, head_size_vo; +}; + +template +class XeSparseSageFwdKernel { + public: + using ProblemShape = ProblemShape_; + using VariableLength = cutlass::fmha::collective::VariableLength; + static constexpr bool is_var_len = cutlass::fmha::collective::is_variable_length_v; + using CollectiveMainloop = CollectiveMainloop_; + using MainloopArguments = typename CollectiveMainloop::Arguments; + using MainloopParams = typename CollectiveMainloop::Params; + using TiledMMAQK = typename CollectiveMainloop::TiledMMAQK; + using TileShapeQK = typename CollectiveMainloop::TileShapeQK; + using SubgroupLayoutQK = typename CollectiveMainloop::SubgroupLayoutQK; + using ElementQ = typename CollectiveMainloop::TensorQ::element_type; + using ElementK = typename CollectiveMainloop::TensorK::element_type; + using ElementV = typename CollectiveMainloop::TensorV::element_type; + using StrideQ = decltype(stride(typename CollectiveMainloop::TensorQ{})); + using StrideK = decltype(stride(typename CollectiveMainloop::TensorK{})); + using StrideV = decltype(stride(typename CollectiveMainloop::TensorV{})); + using SGPerWG = typename CollectiveMainloop::SGPerWG; + using FragA = typename CollectiveMainloop::FragA; + using FragARow = typename CollectiveMainloop::FragARow; + using TileScheduler = TileScheduler_; + using TileSchedulerParams = typename TileScheduler::Params; + using CollectiveEpilogue = CollectiveEpilogue_; + using EpilogueArguments = typename CollectiveEpilogue::Arguments; + using EpilogueParams = typename CollectiveEpilogue::Params; + using TileShapeO = typename CollectiveEpilogue::TileShapeO; + using ElementO = typename CollectiveEpilogue::TensorO::element_type; + using StrideO = decltype(stride(typename CollectiveEpilogue::TensorO{})); + using MainloopSharedStorage = typename CollectiveMainloop::SharedStorage; + using EpilogueSharedStorage = typename CollectiveEpilogue::SharedStorage; + union SharedStorage { + MainloopSharedStorage mainloop; + EpilogueSharedStorage epilogue; + }; + + static constexpr int SharedStorageSize = is_empty_v ? size_t(0) : sizeof(SharedStorage); + + struct KernelArguments { + ProblemShape shape; + const ElementQ* Q; + StrideQ dQ; + const ElementK* K; + StrideK dK; + const ElementV* V; + StrideV dV; + ElementO* O; + StrideO dO; + const ElementK* K_cache; + StrideK dK_cache{}; + const ElementV* V_cache; + StrideV dV_cache{}; + }; + using KernelParams = KernelArguments; + + struct Arguments { + KernelArguments kernel{}; + MainloopArguments mainloop{}; + EpilogueArguments epilogue{}; + KernelHardwareInfo hw_info{}; + }; + + struct Params { + KernelParams kernel; + MainloopParams mainloop; + EpilogueParams epilogue; + TileSchedulerParams scheduler; + }; + + static Params to_underlying_arguments(Arguments const& args, void* workspace) { + return {args.kernel, CollectiveMainloop::to_underlying_arguments(args.mainloop, workspace), + CollectiveEpilogue::to_underlying_arguments(args.epilogue, workspace), + TileScheduler::to_underlying_arguments(args.kernel.shape, args.hw_info, TileShapeO{})}; + } + + static bool can_implement(Arguments const& args) { + return CollectiveMainloop::can_implement(args.mainloop) && CollectiveEpilogue::can_implement(args.epilogue); + } + + static int get_workspace_size(Arguments const& args) { return 0; } + + static cutlass::Status initialize_workspace(Arguments const& args, void* workspace = nullptr, + cudaStream_t stream = nullptr, CudaHostAdapter* cuda_adapter = nullptr) { + return Status::kSuccess; + } + + static dim3 get_grid_shape(Params const& params) { + return TileScheduler::template get_grid_shape(params.scheduler); + } + + static dim3 get_block_shape() { return dim3(SGPerWG::value * intel::sg_size, 1, 1); } + + CUTLASS_DEVICE + Shape get_sequence_length_shape(ProblemShape const& problem_shape, int const& batch) { + if constexpr (is_var_len) { + return cutlass::fmha::collective::apply_variable_length( + Shape{problem_shape.seq_len_qo, problem_shape.seq_len_kv, + problem_shape.seq_len_kv_cache}, + batch); + } else { + return Shape{problem_shape.seq_len_qo, problem_shape.seq_len_kv, problem_shape.seq_len_kv_cache}; + } + } + + CUTLASS_DEVICE + void operator()(Params const& params, char* smem_buf) { + using namespace sycl::ext::oneapi::this_work_item; + + SharedStorage& shared_storage = *reinterpret_cast(smem_buf); + + auto& p = params.kernel; + ProblemShape const& s = p.shape; + int head_group_q = s.num_heads_q / s.num_heads_kv; + + int thr_id = int(ThreadIdxX()); + auto cS = make_identity_tensor(take<0, 2>(TiledMMAQK{}.tile_mnk())); + auto tScS = TiledMMAQK{}.get_slice(thr_id).partition_C(cS); + auto q_offset_wi = get<0>(tScS(0)); + auto q_offset_sg = group_broadcast(sycl::ext::oneapi::this_work_item::get_sub_group(), q_offset_wi, 0); + constexpr int q_sg_tile = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + + TileScheduler tile_scheduler{params.scheduler}; + + CUTLASS_PRAGMA_NO_UNROLL + for (; tile_scheduler.is_valid(); ++tile_scheduler) { + auto [blk_q, blk_v, head_q, idx_b] = tile_scheduler.get_block_coord(); + auto blk_qv = make_coord(blk_q, blk_v); + int head = head_q / head_group_q; + + auto sequence_length_shape = get_sequence_length_shape(s, idx_b); + auto [seq_len_qo, seq_len_kv, seq_len_kv_cache] = sequence_length_shape; + if (blk_q * get<0>(TileShapeQK{}) >= seq_len_qo) continue; + + auto offset = cute::min(seq_len_qo, seq_len_kv); + auto discard_seq_coord = seq_len_qo - offset; + auto full_tile_offset = seq_len_kv - offset; + int seq_coord = cute::min(seq_len_qo, (blk_q * get<0>(TileShapeQK{}) + q_offset_sg)); + + if (CollectiveMainloop::CausalMask && seq_coord < discard_seq_coord) continue; + const int seq_len_new = CollectiveMainloop::CausalMask + ? full_tile_offset + cute::min(seq_len_kv, seq_coord - discard_seq_coord) + q_sg_tile + : seq_len_kv; + const int seq_len = seq_len_new + seq_len_kv_cache; + const int k_blocks = cute::ceil_div(seq_len, get<1>(TileShapeQK{})); + + int offset_q = 0, offset_k = 0, offset_v = 0, offset_o = 0; + int offset_k_cache = 0, offset_v_cache = 0; + if constexpr (is_var_len) { + auto qo_cumulative = s.seq_len_qo.cumulative_length; + auto kv_cumulative = s.seq_len_kv.cumulative_length; + offset_q = s.num_heads_q * s.head_size_qk * qo_cumulative[idx_b]; + offset_k = s.num_heads_kv * s.head_size_qk * kv_cumulative[idx_b]; + offset_v = s.num_heads_kv * s.head_size_vo * kv_cumulative[idx_b]; + offset_o = s.num_heads_q * s.head_size_vo * qo_cumulative[idx_b]; + if (s.seq_len_kv_cache.cumulative_length) { + auto kv_cumulative_cache = s.seq_len_kv_cache.cumulative_length; + offset_k_cache = s.num_heads_kv * s.head_size_qk * kv_cumulative_cache[idx_b]; + offset_v_cache = s.num_heads_kv * s.head_size_vo * kv_cumulative_cache[idx_b]; + } + } + + auto batch_dim = is_var_len ? 1 : s.batch; + auto shape_Q = make_shape(seq_len_qo, s.head_size_qk, s.num_heads_q, batch_dim); + auto shape_K = make_shape(seq_len_kv, s.head_size_qk, s.num_heads_kv, batch_dim); + auto shape_V = make_shape(s.head_size_vo, seq_len_kv, s.num_heads_kv, batch_dim); + auto shape_O = make_shape(seq_len_qo, s.head_size_vo, s.num_heads_q, batch_dim); + auto shape_K_cache = make_shape(seq_len_kv_cache, s.head_size_qk, s.num_heads_kv, batch_dim); + auto shape_V_cache = make_shape(s.head_size_vo, seq_len_kv_cache, s.num_heads_kv, batch_dim); + + auto dcQ = const_cast(p.Q + offset_q); + auto dcK = const_cast(p.K + offset_k); + auto dcV = const_cast(p.V + offset_v); + auto dcK_cache = const_cast(p.K_cache + offset_k_cache); + auto dcV_cache = const_cast(p.V_cache + offset_v_cache); + int seq_q_pad = (seq_len_qo + params.mainloop.scale_block_size - 1) / params.mainloop.scale_block_size; + int seq_kv_total = seq_len_kv + seq_len_kv_cache; + int seq_kv_pad = (seq_kv_total + params.mainloop.scale_block_size - 1) / params.mainloop.scale_block_size; + auto scaleQ = params.mainloop.scale_block_size + ? (float*)params.mainloop.qscale + (idx_b * s.num_heads_q * seq_q_pad + head_q * seq_q_pad) + : nullptr; + auto scaleK = params.mainloop.scale_block_size + ? (float*)params.mainloop.kscale + (idx_b * s.num_heads_kv * seq_kv_pad + head * seq_kv_pad) + : nullptr; + auto scaleV = params.mainloop.scale_block_size && params.mainloop.vscale + ? (float*)params.mainloop.vscale + + ((idx_b * s.num_heads_kv * seq_kv_pad + head * seq_kv_pad) * s.head_size_vo) + : nullptr; + int sparse_q_block = blk_q; + int sparse_q_rows_in_tile = 1; + // Dense FMHA treats blk_q as "the" Q tile handled by this workgroup. + // Sparse FMHA may further subdivide or coarsen that tile into routing rows. + // sparse_q_block_size tells us how many Q tokens share one LUT row. + int sparse_q_block_size = params.mainloop.scale_block_size; + if constexpr (has_sparse_q_block_size::value) { + if (params.mainloop.sparse_q_block_size > 0) { + sparse_q_block_size = params.mainloop.sparse_q_block_size; + } + } + if (sparse_q_block_size > 0) { + // q_blocks_per_tile is the sparse-only notion that differs from the dense + // path: one dense workgroup tile may correspond to multiple sparse routing + // rows. For q_tile=256 + sparse_q_block_size=64 we get 4 sparse rows; + // for q_tile=256 + sparse_q_block_size=256 we get 1 sparse row. + int q_blocks_per_tile = cute::max(1, int(get<0>(TileShapeQK{})) / sparse_q_block_size); + sparse_q_block = blk_q * q_blocks_per_tile; + sparse_q_rows_in_tile = q_blocks_per_tile; + if (params.mainloop.num_q_blocks > 0) { + // Tail tiles can expose fewer sparse rows than the nominal tile shape. + // Clamp both the starting row and the row count so the sparse mainloop + // never reads past the logical routing table. + sparse_q_rows_in_tile = cute::min(sparse_q_rows_in_tile, params.mainloop.num_q_blocks - sparse_q_block); + sparse_q_block = cute::min(sparse_q_block, params.mainloop.num_q_blocks - 1); + } + } + // Each workgroup gets a contiguous slice of valid_block_num / lut rows for + // the sparse Q rows covered by this tile. Dense FMHA has no equivalent + // per-tile routing metadata. + auto valid_blocks_base = params.mainloop.valid_block_num + ? params.mainloop.valid_block_num + + (idx_b * s.num_heads_q + head_q) * params.mainloop.num_q_blocks + + sparse_q_block + : nullptr; + auto lut_rows_base = params.mainloop.lut + ? params.mainloop.lut + + (((idx_b * s.num_heads_q + head_q) * params.mainloop.num_q_blocks + sparse_q_block) * + params.mainloop.num_k_blocks) + : nullptr; + auto ptrO = p.O + offset_o; + + auto stride_q = is_var_len ? cutlass::make_cute_packed_stride(StrideQ{}, shape_Q) : p.dQ; + auto stride_k = is_var_len ? cutlass::make_cute_packed_stride(StrideK{}, shape_K) : p.dK; + auto stride_v = is_var_len ? cutlass::make_cute_packed_stride(StrideV{}, shape_V) : p.dV; + auto stride_o = is_var_len ? cutlass::make_cute_packed_stride(StrideO{}, shape_O) : p.dO; + auto stride_k_cache = is_var_len ? cutlass::make_cute_packed_stride(StrideK{}, shape_K_cache) : p.dK_cache; + auto stride_v_cache = is_var_len ? cutlass::make_cute_packed_stride(StrideV{}, shape_V_cache) : p.dV_cache; + + Tensor Q = make_tensor(make_gmem_ptr(dcQ), make_layout(shape_Q, stride_q)); + Tensor K = make_tensor(make_gmem_ptr(dcK), make_layout(shape_K, stride_k)); + Tensor V = make_tensor(make_gmem_ptr(dcV), make_layout(shape_V, stride_v)); + Tensor K_cache = make_tensor(make_gmem_ptr(dcK_cache), make_layout(shape_K_cache, stride_k_cache)); + Tensor V_cache = make_tensor(make_gmem_ptr(dcV_cache), make_layout(shape_V_cache, stride_v_cache)); + Tensor O = make_tensor(make_gmem_ptr(ptrO), make_layout(shape_O, stride_o)); + + FragA tArA; + FragARow tA_max, tA_sum; + + int l_coord = is_var_len ? 0 : idx_b; + auto mainloop_params = params.mainloop; + if constexpr (has_canonical_nhd_k::value) { + mainloop_params.canonical_nhd_k = + !is_var_len && + int(get<1>(stride_k)) == 1 && + int(get<2>(stride_k)) == s.head_size_qk && + int(get<0>(stride_k)) == s.num_heads_kv * s.head_size_qk; + } + CollectiveMainloop mainloop(mainloop_params, shared_storage.mainloop); + // sparse_q_rows_in_tile is the key extra argument versus dense mainloop: + // it tells the sparse kernel whether this workgroup should walk a single + // LUT row linearly or merge multiple LUT rows inside one Q tile. + mainloop(Q(_, _, head_q, l_coord), K(_, _, head, l_coord), V(_, _, head, l_coord), tArA, tA_max, tA_sum, blk_qv, + 0, k_blocks, k_blocks, thr_id, seq_len, seq_len_kv_cache, idx_b, scaleQ, scaleK, scaleV, + full_tile_offset, discard_seq_coord, lut_rows_base, valid_blocks_base, sparse_q_rows_in_tile, + K_cache(_, _, head, l_coord), + V_cache(_, _, head, l_coord)); + + if constexpr (!is_empty_v && !is_empty_v) { + sycl::group_barrier(get_work_group<3>()); + } + + CollectiveEpilogue epilogue{params.epilogue, shared_storage.epilogue}; + detail::run_fmha_fwd_epilogue(epilogue, O(_, _, head_q, l_coord), tArA, tA_max, tA_sum, blk_qv, thr_id, head_q, idx_b, 0); + } + } +}; + +} // namespace cutlass::fmha::kernel diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sagev1_fwd_mainloop.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sagev1_fwd_mainloop.hpp new file mode 100644 index 000000000..0124fd21d --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sagev1_fwd_mainloop.hpp @@ -0,0 +1,936 @@ +/*************************************************************************************************** + * Copyright (C) 2025 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, + *this list of conditions and the following disclaimer. + * + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * 3. Neither the name of the copyright holder nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE + *POSSIBILITY OF SUCH DAMAGE. + * + **************************************************************************************************/ + +#pragma once + +#include "cute/util/print_tensor.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/dispatch_policy.hpp" + +#include "cute/algorithm/functional.hpp" +#include "cute/algorithm/gemm.hpp" +#include "cute/algorithm/subgroup_algorithms.hpp" +#include "cute/atom/mma_atom.hpp" +#include "flash_attention_v2/collective/fmha_fusion.hpp" +#include +#include +#include +namespace cutlass::sage { + +#ifndef ARK_CUTLASS_SAGE_XEDEFAULT_DEFINED +#define ARK_CUTLASS_SAGE_XEDEFAULT_DEFINED +template +class XeDefault {}; // Default FMHA mainloop, P in registers. +#endif + +}; // namespace cutlass::sage + +namespace cutlass::fmha::collective { + +using namespace cute; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template // Optional TiledCopy for loading V_cache +struct SPARSESAGEV1FwdMainloop { + static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +struct SPARSESAGEV1FwdMainloop, CausalMask_, FullMask_, CachedKV_, PagedKV_, UseInt8PV_, + WriteBackInt8PV_, ExecuteInt8PV_, TiledMMAQK_, TiledMMAPV_, VTiles_, TensorQ_, TensorK_, + TensorV_, TensorK_cache_, TensorV_cache_, TiledCopyQ_, TiledCopyK_, TiledCopyV_, + TiledCopyK_cache_, TiledCopyV_cache_> { + // + // Type Aliases + // + using TiledMMAQK = TiledMMAQK_; + using TiledMMAPV = TiledMMAPV_; + using TileShapeQK = decltype(TiledMMAQK{}.tile_mnk()); + using TileShapePV = decltype(TiledMMAPV{}.tile_mnk()); + static constexpr int VTiles = VTiles_; + using SubgroupLayoutQK = decltype(TiledMMAQK{}.get_atom_layout_mnk()); + using SubgroupLayoutPV = decltype(TiledMMAPV{}.get_atom_layout_mnk()); + using SGPerWG = decltype(product(take<1, 4>(shape(typename TiledMMAQK::ThrLayoutVMNK{})))); + + using TensorQ = TensorQ_; + using TensorK = TensorK_; + using TensorV = TensorV_; + + using TensorQ2D = decltype(TensorQ_{}(append>(make_coord(_, _), 0))); + using TensorK2D = decltype(TensorK_{}(append>(make_coord(_, _), 0))); + using TensorV2D = decltype(TensorV_{}(append>(make_coord(_, _), 0))); + + using TiledCopyQ = + conditional_t, decltype(make_block_2d_copy_A(TiledMMAQK{}, TensorQ2D{})), TiledCopyQ_>; + using TiledCopyK = + conditional_t, decltype(make_block_2d_copy_B(TiledMMAQK{}, TensorK2D{})), TiledCopyK_>; + using TiledCopyV = + conditional_t, decltype(make_block_2d_copy_B(TiledMMAPV{}, TensorV2D{})), TiledCopyV_>; + using TensorK_cache = TensorK_cache_; + using TensorV_cache = TensorV_cache_; + using TensorK_cache2D = decltype(TensorK_cache_{}(append>(make_coord(_, _), 0))); + using TensorV_cache2D = decltype(TensorV_cache_{}(append>(make_coord(_, _), 0))); + using TiledCopyK_cache = + conditional_t, decltype(make_block_2d_copy_B(TiledMMAQK{}, TensorK_cache2D{})), + TiledCopyK_cache_>; + using TiledCopyV_cache = + conditional_t, decltype(make_block_2d_copy_B(TiledMMAPV{}, TensorV_cache2D{})), + TiledCopyV_cache_>; + + // TODO: static_asserts on TiledMMAPV here... + + // + // Accumulator types + // + // FragS: accumulator for Q*K MMA + // FragO: accumulator for P*V MMAs. + // Note: v mode may be split into multiple pieces + // to reduce register pressure. + // Frag*Row types are reductions of the corresponding Frag* types + // over rows. + // + template + using FragC = decltype(TiledMMA{}.get_slice(0).partition_sg_fragment_C( + make_identity_tensor(select<0, 1>(TiledMMA{}.tile_mnk())))); + + using FragS = FragC; + using FragSRow = decltype(reduce<1>(FragS{}, sycl::plus{})); + using ElementS = typename TiledMMAQK::ValTypeD; + using ElementM = typename TiledMMAQK::ValTypeA; + using ElementP = float; + using FragP = decltype(make_subgroup_tensor(make_fragment_like(typename FragS::Base{}), + decltype(FragS{}.tv_layout()){})); + + using FragPRow = decltype(reduce<1>(FragP{}, sycl::plus{})); + using FragPCol = decltype(reduce<0>(FragP{}, sycl::plus{})); + using SingleFragPV = FragC; // (atom val,q',v') + using SingleFragAFloat = std::remove_reference_t(typename SingleFragPV::Base{}), decltype(SingleFragPV{}.tv_layout()){}))>; + using SingleFragA = conditional_t; + using FragA = expand_sg_fragment_t; // (atom val,q',v',VV) + using FragARow = decltype(reduce<1>(FragA{}, sycl::plus{})); + using ElementA = typename SingleFragA::value_type; + + static constexpr int SGTilePVQ = get<0>(shape_div(TileShapePV{}, shape(SubgroupLayoutPV{})))(); + using QuantizedPVOperation = XE_DPAS_TT; + using QuantizedTiledMMAPV = + typename TiledMMAHelper, Layout, SubgroupLayoutPV>::TiledMMA; + using QuantizedSingleFragA = FragC; + + static_assert(size(SingleFragA{}.shape()) == size(QuantizedSingleFragA{}.shape()), + "Quantized PV accumulator must match float PV tile size"); + + static constexpr bool CausalMask = CausalMask_; + static constexpr bool CachedKV = CachedKV_; + static constexpr bool PagedKV = PagedKV_; + // The qtile64 Q-staging experiment stays disabled until the Xe SLM copy layout + // matches the fragment layout expected by make_A_slm_copies(). + static constexpr bool UseInt8PV = UseInt8PV_; + static constexpr bool WriteBackInt8PV = WriteBackInt8PV_; + static constexpr bool ExecuteInt8PV = ExecuteInt8PV_; + + // User-facing arguments + struct Arguments { + float const scale; + float const* mask = nullptr; + // Dense path can treat Q/K as native fp16/bf16. Sparse SAGE consumes INT8 Q/K, + // so every logical block needs dequant scales. scale_block_size is the token + // granularity of those Q/K scale tensors. + int scale_block_size = 0; + float const* qscale = nullptr; + float const* kscale = nullptr; + float const* vscale = nullptr; + // Sparse-only routing metadata. lut stores delta-encoded logical K/V block + // ids per sparse Q row; valid_block_num stores the live length of each row. + // Dense FMHA iterates every K block in-order and does not need either tensor. + int const* lut = nullptr; + int const* valid_block_num = nullptr; + // Logical routing-grid shape, not MMA tile shape. num_q_blocks indexes LUT + // rows, num_k_blocks indexes the logical K/V blocks referenced by those rows. + int num_q_blocks = 0; + int num_k_blocks = 0; + // Sparse-only override for how many Q tokens map to one LUT row. When zero + // we fall back to scale_block_size so routing rows and quant rows coincide. + int sparse_q_block_size = 0; + bool canonical_nhd_k = false; + int const* ptr_page_table = nullptr; + int page_size = 0; + int const* num_pages_per_seq = nullptr; + }; + + // Kernel-facing parameters + using Params = Arguments; + + // SLM data + struct EmptySharedStorage {}; + using SharedStorage = EmptySharedStorage; + + Params params; + SharedStorage& shared_storage; + + // + // Methods + // + + SPARSESAGEV1FwdMainloop(Params const& params_, SharedStorage& shared_storage_) + : params(params_), shared_storage(shared_storage_) {} + + static constexpr Params to_underlying_arguments(Arguments const& args, void* /* workspace */) { + constexpr double kLog2e = 1.4426950408889634074; // log_2(e) + float val = args.scale * static_cast(kLog2e); + return Params{val, args.mask, args.scale_block_size, args.qscale, args.kscale, args.vscale, + args.lut, args.valid_block_num, args.num_q_blocks, args.num_k_blocks, args.sparse_q_block_size, + args.canonical_nhd_k, + args.ptr_page_table, args.page_size, args.num_pages_per_seq}; + } + + CUTLASS_HOST_DEVICE static bool can_implement(Arguments const&) { return true; } + + // DPAS kernels run as SIMD16, so pack two SIMD16 vectors into one temporary SIMD32 vector + // before issuing a single vISA SIMD32 instruction. + CUTLASS_DEVICE static void exp2_pair_simd32_asm(ElementP x0, ElementP x1, ElementP& y0, ElementP& y1) { +#if defined(__SYCL_DEVICE_ONLY__) && defined(SYCL_INTEL_TARGET) + asm( + "{\n" + ".decl OUT0 v_type=G type=F num_elts=16 alias=<%0,0>\n" + ".decl OUT1 v_type=G type=F num_elts=16 alias=<%1,0>\n" + ".decl IN0 v_type=G type=F num_elts=16 alias=<%2,0>\n" + ".decl IN1 v_type=G type=F num_elts=16 alias=<%3,0>\n" + ".decl TMP_IN v_type=G type=F num_elts=32 align=64\n" + ".decl TMP_OUT v_type=G type=F num_elts=32 align=64\n" + "mov (M1_NM, 16) TMP_IN(0,0)<1> IN0(0,0)<1;1,0>\n" + "mov (M1_NM, 16) TMP_IN(1,0)<1> IN1(0,0)<1;1,0>\n" + "exp (M1_NM, 32) TMP_OUT(0,0)<1> TMP_IN(0,0)<1;1,0>\n" + "mov (M1_NM, 16) OUT0(0,0)<1> TMP_OUT(0,0)<1;1,0>\n" + "mov (M1_NM, 16) OUT1(0,0)<1> TMP_OUT(1,0)<1;1,0>\n" + "}\n" + : "=rw"(y0), "=rw"(y1) + : "rw"(x0), "rw"(x1)); +#else + y0 = sycl::native::exp2(x0); + y1 = sycl::native::exp2(x1); +#endif + } + + CUTLASS_DEVICE + int get_physical_k_tile(int K, int l_coord, int seq_len_kv_cache) { + int next_page_logical_idx = K * get<1>(TileShapeQK{}) / params.page_size; + // get<1>(TileShapeQK{}) usually smaller than page_size. + // assuming page_size is multiple of get<1>(TileShapeQK{}) + int tiles_per_page = params.page_size / get<1>(TileShapeQK{}); + int batch_offset = + params.num_pages_per_seq ? params.num_pages_per_seq[l_coord] : l_coord * (seq_len_kv_cache / params.page_size); + + return params.ptr_page_table[batch_offset + next_page_logical_idx] * tiles_per_page + K % tiles_per_page; + } + + CUTLASS_DEVICE + static int logical_block_from_delta_row(int const* row, int valid_blocks, int idx) { + int logical_block = 0; + for (int i = 0; i <= idx && i < valid_blocks; ++i) { + logical_block += row[i]; + } + return logical_block; + } + + template + CUTLASS_DEVICE void operator()(TensorQ2D const& Q_2D, // (q,d) + TensorK2D const& K_2D, // (k,d) + TensorV2D const& V_2D, // (d,k) + FragA& tArA, // Output accumulator (q,v) + FragARow& tA_max, // Softmax row-wise max accumulator + FragARow& tA_sum, // Softmax row-wise sum accumulator + QVCoord blk_qv, // WG tile indices: (Q,V) + int blk_k0, // K block range: [K0,K1) + int blk_k1, + int total_blk, // Total # of K blocks + int thr_id, int seq_len, int seq_len_kv_cache, int l_coord, float* scaleQ, + float* scaleK, float* scaleV, int full_tile_offset, int discard_seq_coord, + int const* lut_rows_base = nullptr, int const* valid_blocks_base = nullptr, + int sparse_q_rows_in_tile = 1, + TensorK_cache2D const& K_cache_2D = TensorK_cache2D{}, + TensorV_cache2D const& V_cache_2D = TensorV_cache2D{}) { + using namespace sycl::ext::oneapi::this_work_item; + + // Short dimension names: + // q = sequence len dimension for Q + // k = sequence len dimension for K + // d = head size dimension for K/Q + // v = head size dimension for V + // VV = MMA tile indices for V + // Capital letters (Q, K, ...) refer to WG block indices. + // Primed letters (q', k', ...) refer to atom block indices. + + auto tile_shape_v = make_shape(get<1>(TileShapePV{}) * C{}, get<2>(TileShapePV{})); + + /* Create proxy coordinate tensors for Q/K/P/V */ + Tensor cQ = make_identity_tensor(Q_2D.shape()); // (q,d) + Tensor cK = make_identity_tensor(K_2D.shape()); // (k,d) + Tensor cV = make_identity_tensor(V_2D.shape()); // (v,k) + Tensor cK_cache = make_identity_tensor(K_cache_2D.shape()); // (k,d) + Tensor cV_cache = make_identity_tensor(V_cache_2D.shape()); // (v,k) + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); // (q,k) + Tensor cPV = make_identity_tensor(select<0, 1>(TileShapePV{})); + + /* Partition global tensors into workgroup tiles */ + Tensor gQ = local_tile(cQ, TileShapeQK{}, append(blk_qv, _), Step<_1, X, _1>{}); // (q,d,D) + Tensor gK = local_tile(cK, TileShapeQK{}, make_coord(_, _, _), Step{}); // (k,d,K,D) + Tensor gV = local_tile(cV, tile_shape_v, make_coord(get<1>(blk_qv), _)); // (v,k,K) + Tensor gV_split = local_tile(gV, TileShapePV{}, make_coord(_, _, 0), Step{}); // (v,k,VV,K) + + Tensor gK_cache = local_tile(cK_cache, TileShapeQK{}, make_coord(_, _, _), Step{}); // (k,d,K,D) + Tensor gV_cache = local_tile(cV_cache, tile_shape_v, make_coord(get<1>(blk_qv), _)); // (v,k,K) + Tensor gV_cache_split = local_tile(gV_cache, TileShapePV{}, make_coord(_, _, 0), Step{}); // (v,k,VV,K) + + /* Create global -> register copies */ + TiledCopyQ copy_q{Q_2D}; + TiledCopyK copy_k{K_2D}; + TiledCopyK_cache copy_k_cache{K_cache_2D}; + [[maybe_unused]] TiledCopyV copy_v{V_2D}; + [[maybe_unused]] TiledCopyV_cache copy_v_cache{V_cache_2D}; + + /* Create MMAs */ + TiledMMAQK mma_qk{}; + [[maybe_unused]] TiledMMAPV mma_pv{}; + [[maybe_unused]] QuantizedTiledMMAPV mma_pv_i8{}; + + /* Slice TiledCopy/TiledMMA operations down to to work-item level */ + auto thr_copy_q = copy_q.get_slice(thr_id); + auto thr_copy_k = copy_k.get_slice(thr_id); + [[maybe_unused]] auto thr_copy_v = copy_v.get_slice(thr_id); + auto thr_copy_k_cache = copy_k_cache.get_slice(thr_id); + [[maybe_unused]] auto thr_copy_v_cache = copy_v_cache.get_slice(thr_id); + auto thr_mma_qk = mma_qk.get_slice(thr_id); + [[maybe_unused]] auto thr_mma_pv = mma_pv.get_slice(thr_id); + [[maybe_unused]] auto thr_mma_pv_i8 = mma_pv_i8.get_slice(thr_id); + + /* Partition coordinate tensors for copy */ + auto tQgQ = thr_copy_q.partition_S(gQ); // (atom_val,q',d',D) + auto tKgK = thr_copy_k.partition_S(gK); // (atom_val,k',d',K,D) + [[maybe_unused]] auto tVgV = thr_copy_v.partition_S(gV_split); // (atom_val,v',k',VV,K) + auto tKgK_cache = thr_copy_k_cache.partition_S(gK_cache); + [[maybe_unused]] auto tVgV_cache = thr_copy_v_cache.partition_S(gV_cache_split); + + /* Create register fragments for MMA and copies */ + auto tQrQ = thr_copy_q.partition_sg_fragment_D(gQ(_, _, 0)); + auto tSrQ = thr_mma_qk.partition_sg_fragment_A(gQ(_, _, 0)); + + auto tKrK = thr_copy_k.partition_sg_fragment_D(gK(_, _, 0, 0)); + auto tSrK = thr_mma_qk.partition_sg_fragment_B(gK(_, _, 0, 0)); + + auto tSrS = thr_mma_qk.partition_sg_fragment_C(cP); + + FragP tPrS = make_subgroup_tensor(make_fragment_like(tSrS.tensor()), tSrS.tv_layout()); + + [[maybe_unused]] auto tArP = thr_mma_pv.partition_sg_fragment_A(cP); + [[maybe_unused]] auto tArP_i8 = thr_mma_pv_i8.partition_sg_fragment_A(cP); + + [[maybe_unused]] auto tVrV = thr_copy_v.partition_sg_fragment_D(gV_split(_, _, 0, 0)); + [[maybe_unused]] auto tArV = thr_mma_pv.partition_sg_fragment_B(gV_split(_, _, 0, 0)); + [[maybe_unused]] auto tCrA = thr_mma_pv_i8.partition_C(cPV); + + /* Create TiledCopy objects for prefetches */ + auto prefetch_q = make_block_2d_prefetch(copy_q); + auto prefetch_k = make_block_2d_prefetch(copy_k); + [[maybe_unused]] auto prefetch_v = make_block_2d_prefetch(copy_v); + auto prefetch_k_cache = make_block_2d_prefetch(copy_k_cache); + [[maybe_unused]] auto prefetch_v_cache = make_block_2d_prefetch(copy_v_cache); + + /* Partition global tensors for prefetch */ + auto pQgQ = prefetch_q.get_slice(thr_id).partition_S(gQ); + auto pKgK = prefetch_k.get_slice(thr_id).partition_S(gK); + [[maybe_unused]] auto pVgV = prefetch_v.get_slice(thr_id).partition_S(gV_split); + auto pKgK_cache = prefetch_k_cache.get_slice(thr_id).partition_S(gK_cache); + [[maybe_unused]] auto pVgV_cache = prefetch_v_cache.get_slice(thr_id).partition_S(gV_cache_split); + + // ------ + // Kernel + // ------ + + /* Initialization steps for first block: Q/K prefetch, O init */ + /* TODO: limit D prefetch for large head size, and reorder K prefetches */ + int kblocks_cache = ceil_div(seq_len_kv_cache, get<1>(TileShapeQK{})); + for (int D = 0; D < size<3>(pQgQ); D++) { + prefetch(prefetch_q, pQgQ(_, _, _, D)); + } + if (lut_rows_base == nullptr) { + for (int D = 0; D < size<4>(pKgK); D++) { + CUTLASS_PRAGMA_UNROLL + for (int K = 0; K < Stages; K++) { + if (K < kblocks_cache) { + if constexpr (PagedKV) { + int physical_K_tile = get_physical_k_tile(K, l_coord, seq_len_kv_cache); + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, physical_K_tile, D)); + } else { + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, K, D)); + } + } else { + prefetch(prefetch_k, pKgK(_, _, _, K - kblocks_cache, D)); + } + } + } + } + if (blk_k0 == 0) { + clear(tArA); + fill(tA_max, cutlass::platform::numeric_limits::lowest()); + clear(tA_sum); + } + + // Sparse row assignment differs from dense FMHA. Dense only needs the MMA tile + // coordinates. Sparse additionally needs to know which routing row each + // subgroup is responsible for within the current Q tile. + bool check_remainder_k = (seq_len % get<1>(TileShapeQK{}) != 0); + int q_sg_tile = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{}))); + int sparse_q_block_size = params.sparse_q_block_size > 0 ? params.sparse_q_block_size : params.scale_block_size; + // How many sparse LUT rows are covered by this workgroup tile. + int q_blocks_per_wg_tile = + sparse_q_block_size > 0 ? cute::max(1, int(get<0>(TileShapeQK{})) / sparse_q_block_size) : 1; + // How many subgroup-row slices belong to one sparse LUT row. + int sg_rows_per_sparse_q_block = + sparse_q_block_size > 0 ? cute::max(1, sparse_q_block_size / q_sg_tile) : 1; + // Map subgroup id -> sparse row id within this workgroup tile. + int subgroup_q_row_in_tile = get_sub_group_id() / sg_rows_per_sparse_q_block; + subgroup_q_row_in_tile = cute::min(subgroup_q_row_in_tile, q_blocks_per_wg_tile - 1); + + float dq_scale = params.scale_block_size + ? scaleQ[(get<0>(blk_qv) * get<0>(TileShapeQK{}) + get_sub_group_id() * q_sg_tile) / + params.scale_block_size] * + params.scale + : 1.f; + auto prefetch_sparse_k_block = [&](int logical_block) { + if (logical_block < 0 || logical_block >= total_blk) return; + for (int D = 0; D < size<4>(pKgK); D++) { + if constexpr (CachedKV) { + if (logical_block < kblocks_cache) { + int physical_block = logical_block; + if constexpr (PagedKV) { + physical_block = get_physical_k_tile(logical_block, l_coord, seq_len_kv_cache); + } + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, physical_block, D)); + } else { + prefetch(prefetch_k, pKgK(_, _, _, logical_block - kblocks_cache, D)); + } + } else { + prefetch(prefetch_k, pKgK(_, _, _, logical_block - kblocks_cache, D)); + } + } + }; // end of prefetch_sparse_k_block + /* Main loop body */ + auto mainloop_body = [&](auto cached_k, int K, bool first_block, bool subgroup_selected, + int sparse_prefetch_block, auto& copy_k_cur, auto& copy_v_cur, auto& prefetch_v_cur, + auto& tKgK_cur, auto& tVgV_cur, auto& pVgV_cur) { + /* Split barrier to keep threads together */ + barrier_arrive(ScopeWorkgroup); + constexpr bool is_cache = decltype(cached_k)::value; + + int k_idx; + if constexpr (is_cache) { + k_idx = K; + if constexpr (PagedKV) { + k_idx = get_physical_k_tile(K, l_coord, seq_len_kv_cache); + } + } else { + k_idx = K - kblocks_cache; + } + int scalek_idx = params.scale_block_size ? K * get<1>(TileShapeQK{}) / params.scale_block_size : 0; + + /* GEMM 1: S = K * Q */ + clear(tSrS); + CUTLASS_PRAGMA_UNROLL + for (int D = 0; D < size<4>(tKgK); D++) { + copy(copy_q, tQgQ(_, _, _, D), tQrQ); + copy(copy_k_cur, tKgK_cur(_, _, _, k_idx, D), tKrK); + reorder(tQrQ, tSrQ); + reorder(tKrK, tSrK); + cute::gemm(mma_qk, tSrQ, tSrK, tSrS); + } + + /* V prefetch for GEMM 2 */ + CUTLASS_PRAGMA_UNROLL + for (int VV = 0; VV < VTiles; VV++) { + prefetch(prefetch_v_cur, pVgV_cur(_, _, _, VV, k_idx)); + } + if (subgroup_selected) { + reorder(tSrS, tPrS); + if constexpr (!CausalMask) { + if (params.scale_block_size) { + if (params.scale_block_size == 1) { + // Need to get global col and row indices to mask the elements + Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); + Tensor gP = local_tile(cPgP, take<0, 2>(TileShapeQK{}), make_coord(get<0>(blk_qv), K)); + auto cS_thread = thr_mma_qk.partition_C(gP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + tPrS(i) *= ElementP(scaleQ[row_idx] * scaleK[col_idx]); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); ++i) { + tPrS(i) *= params.scale; + } + } else { + float _scale = dq_scale * scaleK[scalek_idx]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); i++) tPrS(i) *= _scale; + } + } + } else { + float _scale = dq_scale * scaleK[scalek_idx]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); i++) tPrS(i) *= _scale; + } + + /* Causal masking - only in non-cache mode */ + if constexpr (!is_cache && CausalMask) { + if (K == total_blk - 1) { + // Need to get global col and row indices to mask the elements + Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); + Tensor gP = local_tile(cPgP, take<0, 2>(TileShapeQK{}), make_coord(get<0>(blk_qv), K)); + auto cS_thread = thr_mma_qk.partition_C(gP); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + if (col_idx - seq_len_kv_cache - full_tile_offset > row_idx - discard_seq_coord) { + tPrS(i) = ElementP(-INFINITY); + } + } + } + } else { + if constexpr (FullMask_) { + // Need to get global col and row indices to mask the elements + Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len)); + Tensor gP = local_tile(cPgP, take<0, 2>(TileShapeQK{}), make_coord(get<0>(blk_qv), K)); + auto cS_thread = thr_mma_qk.partition_C(gP); + int row_idx_begin = get<0>(cS_thread(0)); + int row_idx_end = row_idx_begin + q_sg_tile; + int col_idx_begin = get<1>(cS_thread(0)); + int col_idx_end = col_idx_begin + get<1>(TileShapeQK{}); + if (row_idx_end <= seq_len && col_idx_end <= seq_len) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + tPrS(i) += ElementP(params.mask[col_idx + row_idx * seq_len]); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + tPrS(i) += (row_idx < seq_len && col_idx < seq_len) + ? ElementP(params.mask[col_idx + row_idx * seq_len]) + : ElementP(-INFINITY); + } + } + } + } + + /* k masking for remainder tiles */ + if constexpr (!is_cache) { + if (check_remainder_k && K == total_blk - 1) { + FragPCol k_rem_mask; + int k_val = get<0>(tKgK_cur(0, 0, 0, k_idx, 0)) + kblocks_cache * get<1>(TileShapeQK{}); + int k = k_val + get_sub_group().get_local_id()[0]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < k_rem_mask.size(); i++, k += intel::sg_size) { + k_rem_mask(i) = (k < seq_len) ? ElementP(sycl::nan(0u)) : ElementP(-INFINITY); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tPrS.size(); i++) { + tPrS(i) = sycl::fmin(tPrS(i), broadcast<1>(k_rem_mask, tPrS, i)); + } + } + } + + /* Apply softmax and scaling (tA rescaling fused into GEMM2 VTile loop) */ + auto rescale = softmax(first_block, tPrS, tA_max, tA_sum); + if constexpr (!UseInt8PV) { + reorder(tPrS, tArP); + } else { + reorder(tPrS, tArP_i8); + } + + /* GEMM 2: A += P * V, split in v dimension. + tArA rescaling is fused to per-VTile */ + if constexpr (!UseInt8PV) { + CUTLASS_PRAGMA_UNROLL + for (int VV = 0; VV < VTiles; VV++) { + copy(copy_v_cur, tVgV_cur(_, _, _, VV, k_idx), tVrV); + reorder(tVrV, tArV); + if (!first_block) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tArA.size() / VTiles; i++) tArA(_, _, _, VV)(i) *= broadcast<0>(rescale, tArA, i); + } + + cute::gemm(mma_pv, tArP, tArV, tArA(_, _, _, VV)); + } + } else { + constexpr ElementA kInvVQuantScale = ElementA(1.0f / 127.0f); + int scalev_head_dim = int(size<0>(V_2D)); + int v_block_base = int(get<1>(blk_qv)) * int(get<1>(TileShapePV{})) * VTiles; + int scalev_block_base = scalek_idx * scalev_head_dim; + CUTLASS_PRAGMA_UNROLL + for (int VV = 0; VV < VTiles; VV++) { + auto tArA_v = tArA(_, _, _, VV); + if (!first_block) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tArA.size() / VTiles; i++) tArA_v(i) *= broadcast<0>(rescale, tArA, i); + } + + auto tArAcc = thr_mma_pv_i8.partition_sg_fragment_C(cPV); + auto tVrV_i8 = thr_copy_v.partition_sg_fragment_D(gV_split(_, _, 0, 0)); + auto tArV_i8 = thr_mma_pv_i8.partition_sg_fragment_B(gV_split(_, _, 0, 0)); + copy(copy_v_cur, tVgV_cur(_, _, _, VV, k_idx), tVrV_i8); + reorder(tVrV_i8, tArV_i8); + clear(tArAcc); + + if constexpr (ExecuteInt8PV) { + cute::gemm(mma_pv_i8, tArP_i8, tArV_i8, tArAcc); + } + + if constexpr (WriteBackInt8PV) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tArAcc.size(); ++i) { + int local_v = int(get<1>(tCrA(i))); + int scalev_idx = scalev_block_base + v_block_base + VV * int(get<1>(TileShapePV{})) + local_v; + tArA_v(i) += ElementA(tArAcc(i)) * ElementA(scaleV[scalev_idx]) * kInvVQuantScale; + } + } + } + } + } + + /* K prefetch */ + if (lut_rows_base == nullptr) { + int K_next = K + Stages; + for (int D = 0; D < size<4>(pKgK); D++) { + if constexpr (is_cache) { + bool is_cache_next = K_next < kblocks_cache; + int physical_K_next = K_next; + if constexpr (PagedKV) { + if (is_cache_next) { + physical_K_next = get_physical_k_tile(K_next, l_coord, seq_len_kv_cache); + } + } + if (is_cache_next) { + prefetch(prefetch_k_cache, pKgK_cache(_, _, _, physical_K_next, D)); + } else { + prefetch(prefetch_k, pKgK(_, _, _, K_next - kblocks_cache, D)); + } + } else { + prefetch(prefetch_k, pKgK(_, _, _, K_next - kblocks_cache, D)); + } + } + } else { + if (sparse_prefetch_block < total_blk) { + prefetch_sparse_k_block(sparse_prefetch_block); + } + } + + barrier_wait(ScopeWorkgroup); + }; // end of mainloop_body + + /* Main loop, blocked in k. */ + if (lut_rows_base != nullptr && valid_blocks_base != nullptr) { + bool subgroup_started = false; + if (sparse_q_rows_in_tile == 1) { + // Linear sparse traversal: one workgroup owns exactly one sparse LUT row. + // This is the closest sparse analogue to the dense path because all + // participating subgroups walk the same sequence of selected K blocks. + int const* row_ptr = lut_rows_base; + int row_valid = valid_blocks_base[0]; + int row_pos = 0; + int row_cur_block = row_valid > 0 ? logical_block_from_delta_row(row_ptr, row_valid, 0) : total_blk; + int prefetch_pos = row_pos; + int prefetch_cur_block = row_cur_block; + + auto advance_single_sparse_block = [&](int& frontier_pos, int& frontier_cur_block) { + if (frontier_cur_block >= total_blk) return; + frontier_pos += 1; + if (frontier_pos < row_valid) { + frontier_cur_block += row_ptr[frontier_pos]; + } else { + frontier_cur_block = total_blk; + } + }; + + auto pop_single_sparse_block = [&](int& frontier_pos, int& frontier_cur_block) { + int block = frontier_cur_block; + advance_single_sparse_block(frontier_pos, frontier_cur_block); + return block; + }; + + for (int stage = 0; stage < Stages; ++stage) { + int sparse_prefetch_block = pop_single_sparse_block(prefetch_pos, prefetch_cur_block); + if (sparse_prefetch_block >= total_blk) break; + prefetch_sparse_k_block(sparse_prefetch_block); + } + + while (row_cur_block < total_blk) { + int next_block = row_cur_block; + bool subgroup_selected = (subgroup_q_row_in_tile == 0); + bool first_selected_block = !subgroup_started; + int sparse_prefetch_block = pop_single_sparse_block(prefetch_pos, prefetch_cur_block); + int K = next_block; + if constexpr (CachedKV) { + if (K < kblocks_cache) { + if (K >= blk_k0 && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k_cache, copy_v_cache, prefetch_v_cache, tKgK_cache, + tVgV_cache, pVgV_cache); + subgroup_started = true; + } + } else { + K += kblocks_cache; + if (K >= (blk_k0 > kblocks_cache ? blk_k0 : kblocks_cache) && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k, copy_v, prefetch_v, tKgK, tVgV, pVgV); + subgroup_started = true; + } + } + } else { + if (K >= blk_k0 && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k, copy_v, prefetch_v, tKgK, tVgV, pVgV); + subgroup_started = true; + } + } + + advance_single_sparse_block(row_pos, row_cur_block); + } + } else { + // Merged multi-row sparse traversal: one workgroup covers multiple sparse + // Q rows, so each row has its own LUT frontier. We merge those frontiers + // by logical K block and let each subgroup participate only when the + // current merged block belongs to its assigned sparse row. + static constexpr int kMaxSparseRowsPerTile = cute::max(1, int(get<0>(TileShapeQK{})) / 64); + int const* row_ptrs[kMaxSparseRowsPerTile]; + int active_rows[kMaxSparseRowsPerTile]; + int active_row_count = 0; + int row_valid[kMaxSparseRowsPerTile]; + int row_pos[kMaxSparseRowsPerTile]; + int row_cur_block[kMaxSparseRowsPerTile]; + + for (int row = 0; row < kMaxSparseRowsPerTile; ++row) { + row_ptrs[row] = lut_rows_base + row * params.num_k_blocks; + if (row < sparse_q_rows_in_tile) { + row_valid[row] = valid_blocks_base[row]; + row_pos[row] = 0; + row_cur_block[row] = + row_valid[row] > 0 ? logical_block_from_delta_row(row_ptrs[row], row_valid[row], 0) : total_blk; + if (row_valid[row] > 0) { + active_rows[active_row_count++] = row; + } + } else { + row_valid[row] = 0; + row_pos[row] = 0; + row_cur_block[row] = total_blk; + } + } + + auto find_sparse_block = [&](int* frontier_pos, int* frontier_cur_block) { + int block = total_blk; + for (int active = 0; active < active_row_count; ++active) { + int row = active_rows[active]; + if (frontier_pos[row] < row_valid[row]) { + block = cute::min(block, frontier_cur_block[row]); + } + } + return block; + }; + + auto advance_sparse_block = [&](int block, int* frontier_pos, int* frontier_cur_block) { + if (block >= total_blk) return; + for (int active = 0; active < active_row_count; ++active) { + int row = active_rows[active]; + if (frontier_pos[row] < row_valid[row] && frontier_cur_block[row] == block) { + frontier_pos[row] += 1; + if (frontier_pos[row] < row_valid[row]) { + frontier_cur_block[row] += row_ptrs[row][frontier_pos[row]]; + } else { + frontier_cur_block[row] = total_blk; + } + } + } + }; + + auto pop_sparse_block = [&](int* frontier_pos, int* frontier_cur_block) { + int block = find_sparse_block(frontier_pos, frontier_cur_block); + advance_sparse_block(block, frontier_pos, frontier_cur_block); + return block; + }; + + int prefetch_pos[kMaxSparseRowsPerTile]; + int prefetch_cur_block[kMaxSparseRowsPerTile]; + for (int row = 0; row < kMaxSparseRowsPerTile; ++row) { + prefetch_pos[row] = row_pos[row]; + prefetch_cur_block[row] = row_cur_block[row]; + } + + for (int stage = 0; stage < Stages; ++stage) { + int sparse_prefetch_block = pop_sparse_block(prefetch_pos, prefetch_cur_block); + if (sparse_prefetch_block >= total_blk) break; + prefetch_sparse_k_block(sparse_prefetch_block); + } + + int next_block = find_sparse_block(row_pos, row_cur_block); + while (next_block < total_blk) { + // subgroup_q_row_in_tile is the row assignment computed earlier. + // subgroup_selected gates the dense GEMM/softmax/PV work so only the + // subgroups whose sparse row requested next_block contribute to it. + bool subgroup_selected = subgroup_q_row_in_tile < sparse_q_rows_in_tile && + row_pos[subgroup_q_row_in_tile] < row_valid[subgroup_q_row_in_tile] && + row_cur_block[subgroup_q_row_in_tile] == next_block; + bool first_selected_block = !subgroup_started; + int sparse_prefetch_block = pop_sparse_block(prefetch_pos, prefetch_cur_block); + int K = next_block; + if constexpr (CachedKV) { + if (K < kblocks_cache) { + if (K >= blk_k0 && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k_cache, copy_v_cache, prefetch_v_cache, tKgK_cache, + tVgV_cache, pVgV_cache); + subgroup_started = subgroup_started || subgroup_selected; + } + } else { + K += kblocks_cache; + if (K >= (blk_k0 > kblocks_cache ? blk_k0 : kblocks_cache) && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k, copy_v, prefetch_v, tKgK, tVgV, pVgV); + subgroup_started = subgroup_started || subgroup_selected; + } + } + } else { + if (K >= blk_k0 && K < blk_k1) { + mainloop_body(std::bool_constant{}, K, first_selected_block, subgroup_selected, + sparse_prefetch_block, copy_k, copy_v, prefetch_v, tKgK, tVgV, pVgV); + subgroup_started = subgroup_started || subgroup_selected; + } + } + + advance_sparse_block(next_block, row_pos, row_cur_block); + next_block = find_sparse_block(row_pos, row_cur_block); + } + } + } else { + if constexpr (CachedKV) { + for (int K = blk_k0; K < kblocks_cache; K++) { + mainloop_body(std::bool_constant{}, K, K == blk_k0, true, total_blk, copy_k_cache, copy_v_cache, + prefetch_v_cache, tKgK_cache, tVgV_cache, pVgV_cache); + } + } + for (int K = (blk_k0 > kblocks_cache ? blk_k0 : kblocks_cache); K < blk_k1; K++) { + mainloop_body(std::bool_constant{}, K, + K == (blk_k0 > kblocks_cache ? blk_k0 : kblocks_cache), true, total_blk, copy_k, copy_v, + prefetch_v, tKgK, tVgV, pVgV); + } + } + } + + // Single step of blocked softmax. + CUTLASS_DEVICE + FragARow softmax(bool first_block, // First softmax block? + FragP& tS, // Softmax src/dst block + FragARow& tS_max, // Softmax row-wise max accumulator + FragARow& tS_sum) { // Softmax row-wise sum accumulator + /* Compute row-wise maxima for this block */ + + auto tS_bmax = reduce<1>(tS, sycl::maximum{}); + FragARow rescale; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_max.size(); i++) { + ElementA new_max = sycl::max(tS_max(i), ElementA(tS_bmax(i))); + rescale(i) = sycl::native::exp2(tS_max(i) - new_max); + tS_max(i) = new_max; + } + + /* Scale S and subtract maxima, then exponentiate */ + static_assert(FragP{}.size() % 2 == 0, "FragP size must be even for pairwise SIMD32 exp."); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS.size(); i += 2) { + ElementP x0 = tS(i) - broadcast<0>(tS_max, tS, i); + ElementP x1 = tS(i + 1) - broadcast<0>(tS_max, tS, i + 1); + exp2_pair_simd32_asm(x0, x1, tS(i), tS(i + 1)); + } + + if constexpr (UseInt8PV) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS.size(); i++) { + tS(i) *= ElementP(127.0f); + } + } + + /* Rescale existing S sums */ + if (!first_block) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_sum.size(); i++) { + tS_sum(i) *= rescale(i); + } + } + + /* Update sums */ + auto tS_bsum = reduce<1>(tS, sycl::plus{}); + if constexpr (UseInt8PV) { + constexpr ElementA kInvPQuantScale = ElementA(1.0f / 127.0f); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_sum.size(); i++) tS_sum(i) += ElementA(tS_bsum(i)) * kInvPQuantScale; + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_sum.size(); i++) tS_sum(i) += tS_bsum(i); + } + + return rescale; + } +}; +} // namespace cutlass::fmha::collective + +///////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp index 5effa7975..792e77bf6 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp @@ -228,6 +228,44 @@ void sage_prefill_varlen(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, int head_dim, float softmax_scale, bool is_causal, const int* cu_seqlens_q, const int* cu_seqlens_k, float* lse = nullptr); + +void sdpa_impl_qks8_sparse_d64_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, + BTLA_DTYPE pv_dtype = BTLA_DTYPE::F16); + +void sdpa_impl_qks8_sparse_row_linear_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, + BTLA_DTYPE pv_dtype = BTLA_DTYPE::F16); + +void sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, int num_q_blocks, int num_k_blocks, + int q_tile_override, int q_stride_s, int q_stride_d, int q_stride_h, int q_stride_b, int k_stride_s, + int k_stride_d, int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, int v_stride_h, int v_stride_b, + int o_stride_s, int o_stride_d, int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int head_dim, float softmax_scale, bool is_causal, + BTLA_DTYPE pv_dtype = BTLA_DTYPE::F16); + +void sdpa_impl_qks8_sparse_decode_pvhalf(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* K_cache_ptr, + void* V_cache_ptr, void* O_ptr, void* mask, int scale_block_size, + void* qscale, void* kscale, void* lut, void* valid_block_num, + int num_q_blocks, int num_k_blocks, int q_stride_s, int q_stride_d, + int q_stride_h, int q_stride_b, int k_stride_s, int k_stride_d, + int k_stride_h, int k_stride_b, int v_stride_d, int v_stride_s, + int v_stride_h, int v_stride_b, int o_stride_s, int o_stride_d, + int o_stride_h, int o_stride_b, int batch, int num_heads_q, int num_heads_kv, + int seq_len_q, int seq_len_kv, int seq_len_kv_cache, int head_dim, + float softmax_scale, bool is_causal, BTLA_DTYPE pv_dtype = BTLA_DTYPE::F16); #endif // ARK_SYCL_TLA } // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa_sparse.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa_sparse.hpp new file mode 100644 index 000000000..a8b4d2d10 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa_sparse.hpp @@ -0,0 +1,1291 @@ +// SYCL-TLA Flash Attention Wrapper (Prefill + Decode) +// Provides both prefill and decode entry points while accounting for convert_type ODR concerns in sycl-tla +// +// Copyright (C) 2026 Intel Corporation +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include "bestla/bestla.h" +#ifdef ARK_XPU +#include +#endif + +#if defined(ARK_SYCL_TLA) + +#ifndef ARK_SDPA_ENABLE_SPARSE +#define ARK_SDPA_ENABLE_SPARSE 1 +#endif + +#include "cutlass/epilogue/collective/default_epilogue.hpp" +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/util/packed_stride.hpp" +#if defined(ARK_SDPA_ENABLE_DENSE) +#include "stla/xe_sdpa_fwd_mainloop.hpp" +#include "stla/xe_sagev1_fwd_mainloop.hpp" +#include "stla/xe_sage_fwd_kernel.hpp" +#endif +#if defined(ARK_SDPA_ENABLE_SPARSE) +#include "stla/xe_sparse_fmha_fwd_epilogue.hpp" +#include "stla/xe_sparse_sagev1_fwd_mainloop.hpp" +#include "stla/xe_sparse_sage_fwd_kernel.hpp" +#endif +#include "flash_attention_v2/collective/fmha_fusion.hpp" +#if defined(ARK_SDPA_ENABLE_DENSE) +#include "flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp" +#include "flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp" +#endif +#include "flash_attention_v2/kernel/xe_tile_scheduler.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/sycl_event_manager.hpp" +#include +#include + +#include "helper.h" +#include "cutlass/util/command_line.h" +#include "cutlass/util/device_memory.h" +#include "cutlass/util/reference/device/gemm_complex.h" +#include "cutlass/util/reference/device/tensor_compare.h" +#include "sycl_common.hpp" + +#include + +#endif // ARK_SYCL_TLA + +namespace ark { + +/// Flash Attention data type codes (matches Python side) +enum class FlashAttnDtype : int { + FP16 = 0, + BF16 = 1, + FP32 = 2, + FP8_E4M3 = 3, + FP8_E5M2 = 4, +}; + +#if defined(ARK_SYCL_TLA) + +namespace sparse_detail { + +using namespace cute; + +template +struct has_canonical_nhd_k : std::false_type {}; + +template +struct has_canonical_nhd_k().canonical_nhd_k)>> : std::true_type {}; + +template +struct has_sparse_q_block_size : std::false_type {}; + +template +struct has_sparse_q_block_size().sparse_q_block_size)>> : std::true_type {}; + +// Sparse SDPA launch options. +// +// Dense SDPA mainly needs tensor pointers, logical shapes, and an optional mask. +// Sparse SDPA adds routing metadata that tells each workgroup which logical K/V +// blocks are worth visiting for each sparse Q block. The fields below are the +// sparse-only knobs that do not exist, or do not matter, in the dense path. +struct Options { + const void *q = nullptr, *k = nullptr, *v = nullptr; + void* mask = nullptr; + void* o = nullptr; + // Selects the Q tile shape at kernel-dispatch time. In the sparse path this + // also controls how many sparse Q rows one workgroup may cover. + int q_tile_override = 0; + // Quantization / scaling granularity for the dense Q/K tensors. Sparse routing + // uses the same block unit by default unless sparse_q_block_size overrides it. + int scale_block_size = 0; + // Sparse routing granularity along Q. When this differs from scale_block_size + // we decouple "how Q/K are quantized" from "how many Q tokens share one LUT row". + // Example: q_tile=256 with sparse_q_block_size=256 means one sparse row covers + // the full workgroup tile even though scales still live at 64-token granularity. + int sparse_q_block_size = 0; + const void *qscale = nullptr, *kscale = nullptr, *vscale = nullptr; + // Sparse routing metadata is stored per (batch, q_head, q_block). Each LUT row is delta-encoded + // over logical KV blocks and valid_block_num says how many entries in that row are live. + const int *lut = nullptr, *valid_block_num = nullptr; + // Number of logical sparse Q rows and logical sparse K columns represented by + // the LUT. These are not tile counts in the dense GEMM sense; they describe the + // routing grid that the sparse mainloop walks. + int num_q_blocks = 0, num_k_blocks = 0; + const void *block_K = nullptr, *block_V = nullptr; + const int *page_table = nullptr, *num_pages_per_seq = nullptr; + bool is_causal = false; + bool varlen = false; + bool use_paged_kv = false; + bool use_tensor_strides = false; + int batch = 0, num_heads_q = 0, num_heads_kv = 0, seq_len_qo = 0, seq_len_kv = 0, seq_len_kv_cache = 0, page_size = 0, + head_size_qk = 0, head_size_vo = 0; + int total_seqlen_q = 0, total_seqlen_kv = 0, total_seqlen_kv_cache = 0; + int max_seqlen_q = 0, max_seqlen_kv = 0, max_seqlen_kv_cache = 0; + int q_stride_s = 0, q_stride_d = 1, q_stride_h = 0, q_stride_b = 0; + int k_stride_s = 0, k_stride_d = 1, k_stride_h = 0, k_stride_b = 0; + int v_stride_d = 1, v_stride_s = 0, v_stride_h = 0, v_stride_b = 0; + int o_stride_s = 0, o_stride_d = 1, o_stride_h = 0, o_stride_b = 0; + float softmax_scale = 0.0f; + bool persistent = false; + + void print(std::ostream& os = std::cout) const { + os << std::boolalpha << "Options {\n" + << " q: " << q << "\n" + << " k: " << k << "\n" + << " v: " << v << "\n" + << " o: " << o << "\n" + << " mask: " << mask << "\n" + << " q_tile_override: " << q_tile_override << "\n" + << " sparse_q_block_size: " << sparse_q_block_size << "\n" + << " lut: " << lut << "\n" + << " valid_block_num: " << valid_block_num << "\n" + << " num_q_blocks: " << num_q_blocks << "\n" + << " num_k_blocks: " << num_k_blocks << "\n" + << " block_K: " << block_K << "\n" + << " block_V: " << block_V << "\n" + << " page_table: " << page_table << "\n" + << " num_pages_per_seq: " << num_pages_per_seq << "\n" + << " is_causal: " << is_causal << "\n" + << " varlen: " << varlen << "\n" + << " use_paged_kv: " << use_paged_kv << "\n" + << " use_tensor_strides: " << use_tensor_strides << "\n" + << " batch: " << batch << "\n" + << " num_heads_q: " << num_heads_q << "\n" + << " num_heads_kv: " << num_heads_kv << "\n" + << " seq_len_qo: " << seq_len_qo << "\n" + << " seq_len_kv: " << seq_len_kv << "\n" + << " seq_len_kv_cache: " << seq_len_kv_cache << "\n" + << " page_size: " << page_size << "\n" + << " head_size_qk: " << head_size_qk << "\n" + << " head_size_vo: " << head_size_vo << "\n" + << " total_seqlen_q: " << total_seqlen_q << "\n" + << " total_seqlen_kv: " << total_seqlen_kv << "\n" + << " total_seqlen_kv_cache: " << total_seqlen_kv_cache << "\n" + << " max_seqlen_q: " << max_seqlen_q << "\n" + << " max_seqlen_kv: " << max_seqlen_kv << "\n" + << " max_seqlen_kv_cache: " << max_seqlen_kv_cache << "\n" + << " q_stride: (" << q_stride_s << ", " << q_stride_d << ", " << q_stride_h << ", " << q_stride_b << ")\n" + << " k_stride: (" << k_stride_s << ", " << k_stride_d << ", " << k_stride_h << ", " << k_stride_b << ")\n" + << " v_stride: (" << v_stride_d << ", " << v_stride_s << ", " << v_stride_h << ", " << v_stride_b << ")\n" + << " o_stride: (" << o_stride_s << ", " << o_stride_d << ", " << o_stride_h << ", " << o_stride_b << ")\n" + << " softmax_scale: " << softmax_scale << "\n" + << " persistent: " << persistent << "\n" + << "}\n"; + } +}; + +template +MainloopArguments make_sage_mainloop_arguments(const Options& options) { + if constexpr (has_canonical_nhd_k::value) { + // Sparse mainloop extends the dense-style argument bundle with routing state. + // The extra fields below are exactly what makes the sparse path different: + // scale_block_size / qscale / kscale: dequantization metadata for INT8 Q/K + // lut / valid_block_num: delta-encoded sparse routing table + // num_q_blocks / num_k_blocks: logical routing-grid shape + // sparse_q_block_size: how many Q tokens share one LUT row + return {options.softmax_scale, + static_cast(options.mask), + options.scale_block_size, + static_cast(options.qscale), + static_cast(options.kscale), + static_cast(options.vscale), + options.lut, + options.valid_block_num, + options.num_q_blocks, + // num_k_blocks counts the logical KV blocks visible to sparse routing. For cached decode this + // spans the concatenated cache + current KV space even though the kernel may source data from + // different tensors under the hood. + options.num_k_blocks, + options.sparse_q_block_size, + false, + options.use_paged_kv ? options.page_table : nullptr, + options.use_paged_kv ? options.page_size : 0, + options.use_paged_kv ? options.num_pages_per_seq : nullptr}; + } else { + return {options.softmax_scale, + static_cast(options.mask), + options.scale_block_size, + static_cast(options.qscale), + static_cast(options.kscale), + static_cast(options.vscale), + options.use_paged_kv ? options.page_table : nullptr, + options.use_paged_kv ? options.page_size : 0, + options.use_paged_kv ? options.num_pages_per_seq : nullptr}; + } +} + + template + inline void set_tensor_strides_from_options(const Options& options, StrideQ& stride_Q, StrideK& stride_K, + StrideV& stride_V, StrideO& stride_O) { + stride_Q = cute::make_stride(options.q_stride_s, cute::_1{}, options.q_stride_h, options.q_stride_b); + stride_K = cute::make_stride(options.k_stride_s, cute::_1{}, options.k_stride_h, options.k_stride_b); + stride_V = cute::make_stride(cute::_1{}, options.v_stride_s, options.v_stride_h, options.v_stride_b); + stride_O = cute::make_stride(options.o_stride_s, cute::_1{}, options.o_stride_h, options.o_stride_b); + } + +// 3 input matrices: (K)eys, (Q)ueries and (V)alues. +using LayoutQ = cutlass::layout::RowMajor; +using LayoutK = cutlass::layout::ColumnMajor; +using LayoutV = cutlass::layout::RowMajor; +using LayoutO = cutlass::layout::RowMajor; + +#if defined(ARK_SDPA_ENABLE_DENSE) +template +struct KernelRunner { + using StrideQ = typename FMHAKernel::StrideQ; + using StrideK = typename FMHAKernel::StrideK; + using StrideV = typename FMHAKernel::StrideV; + using StrideO = typename FMHAKernel::StrideO; + + using ElementQ = typename FMHAKernel::ElementQ; + using ElementK = typename FMHAKernel::ElementK; + using ElementV = typename FMHAKernel::ElementV; + using ElementO = typename FMHAKernel::ElementO; + + using CollectiveMainloop = typename FMHAKernel::CollectiveMainloop; + using ElementS = typename CollectiveMainloop::ElementS; + + using ProblemShapeType = cutlass::fmha::kernel::FMHAProblemShape; + + // + // Data members + // + + /// Initialization + StrideQ stride_Q; + StrideK stride_K; + StrideV stride_V; + StrideK stride_K_cache; + StrideV stride_V_cache; + StrideO stride_O; + uint64_t seed = 0; + + // + // Methods + // + + // Note that the GemmUniversalAdapter currently doesn't support flash attention, which is why this + // secondary `run` function is required to launch the kernel. + static void run(typename FMHAKernel::Params params) { + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + + dim3 const block = FMHAKernel::get_block_shape(); + dim3 const grid = FMHAKernel::get_grid_shape(params); + + // configure smem size and carveout + int smem_size = FMHAKernel::SharedStorageSize; + + const auto sycl_block = compat::dim3(block.x, block.y, block.z); + const auto sycl_grid = compat::dim3(grid.x, grid.y, grid.z); + + // Launch parameters depend on whether SYCL compiler supports work-group scratch memory extension + compat::experimental::launch_properties launch_props{ + syclex::work_group_scratch_size(smem_size), + }; + compat::experimental::kernel_properties kernel_props{syclex::sub_group_size, + intelex::grf_size<256>}; + compat::experimental::launch_policy policy{sycl_grid, sycl_block, launch_props, kernel_props}; + auto event = compat::experimental::launch, FMHAKernel>(policy, params); + + EventManager::getInstance().addEvent(event); + } + + template + auto initialize_varlen(const ProblemShape& problem_size, const Options& options) { + ProblemShape problem_size_for_init = problem_size; + get<0>(problem_size_for_init) = 1; + get<3>(problem_size_for_init) = options.total_seqlen_q; + get<4>(problem_size_for_init) = options.total_seqlen_kv; + get<5>(problem_size_for_init) = options.total_seqlen_kv_cache; + + ProblemShapeType problem_size_for_launch; + problem_size_for_launch.batch = get<0>(problem_size); + problem_size_for_launch.num_heads_q = get<1>(problem_size); + problem_size_for_launch.num_heads_kv = get<2>(problem_size); + problem_size_for_launch.seq_len_qo = cutlass::fmha::collective::VariableLength{options.max_seqlen_q}; + problem_size_for_launch.seq_len_kv = cutlass::fmha::collective::VariableLength{options.max_seqlen_kv}; + problem_size_for_launch.seq_len_kv_cache = cutlass::fmha::collective::VariableLength{options.max_seqlen_kv_cache}; + problem_size_for_launch.head_size_qk = get<6>(problem_size); + problem_size_for_launch.head_size_vo = get<7>(problem_size); + + return cute::make_tuple(problem_size_for_init, problem_size_for_launch); + } + + ProblemShapeType initialize(const Options& options) { + auto problem_shape_in = + cute::make_tuple(options.batch, options.num_heads_q, options.num_heads_kv, options.seq_len_qo, + options.seq_len_kv, options.seq_len_kv_cache, options.head_size_qk, options.head_size_vo); + ProblemShapeType shape; + + decltype(problem_shape_in) problem_size; + + if constexpr (isVarLen) { + auto [problem_shape_init, problem_shape_launch] = initialize_varlen(problem_shape_in, options); + problem_size = problem_shape_init; + shape = problem_shape_launch; + } else { + problem_size = problem_shape_in; + shape.batch = options.batch; + shape.num_heads_q = options.num_heads_q; + shape.num_heads_kv = options.num_heads_kv; + shape.seq_len_qo = options.seq_len_qo; + shape.seq_len_kv = options.seq_len_kv; + shape.seq_len_kv_cache = options.seq_len_kv_cache; + shape.head_size_qk = options.head_size_qk; + shape.head_size_vo = options.head_size_vo; + } + + auto [batch, num_heads_q, num_heads_kv, seq_len_qo, seq_len_kv, seq_len_kv_cache, head_size_qk, head_size_vo] = + problem_size; + auto shape_Q = cute::make_shape(seq_len_qo, head_size_qk, num_heads_q, batch); + auto shape_K = cute::make_shape(seq_len_kv, head_size_qk, num_heads_kv, batch); + auto shape_V = cute::make_shape(head_size_vo, seq_len_kv, num_heads_kv, batch); + auto shape_K_cache = cute::make_shape(seq_len_kv_cache, head_size_qk, num_heads_kv, batch); + auto shape_V_cache = cute::make_shape(head_size_vo, seq_len_kv_cache, num_heads_kv, batch); + auto shape_O = cute::make_shape(seq_len_qo, head_size_vo, num_heads_q, batch); + + if constexpr (!isVarLen) { + if (options.use_tensor_strides) { + set_tensor_strides_from_options(options, stride_Q, stride_K, stride_V, stride_O); + } else { + stride_Q = cutlass::make_cute_packed_stride(StrideQ{}, shape_Q); + stride_K = cutlass::make_cute_packed_stride(StrideK{}, shape_K); + stride_V = cutlass::make_cute_packed_stride(StrideV{}, shape_V); + stride_O = cutlass::make_cute_packed_stride(StrideO{}, shape_O); + } + } else { + stride_Q = cutlass::make_cute_packed_stride(StrideQ{}, shape_Q); + stride_K = cutlass::make_cute_packed_stride(StrideK{}, shape_K); + stride_V = cutlass::make_cute_packed_stride(StrideV{}, shape_V); + stride_O = cutlass::make_cute_packed_stride(StrideO{}, shape_O); + } + stride_K_cache = cutlass::make_cute_packed_stride(StrideK{}, shape_K_cache); + stride_V_cache = cutlass::make_cute_packed_stride(StrideV{}, shape_V_cache); + + if constexpr (isVarLen) { + // shape.seq_len_qo.cumulative_length = device_cumulative_seqlen_q.get(); + // shape.seq_len_kv.cumulative_length = device_cumulative_seqlen_kv.get(); + // shape.seq_len_kv_cache.cumulative_length = device_cumulative_seqlen_kv_cache.get(); + } + return shape; + } + + cutlass::Status run(const Options& options, const cutlass::KernelHardwareInfo& hw_info) { + ProblemShapeType shape = initialize(options); + + typename FMHAKernel::Arguments arguments{ + { + shape, + static_cast(options.q), + stride_Q, + static_cast(options.k), + stride_K, + static_cast(options.v), + stride_V, + static_cast(options.o), + stride_O, + static_cast(options.block_K), + stride_K_cache, + static_cast(options.block_V), + stride_V_cache, + }, + {options.softmax_scale, static_cast(options.mask), + options.use_paged_kv ? options.page_table : nullptr, options.use_paged_kv ? options.page_size : 0, + options.use_paged_kv ? options.num_pages_per_seq : nullptr}, + {}, + hw_info}; + // Define device-global scratch memory + size_t workspace_size = FMHAKernel::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + if (!FMHAKernel::can_implement(arguments)) { + std::cout << "Invalid Problem Size: " << options.batch << 'x' << options.num_heads_q << 'x' << options.seq_len_qo + << 'x' << options.seq_len_kv << 'x' << options.head_size_qk << 'x' << options.head_size_vo + << (options.is_causal ? "xCausal" : "xNonCausal") << std::endl; + return cutlass::Status::kErrorInvalidProblem; + } + + // Initialize the workspace + CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace.get())); + + // Convert host-side arguments to device-side arguments to be passed to the kernel + auto params = FMHAKernel::to_underlying_arguments(arguments, workspace.get()); + + run(params); + + return cutlass::Status::kSuccess; + } +}; + +template default */ + int PipelineStages, bool persistent, typename ElementQ = bfloat16_t, typename ElementK = bfloat16_t, + typename ElementV = bfloat16_t, typename ElementO = ElementQ, + typename MMAOperation_ = void, /* void -> default */ + typename StrideQ = Stride, typename StrideK = Stride, + typename StrideV = Stride<_1, int, int, int>, typename StrideO = Stride, + typename GmemTiledCopyQ = void, /* void -> default block 2D */ + typename GmemTiledCopyK = void, typename GmemTiledCopyV = void, typename GmemTiledCopyO = void> +struct FMHAConfig { + static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + using MMAOperation = cute::conditional_t< + is_void_v, + typename cute::conditional_t< + cute::is_same_v || cute::is_same_v, + XE_DPAS_TT, XE_DPAS_TT>, + MMAOperation_>; + using SubgroupLayoutPV = + cute::conditional_t, + decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; + + template + static int run(const Options& options) { + // + // Run examples + // + + // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This + // information is used by the underlying kernel. + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + using ProblemShapeType = cutlass::fmha::kernel::FMHAProblemShape; + + using TiledMMAQK = typename TiledMMAHelper, Layout, SubgroupLayoutQK>::TiledMMA; + using TiledMMAPV = typename TiledMMAHelper, Layout, SubgroupLayoutPV>::TiledMMA; + + static_assert(get<0>(TileShapeOutput{}) == get<0>(TileShapePV{}), + "Output tile and P*V tile have different sizes in Q dimension"); + constexpr int VTiles = get<1>(TileShapeOutput{}) / get<1>(TileShapePV{}); + + auto make_dummy_tensor = [&](auto val, auto stride) { + return make_tensor(make_gmem_ptr(&val), make_layout(repeat>(1), stride)); + }; + + using TensorQ = decltype(make_dummy_tensor(ElementQ{}, StrideQ{})); + using TensorK = decltype(make_dummy_tensor(ElementK{}, StrideK{})); + using TensorV = decltype(make_dummy_tensor(ElementV{}, StrideV{})); + using TensorO = decltype(make_dummy_tensor(ElementO{}, StrideO{})); + using TensorK_cache = TensorK; + using TensorV_cache = TensorV; + using GmemTiledCopyK_cache = GmemTiledCopyK; + using GmemTiledCopyV_cache = GmemTiledCopyV; + using MainloopDispatchPolicy = cutlass::sdpa::XeDefault; + if constexpr (Causal) { + // Mainloop + using CollectiveMainloop = + cutlass::fmha::collective::SDPAFwdMainloop; + + // Epilogue + using CollectiveEpilogue = cutlass::fmha::collective::SparseFMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = conditional_t< + is_same_v, + cutlass::fmha::kernel::XeFMHAFwdDynamicSplitKernel, + cutlass::fmha::kernel::XeFMHAFwdKernel>; + + KernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + if (options.mask) { + // Mainloop + using CollectiveMainloop = + cutlass::fmha::collective::SDPAFwdMainloop; + + // Epilogue + using CollectiveEpilogue = cutlass::fmha::collective::SparseFMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = + conditional_t, + cutlass::fmha::kernel::XeFMHAFwdDynamicSplitKernel, + cutlass::fmha::kernel::XeFMHAFwdKernel>; + + KernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + // Mainloop + using CollectiveMainloop = + cutlass::fmha::collective::SDPAFwdMainloop; + + // Epilogue + using CollectiveEpilogue = cutlass::fmha::collective::SparseFMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = + conditional_t, + cutlass::fmha::kernel::XeFMHAFwdDynamicSplitKernel, + cutlass::fmha::kernel::XeFMHAFwdKernel>; + + KernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } + } + + return 0; + } + + static int run(const Options& options) { + bool cached_kv = options.seq_len_kv_cache > 0; + if constexpr (persistent) { + if (options.use_paged_kv || options.seq_len_kv_cache > 0) { + std::cerr + << "Error: Persistent kernel does not support paged/cached KV cache (use_paged_kv or seq_len_kv_cache > 0)." + << std::endl; + return -1; + } + return run(options); + } else if (options.use_paged_kv && !options.varlen) { + throw std::runtime_error("Paged KV without varlen is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && options.varlen && !cached_kv) { + throw std::runtime_error("Varlen without cached KV is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && !options.varlen && !cached_kv) { + return run(options); + } else if (!options.use_paged_kv && options.varlen && cached_kv) { + throw std::runtime_error("Varlen with cached KV but without paged KV is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && !options.varlen && cached_kv) { + throw std::runtime_error("Cached KV without varlen is not supported yet"); + // return run(options); + } else { + throw std::runtime_error("The combination of options is not supported yet"); + // return run(options); + } + } +}; +#endif // ARK_SDPA_ENABLE_DENSE + +template +struct SageKernelRunner { + using StrideQ = typename FMHAKernel::StrideQ; + using StrideK = typename FMHAKernel::StrideK; + using StrideV = typename FMHAKernel::StrideV; + using StrideO = typename FMHAKernel::StrideO; + + using ElementQ = typename FMHAKernel::ElementQ; + using ElementK = typename FMHAKernel::ElementK; + using ElementV = typename FMHAKernel::ElementV; + using ElementO = typename FMHAKernel::ElementO; + + using CollectiveMainloop = typename FMHAKernel::CollectiveMainloop; + using ElementS = typename CollectiveMainloop::ElementS; + + using ProblemShapeType = cutlass::fmha::kernel::SparseSageProblemShape; + + // + // Data members + // + + /// Initialization + StrideQ stride_Q; + StrideK stride_K; + StrideV stride_V; + StrideK stride_K_cache; + StrideV stride_V_cache; + StrideO stride_O; + uint64_t seed = 0; + + // + // Methods + // + + // Note that the GemmUniversalAdapter currently doesn't support flash attention, which is why this + // secondary `run` function is required to launch the kernel. + static void run(typename FMHAKernel::Params params) { + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + + dim3 const block = FMHAKernel::get_block_shape(); + dim3 const grid = FMHAKernel::get_grid_shape(params); + + // configure smem size and carveout + int smem_size = FMHAKernel::SharedStorageSize; + + const auto sycl_block = compat::dim3(block.x, block.y, block.z); + const auto sycl_grid = compat::dim3(grid.x, grid.y, grid.z); + + // Launch parameters depend on whether SYCL compiler supports work-group scratch memory extension + compat::experimental::launch_properties launch_props{ + syclex::work_group_scratch_size(smem_size), + }; + compat::experimental::kernel_properties kernel_props{syclex::sub_group_size, + intelex::grf_size<256>}; + compat::experimental::launch_policy policy{sycl_grid, sycl_block, launch_props, kernel_props}; + auto event = compat::experimental::launch, FMHAKernel>(policy, params); + + EventManager::getInstance().addEvent(event); + } + + template + auto initialize_varlen(const ProblemShape& problem_size, const Options& options) { + ProblemShape problem_size_for_init = problem_size; + get<0>(problem_size_for_init) = 1; + get<3>(problem_size_for_init) = options.total_seqlen_q; + get<4>(problem_size_for_init) = options.total_seqlen_kv; + get<5>(problem_size_for_init) = options.total_seqlen_kv_cache; + + ProblemShapeType problem_size_for_launch; + problem_size_for_launch.batch = get<0>(problem_size); + problem_size_for_launch.num_heads_q = get<1>(problem_size); + problem_size_for_launch.num_heads_kv = get<2>(problem_size); + problem_size_for_launch.seq_len_qo = cutlass::fmha::collective::VariableLength{options.max_seqlen_q}; + problem_size_for_launch.seq_len_kv = cutlass::fmha::collective::VariableLength{options.max_seqlen_kv}; + problem_size_for_launch.seq_len_kv_cache = cutlass::fmha::collective::VariableLength{options.max_seqlen_kv_cache}; + problem_size_for_launch.head_size_qk = get<6>(problem_size); + problem_size_for_launch.head_size_vo = get<7>(problem_size); + + return cute::make_tuple(problem_size_for_init, problem_size_for_launch); + } + + ProblemShapeType initialize(const Options& options) { + auto problem_shape_in = + cute::make_tuple(options.batch, options.num_heads_q, options.num_heads_kv, options.seq_len_qo, + options.seq_len_kv, options.seq_len_kv_cache, options.head_size_qk, options.head_size_vo); + ProblemShapeType shape; + + decltype(problem_shape_in) problem_size; + + if constexpr (isVarLen) { + auto [problem_shape_init, problem_shape_launch] = initialize_varlen(problem_shape_in, options); + problem_size = problem_shape_init; + shape = problem_shape_launch; + } else { + problem_size = problem_shape_in; + shape.batch = options.batch; + shape.num_heads_q = options.num_heads_q; + shape.num_heads_kv = options.num_heads_kv; + shape.seq_len_qo = options.seq_len_qo; + shape.seq_len_kv = options.seq_len_kv; + shape.seq_len_kv_cache = options.seq_len_kv_cache; + shape.head_size_qk = options.head_size_qk; + shape.head_size_vo = options.head_size_vo; + } + + auto [batch, num_heads_q, num_heads_kv, seq_len_qo, seq_len_kv, seq_len_kv_cache, head_size_qk, head_size_vo] = + problem_size; + auto shape_Q = cute::make_shape(seq_len_qo, head_size_qk, num_heads_q, batch); + auto shape_K = cute::make_shape(seq_len_kv, head_size_qk, num_heads_kv, batch); + auto shape_V = cute::make_shape(head_size_vo, seq_len_kv, num_heads_kv, batch); + auto shape_K_cache = cute::make_shape(seq_len_kv_cache, head_size_qk, num_heads_kv, batch); + auto shape_V_cache = cute::make_shape(head_size_vo, seq_len_kv_cache, num_heads_kv, batch); + auto shape_O = cute::make_shape(seq_len_qo, head_size_vo, num_heads_q, batch); + + if constexpr (!isVarLen) { + if (options.use_tensor_strides) { + set_tensor_strides_from_options(options, stride_Q, stride_K, stride_V, stride_O); + } else { + stride_Q = cutlass::make_cute_packed_stride(StrideQ{}, shape_Q); + stride_K = cutlass::make_cute_packed_stride(StrideK{}, shape_K); + stride_V = cutlass::make_cute_packed_stride(StrideV{}, shape_V); + stride_O = cutlass::make_cute_packed_stride(StrideO{}, shape_O); + } + } else { + stride_Q = cutlass::make_cute_packed_stride(StrideQ{}, shape_Q); + stride_K = cutlass::make_cute_packed_stride(StrideK{}, shape_K); + stride_V = cutlass::make_cute_packed_stride(StrideV{}, shape_V); + stride_O = cutlass::make_cute_packed_stride(StrideO{}, shape_O); + } + stride_K_cache = cutlass::make_cute_packed_stride(StrideK{}, shape_K_cache); + stride_V_cache = cutlass::make_cute_packed_stride(StrideV{}, shape_V_cache); + + if constexpr (isVarLen) { + // shape.seq_len_qo.cumulative_length = device_cumulative_seqlen_q.get(); + // shape.seq_len_kv.cumulative_length = device_cumulative_seqlen_kv.get(); + // shape.seq_len_kv_cache.cumulative_length = device_cumulative_seqlen_kv_cache.get(); + } + return shape; + } + + cutlass::Status run(const Options& options, const cutlass::KernelHardwareInfo& hw_info) { + ProblemShapeType shape = initialize(options); + + typename FMHAKernel::Arguments arguments{ + { + shape, + static_cast(options.q), + stride_Q, + static_cast(options.k), + stride_K, + static_cast(options.v), + stride_V, + static_cast(options.o), + stride_O, + static_cast(options.block_K), + stride_K_cache, + static_cast(options.block_V), + stride_V_cache, + }, + make_sage_mainloop_arguments(options), + {}, + hw_info}; + // Define device-global scratch memory + size_t workspace_size = FMHAKernel::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + if (!FMHAKernel::can_implement(arguments)) { + std::cout << "Invalid Problem Size: " << options.batch << 'x' << options.num_heads_q << 'x' << options.seq_len_qo + << 'x' << options.seq_len_kv << 'x' << options.head_size_qk << 'x' << options.head_size_vo + << (options.is_causal ? "xCausal" : "xNonCausal") << std::endl; + return cutlass::Status::kErrorInvalidProblem; + } + + // Initialize the workspace + CUTLASS_CHECK(FMHAKernel::initialize_workspace(arguments, workspace.get())); + + // Convert host-side arguments to device-side arguments to be passed to the kernel + auto params = FMHAKernel::to_underlying_arguments(arguments, workspace.get()); + + run(params); + + return cutlass::Status::kSuccess; + } +}; + +#if defined(ARK_SDPA_ENABLE_DENSE) +template default */ + int PipelineStages, bool persistent, typename ElementQ = bfloat16_t, typename ElementK = bfloat16_t, + typename ElementV = bfloat16_t, typename ElementO = ElementQ, + typename MMAOperation_ = void, /* void -> default */ + typename StrideQ = Stride, typename StrideK = Stride, + typename StrideV = Stride<_1, int, int, int>, typename StrideO = Stride, + typename GmemTiledCopyQ = void, /* void -> default block 2D */ + typename GmemTiledCopyK = void, typename GmemTiledCopyV = void, typename GmemTiledCopyO = void> +struct SageConfig { + static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + using MMAOperation = + cute::conditional_t, XE_DPAS_TT, MMAOperation_>; + // The PV "float" tiled MMA operates on the output element type. For UseInt8PV the kernel also + // constructs a separate int8 quantized PV MMA internally, while this float MMA only needs to + // describe the tile shape used for the accumulator and dequantized path. For the non-int8-PV + // path, V is consumed directly so we use ElementO (== ElementV) which supports half_t / bfloat16_t. + using MMAOperationPV = cute::conditional_t, + XE_DPAS_TT, MMAOperation_>; + using SubgroupLayoutPV = + cute::conditional_t, + decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; + + template + static int run(const Options& options) { + // + // Run examples + // + + // The KernelHardwareInfo struct holds the number of EUs on the GPU with a given device ID. This + // information is used by the underlying kernel. + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + + using ProblemShapeType = cutlass::fmha::kernel::SparseSageProblemShape; + + using TiledMMAQK = typename TiledMMAHelper, Layout, SubgroupLayoutQK>::TiledMMA; + using TiledMMAPV = + typename TiledMMAHelper, Layout, SubgroupLayoutPV>::TiledMMA; + + static_assert(get<0>(TileShapeOutput{}) == get<0>(TileShapePV{}), + "Output tile and P*V tile have different sizes in Q dimension"); + constexpr int VTiles = get<1>(TileShapeOutput{}) / get<1>(TileShapePV{}); + + auto make_dummy_tensor = [&](auto val, auto stride) { + return make_tensor(make_gmem_ptr(&val), make_layout(repeat>(1), stride)); + }; + + using TensorQ = decltype(make_dummy_tensor(ElementQ{}, StrideQ{})); + using TensorK = decltype(make_dummy_tensor(ElementK{}, StrideK{})); + using TensorV = decltype(make_dummy_tensor(ElementV{}, StrideV{})); + using TensorO = decltype(make_dummy_tensor(ElementO{}, StrideO{})); + using TensorK_cache = TensorK; + using TensorV_cache = TensorV; + using GmemTiledCopyK_cache = GmemTiledCopyK; + using GmemTiledCopyV_cache = GmemTiledCopyV; + + // Mainloop + using MainloopDispatchPolicy = cutlass::sage::XeDefault; + if constexpr (Causal) { + using CollectiveMainloop = + cutlass::fmha::collective::SAGEV1FwdMainloop; + + // Epilogue + using CollectiveEpilogue = + cutlass::fmha::collective::FMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = + cutlass::fmha::kernel::XeSageFwdKernel; + + SageKernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + if (options.mask) { + using CollectiveMainloop = + cutlass::fmha::collective::SAGEV1FwdMainloop; + + // Epilogue + using CollectiveEpilogue = + cutlass::fmha::collective::FMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = cutlass::fmha::kernel::XeSageFwdKernel; + + SageKernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + using CollectiveMainloop = + cutlass::fmha::collective::SAGEV1FwdMainloop; + + // Epilogue + using CollectiveEpilogue = + cutlass::fmha::collective::FMHAFwdEpilogue; + + static_assert(!(persistent & Causal), "persistent SDPA kernel not support Causal yet"); + using FMHAKernel = cutlass::fmha::kernel::XeSageFwdKernel; + + SageKernelRunner runner; + + CUTLASS_CHECK(runner.run(options, hw_info)); + } + } + return 0; + } + + static int run(const Options& options) { + bool cached_kv = options.seq_len_kv_cache > 0; + if constexpr (persistent) { + if (options.use_paged_kv || options.seq_len_kv_cache > 0) { + std::cerr + << "Error: Persistent kernel does not support paged/cached KV cache (use_paged_kv or seq_len_kv_cache > 0)." + << std::endl; + return -1; + } + return run(options); + } else if (options.use_paged_kv && !options.varlen) { + throw std::runtime_error("Paged KV without varlen is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && options.varlen && !cached_kv) { + throw std::runtime_error("Varlen without cached KV is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && !options.varlen && !cached_kv) { + return run(options); + } else if (!options.use_paged_kv && options.varlen && cached_kv) { + throw std::runtime_error("Varlen with cached KV but without paged KV is not supported yet"); + // return run(options); + } else if (!options.use_paged_kv && !options.varlen && cached_kv) { + throw std::runtime_error("Cached KV without varlen is not supported yet"); + // return run(options); + } else { + throw std::runtime_error("The combination of options is not supported yet"); + // return run(options); + } + } +}; +#endif + +#if defined(ARK_SDPA_ENABLE_SPARSE) +template , + typename StrideK = Stride, typename StrideV = Stride<_1, int, int, int>, + typename StrideO = Stride, typename GmemTiledCopyQ = void, + typename GmemTiledCopyK = void, typename GmemTiledCopyV = void, typename GmemTiledCopyO = void> +struct SparseSageConfig { + static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + using MMAOperation = + cute::conditional_t, XE_DPAS_TT, MMAOperation_>; + using MMAOperationPV = cute::conditional_t, + XE_DPAS_TT, MMAOperation_>; + using SubgroupLayoutPV = + cute::conditional_t, + decltype(cutlass::fmha::collective::get_sg_layout_pv(SubgroupLayoutQK{})), SubgroupLayoutPV_>; + + template + static int run(const Options& options) { + cutlass::KernelHardwareInfo hw_info; + hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id); + using ProblemShapeType = cutlass::fmha::kernel::SparseSageProblemShape; + using TiledMMAQK = typename TiledMMAHelper, Layout, SubgroupLayoutQK>::TiledMMA; + using TiledMMAPV = + typename TiledMMAHelper, Layout, SubgroupLayoutPV>::TiledMMA; + static_assert(get<0>(TileShapeOutput{}) == get<0>(TileShapePV{}), + "Output tile and P*V tile have different sizes in Q dimension"); + constexpr int VTiles = get<1>(TileShapeOutput{}) / get<1>(TileShapePV{}); + auto make_dummy_tensor = [&](auto val, auto stride) { + return make_tensor(make_gmem_ptr(&val), make_layout(repeat>(1), stride)); + }; + using TensorQ = decltype(make_dummy_tensor(ElementQ{}, StrideQ{})); + using TensorK = decltype(make_dummy_tensor(ElementK{}, StrideK{})); + using TensorV = decltype(make_dummy_tensor(ElementV{}, StrideV{})); + using TensorO = decltype(make_dummy_tensor(ElementO{}, StrideO{})); + using TensorK_cache = TensorK; + using TensorV_cache = TensorV; + using GmemTiledCopyK_cache = GmemTiledCopyK; + using GmemTiledCopyV_cache = GmemTiledCopyV; + using MainloopDispatchPolicy = cutlass::sage::XeDefault; + if constexpr (Causal) { + using CollectiveMainloop = + cutlass::fmha::collective::SPARSESAGEV1FwdMainloop; + using CollectiveEpilogue = + cutlass::fmha::collective::SparseFMHAFwdEpilogue; + using FMHAKernel = cutlass::fmha::kernel::XeSparseSageFwdKernel; + SageKernelRunner runner; + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + if (options.mask) { + using CollectiveMainloop = + cutlass::fmha::collective::SPARSESAGEV1FwdMainloop; + using CollectiveEpilogue = + cutlass::fmha::collective::SparseFMHAFwdEpilogue; + using FMHAKernel = cutlass::fmha::kernel::XeSparseSageFwdKernel; + SageKernelRunner runner; + CUTLASS_CHECK(runner.run(options, hw_info)); + } else { + using CollectiveMainloop = + cutlass::fmha::collective::SPARSESAGEV1FwdMainloop; + using CollectiveEpilogue = + cutlass::fmha::collective::SparseFMHAFwdEpilogue; + using FMHAKernel = cutlass::fmha::kernel::XeSparseSageFwdKernel; + SageKernelRunner runner; + CUTLASS_CHECK(runner.run(options, hw_info)); + } + } + return 0; + } + + static int run(const Options& options) { + if (options.use_paged_kv || options.varlen) { + throw std::runtime_error("Sparse Sage does not support paged KV or varlen in v1"); + } + if (options.block_K != nullptr || options.block_V != nullptr) { + if (options.seq_len_kv_cache <= 0 || options.seq_len_qo != 1) { + throw std::runtime_error("Sparse Sage only supports block_K/block_V for seq_len_q == 1 cached decode in v1"); + } + } else if (options.seq_len_kv_cache > 0) { + throw std::runtime_error("Sparse Sage cached decode requires block_K/block_V cache tensors"); + } + if (options.lut == nullptr || options.valid_block_num == nullptr) { + throw std::runtime_error("Sparse Sage requires lut and valid_block_num"); + } + if (options.seq_len_kv_cache > 0) { + return run(options); + } + return run(options); + } +}; +#endif // ARK_SDPA_ENABLE_SPARSE + +// ======================================================================== +// Prefill Kernel Launch +// ======================================================================== + +#if defined(ARK_SDPA_ENABLE_DENSE) +template +inline int launch_prefill_kernel_128(Options const& options) { + constexpr int PipelineStages = 2; + constexpr int PipelineStages1 = 2; + using ShapeQK = Shape<_256, _32, _32>; + using ShapePV = Shape<_256, _32, _32>; + using ShapeOut = Shape<_256, _128>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_256, _32, _32>; + using ShapePV1 = Shape<_256, _32, _32>; + using ShapeOut1 = Shape<_256, _128>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_sage_prefill_kernel_128(Options const& options) { + constexpr int PipelineStages = 2; + constexpr int PipelineStages1 = 2; + using ShapeQK = Shape<_256, _64, _32>; + using ShapePV = Shape<_256, _32, _64>; + using ShapeOut = Shape<_256, _128>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_256, _64, _32>; + using ShapePV1 = Shape<_256, _32, _64>; + using ShapeOut1 = Shape<_256, _128>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? SageConfig::run(options) + : SageConfig::run(options); +} + +template +inline int launch_sage_prefill_kernel_64(Options const& options) { + constexpr int PipelineStages = 2; + constexpr int PipelineStages1 = 2; + using ShapeQK = Shape<_128, _64, _32>; + using ShapePV = Shape<_128, _32, _64>; + using ShapeOut = Shape<_128, _64>; + using SubgroupLayoutQK = Layout>; + using SubgroupLayoutPV = void; + return options.is_causal + ? SageConfig::run(options) + : SageConfig::run(options); +} +#endif // ARK_SDPA_ENABLE_DENSE + +#if defined(ARK_SDPA_ENABLE_SPARSE) +template +inline int launch_sparse_sage_prefill_kernel_128(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_256, _64, _32>; + using ShapePV = Shape<_256, _32, _64>; + using ShapeOut = Shape<_256, _128>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_256, _64, _32>; + using ShapePV1 = Shape<_256, _32, _64>; + using ShapeOut1 = Shape<_256, _128>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? SparseSageConfig::run(options) + : SparseSageConfig::run(options); +} + +template +inline int launch_sparse_sage_prefill_kernel_128_qtile128(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_128, _64, _32>; + using ShapePV = Shape<_128, _32, _64>; + using ShapeOut = Shape<_128, _128>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_128, _64, _32>; + using ShapePV1 = Shape<_128, _32, _64>; + using ShapeOut1 = Shape<_128, _128>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? SparseSageConfig::run(options) + : SparseSageConfig::run(options); +} + +template +inline int launch_sparse_sage_prefill_kernel_128_qtile64(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_64, _64, _32>; + using ShapePV = Shape<_64, _32, _64>; + using ShapeOut = Shape<_64, _128>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_64, _64, _32>; + using ShapePV1 = Shape<_64, _32, _64>; + using ShapeOut1 = Shape<_64, _128>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? SparseSageConfig::run(options) + : SparseSageConfig::run(options); +} + +template +inline int launch_sparse_sage_prefill_kernel_64(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_128, _64, _32>; + using ShapePV = Shape<_128, _32, _64>; + using ShapeOut = Shape<_128, _64>; + using SubgroupLayoutQK = Layout>; + using SubgroupLayoutPV = void; + return options.is_causal + ? SparseSageConfig::run(options) + : SparseSageConfig::run(options); +} +#endif // ARK_SDPA_ENABLE_SPARSE + +#if defined(ARK_SDPA_ENABLE_DENSE) +template +inline int launch_prefill_kernel_64(Options const& options) { + constexpr int PipelineStages = 1; + using ShapeQK = Shape<_256, _64, _32>; + using ShapePV = Shape<_256, _32, _64>; + using ShapeOut = Shape<_256, _64>; + using SubgroupLayoutQK = Layout>; + using ShapeQK1 = Shape<_128, _64, _32>; + using ShapePV1 = Shape<_128, _32, _64>; + using ShapeOut1 = Shape<_128, _64>; + using SubgroupLayoutQK1 = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_prefill_kernel_192(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_256, _64, _32>; + using ShapePV = Shape<_256, _32, _64>; + using ShapeOut = Shape<_256, _192>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_prefill_kernel_96(Options const& options) { + constexpr int PipelineStages = 2; + using ShapeQK = Shape<_128, _64, _32>; + using ShapePV = Shape<_128, _32, _64>; + using ShapeOut = Shape<_128, _96>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_decode_kernel_128(Options const& options) { + constexpr int PipelineStages = 1; + using NUM_SG = std::conditional_t; + using KV_TILE_SIZE = std::conditional_t; + using ShapeQK = Shape<_1, KV_TILE_SIZE, _64>; + using ShapePV = Shape<_1, _32, KV_TILE_SIZE>; + using ShapeOut = Shape<_1, _128>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_decode_kernel_64(Options const& options) { + constexpr int PipelineStages = 1; + using NUM_SG = std::conditional_t; + using KV_TILE_SIZE = std::conditional_t; + using ShapeQK = Shape<_1, KV_TILE_SIZE, _64>; + using ShapePV = Shape<_1, _32, KV_TILE_SIZE>; + using ShapeOut = Shape<_1, _64>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_decode_kernel_96(Options const& options) { + constexpr int PipelineStages = 1; + using NUM_SG = std::conditional_t; + using KV_TILE_SIZE = std::conditional_t; + using ShapeQK = Shape<_1, KV_TILE_SIZE, _64>; + using ShapePV = Shape<_1, _32, KV_TILE_SIZE>; + using ShapeOut = Shape<_1, _96>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} + +template +inline int launch_decode_kernel_192(Options const& options) { + constexpr int PipelineStages = 1; + using NUM_SG = std::conditional_t; + using KV_TILE_SIZE = std::conditional_t; + using ShapeQK = Shape<_1, KV_TILE_SIZE, _64>; + using ShapePV = Shape<_1, _32, KV_TILE_SIZE>; + using ShapeOut = Shape<_1, _192>; + using SubgroupLayoutQK = Layout>; + return options.is_causal ? FMHAConfig::run(options) + : FMHAConfig::run(options); +} +#endif // ARK_SDPA_ENABLE_DENSE + +} // namespace sparse_detail +#endif // ARK_SYCL_TLA + +} // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sage_sparse_prefill_e2e.py b/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sage_sparse_prefill_e2e.py new file mode 100644 index 000000000..1c3cbcc76 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sage_sparse_prefill_e2e.py @@ -0,0 +1,297 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import math +import sys +from pathlib import Path + +import torch + +REPO_PARENT = Path(__file__).resolve().parents[3] +if str(REPO_PARENT) not in sys.path: + sys.path.insert(0, str(REPO_PARENT)) + +import auto_round_kernel as ark + + +def ensure_sparse_binding() -> None: + if getattr(ark, "xpu_lib", None) is not None and hasattr(ark.xpu_lib, "sage_sparse"): + return + candidates = sorted((REPO_PARENT / "auto_round_kernel" / "xbuild").glob("auto_round_kernel_xpu*.so")) + if not candidates: + raise RuntimeError("Unable to locate built XPU extension with sage_sparse in xbuild/") + ext_path = candidates[-1] + spec = importlib.util.spec_from_file_location("auto_round_kernel_xpu", ext_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load extension spec from {ext_path}") + module = importlib.util.module_from_spec(spec) + sys.modules["auto_round_kernel_xpu"] = module + spec.loader.exec_module(module) + if not hasattr(module, "sage_sparse"): + raise RuntimeError(f"Loaded extension does not expose sage_sparse: {ext_path}") + ark.xpu_lib = module + + +def quantize_qk(tensor: torch.Tensor, block_size: int) -> tuple[torch.Tensor, torch.Tensor]: + batch, heads, seq_len, head_dim = tensor.shape + num_rows = batch * heads * seq_len + num_blocks = num_rows // block_size + q_i8 = torch.empty_like(tensor, dtype=torch.int8) + scale = torch.empty(num_blocks, dtype=torch.float32, device=tensor.device) + lib = ark.get_lib(tensor) + stream = ark.get_stream(tensor) + lib.sage_dynamic_quant( + stream, + tensor.data_ptr(), + 0, + q_i8.data_ptr(), + scale.data_ptr(), + num_rows, + head_dim, + block_size, + ) + return q_i8, scale.reshape(batch, heads, seq_len // block_size, 1) + + +def build_sparse_metadata_and_mask( + batch: int, + heads: int, + seq_len: int, + block_size: int, + query_tile_tokens: int, + per_query_tile_selection: list[list[int]], + device: torch.device, + is_causal: bool = False, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_blocks = seq_len // block_size + kv_blocks = q_blocks + active_query_tiles = (seq_len + query_tile_tokens - 1) // query_tile_tokens + assert len(per_query_tile_selection) == active_query_tiles + + lut = torch.zeros((batch, heads, q_blocks, kv_blocks), dtype=torch.int32, device=device) + valid = torch.zeros((batch, heads, q_blocks), dtype=torch.int32, device=device) + mask = torch.full((batch, 1, seq_len, seq_len), -1.0e9, dtype=torch.float32, device=device) + + q_blocks_per_query_tile = max(1, query_tile_tokens // block_size) + for qblk in range(q_blocks): + qtile = min(qblk // q_blocks_per_query_tile, active_query_tiles - 1) + selected_blocks = per_query_tile_selection[qtile] + previous = 0 + for i, selected in enumerate(selected_blocks): + lut[..., qblk, i] = selected if i == 0 else (selected - previous) + previous = selected + valid[..., qblk] = len(selected_blocks) + + for qtile, selected_blocks in enumerate(per_query_tile_selection): + q_start = qtile * query_tile_tokens + q_end = min(q_start + query_tile_tokens, seq_len) + for qt in range(q_start, q_end): + for selected in selected_blocks: + k_start = selected * block_size + k_end = min(k_start + block_size, seq_len) + if not is_causal: + mask[:, :, qt : qt + 1, k_start:k_end] = 0.0 + else: + visible_end = min(k_end, qt + 1) + if visible_end > k_start: + mask[:, :, qt : qt + 1, k_start:visible_end] = 0.0 + + return lut.contiguous(), valid.contiguous(), mask.contiguous() + + +def run_case(head_dim: int, block_size: int = 64, is_causal: bool = False) -> None: + device = torch.device("xpu") + batch = 1 + heads = 4 + seq_len = 256 + scale = 1.0 / math.sqrt(head_dim) + query_tile_tokens = 128 if head_dim == 64 else 256 + kv_blocks = seq_len // block_size + if not is_causal: + per_query_tile_selection = [[0, 1], [1, 3]] if head_dim == 64 else [[0, 2]] + case_name = "python_prefill" + else: + per_query_tile_selection = [[0, 1, 2], [1, 3]] if head_dim == 64 else [[0, 2, 3]] + case_name = "python_prefill_causal" + + torch.manual_seed(2026 + head_dim) + query = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + key = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + value = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + + q_i8, q_scale = quantize_qk(query, block_size) + k_i8, k_scale = quantize_qk(key, block_size) + lut, valid, dense_mask = build_sparse_metadata_and_mask( + batch, heads, seq_len, block_size, query_tile_tokens, per_query_tile_selection, device, is_causal=is_causal + ) + + dense_out = ark.sage( + q_i8, + k_i8, + value, + attn_mask=dense_mask, + is_causal=False, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + sparse_out = ark.sage_sparse( + q_i8, + k_i8, + value, + lut, + valid, + is_causal=is_causal, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + torch.xpu.synchronize() + + diff = (dense_out.float() - sparse_out.float()).abs() + max_diff = float(diff.max().cpu()) + mean_diff = float(diff.mean().cpu()) + print(f"[sage_sparse][{case_name}] D={head_dim} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + if max_diff > 5e-3 or mean_diff > 5e-4: + raise RuntimeError(f"sage_sparse python prefill mismatch for D={head_dim}, causal={is_causal}") + + +def run_multi_row_tile_case() -> None: + device = torch.device("xpu") + batch = 1 + heads = 4 + head_dim = 128 + seq_len = 256 + block_size = 64 + scale = 1.0 / math.sqrt(head_dim) + query_tile_tokens = 64 + per_query_tile_selection = [ + [0, 1], + [1, 3], + [0, 2, 3], + [2, 3], + ] + + torch.manual_seed(4026) + query = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + key = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + value = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + + q_i8, q_scale = quantize_qk(query, block_size) + k_i8, k_scale = quantize_qk(key, block_size) + lut, valid, dense_mask = build_sparse_metadata_and_mask( + batch, heads, seq_len, block_size, query_tile_tokens, per_query_tile_selection, device, is_causal=False + ) + + dense_out = ark.sage( + q_i8, + k_i8, + value, + attn_mask=dense_mask, + is_causal=False, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + sparse_out = ark.sage_sparse( + q_i8, + k_i8, + value, + lut, + valid, + is_causal=False, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + torch.xpu.synchronize() + + diff = (dense_out.float() - sparse_out.float()).abs() + max_diff = float(diff.max().cpu()) + mean_diff = float(diff.mean().cpu()) + print(f"[sage_sparse][python_prefill_multi_row_tile] D=128 max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + if max_diff > 5e-3 or mean_diff > 5e-4: + raise RuntimeError("sage_sparse multi-row tile mismatch for D=128") + + +def run_all_selected_case(head_dim: int, block_size: int = 64) -> None: + device = torch.device("xpu") + batch = 1 + heads = 4 + seq_len = 256 + scale = 1.0 / math.sqrt(head_dim) + query_tile_tokens = 128 if head_dim == 64 else 256 + kv_blocks = seq_len // block_size + active_query_tiles = (seq_len + query_tile_tokens - 1) // query_tile_tokens + per_query_tile_selection = [list(range(kv_blocks)) for _ in range(active_query_tiles)] + + torch.manual_seed(3026 + head_dim) + query = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + key = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + value = torch.randn((batch, heads, seq_len, head_dim), dtype=torch.float16, device=device) + + q_i8, q_scale = quantize_qk(query, block_size) + k_i8, k_scale = quantize_qk(key, block_size) + lut, valid, dense_mask = build_sparse_metadata_and_mask( + batch, heads, seq_len, block_size, query_tile_tokens, per_query_tile_selection, device, is_causal=False + ) + + dense_out = ark.sage( + q_i8, + k_i8, + value, + attn_mask=dense_mask, + is_causal=False, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + sparse_out = ark.sage_sparse( + q_i8, + k_i8, + value, + lut, + valid, + is_causal=False, + scale=scale, + quant_block_size=block_size, + qscale=q_scale, + kscale=k_scale, + tensor_layout="HND", + ) + torch.xpu.synchronize() + + diff = (dense_out.float() - sparse_out.float()).abs() + max_diff = float(diff.max().cpu()) + mean_diff = float(diff.mean().cpu()) + print(f"[sage_sparse][python_prefill_all_selected] D={head_dim} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + if max_diff > 5e-3 or mean_diff > 5e-4: + raise RuntimeError(f"sage_sparse python all-selected mismatch for D={head_dim}") + + +def main() -> None: + ensure_sparse_binding() + if not torch.xpu.is_available(): + raise RuntimeError("XPU device is required") + run_all_selected_case(64) + run_all_selected_case(128) + run_case(64) + run_case(128) + run_multi_row_tile_case() + run_case(64, is_causal=True) + run_case(128, is_causal=True) + + +if __name__ == "__main__": + main() diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sparge_preprocess_topk_e2e.py b/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sparge_preprocess_topk_e2e.py new file mode 100644 index 000000000..a6b1de13b --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/test/test_sparge_preprocess_topk_e2e.py @@ -0,0 +1,313 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import importlib.util +import math +import sys +from pathlib import Path + +import torch + +REPO_PARENT = Path(__file__).resolve().parents[3] +if str(REPO_PARENT) not in sys.path: + sys.path.insert(0, str(REPO_PARENT)) + +import auto_round_kernel as ark + + +def ensure_sparse_binding() -> None: + if getattr(ark, "xpu_lib", None) is not None and hasattr(ark.xpu_lib, "sage_sparse"): + return + candidates = sorted((REPO_PARENT / "auto_round_kernel" / "xbuild").glob("auto_round_kernel_xpu*.so")) + if not candidates: + raise RuntimeError("Unable to locate built XPU extension with sage_sparse in xbuild/") + ext_path = candidates[-1] + spec = importlib.util.spec_from_file_location("auto_round_kernel_xpu", ext_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load extension spec from {ext_path}") + module = importlib.util.module_from_spec(spec) + sys.modules["auto_round_kernel_xpu"] = module + spec.loader.exec_module(module) + if not hasattr(module, "sage_sparse"): + raise RuntimeError(f"Loaded extension does not expose sage_sparse: {ext_path}") + ark.xpu_lib = module + + +def _to_layout(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: + if tensor_layout == "HND": + return tensor.contiguous() + return tensor.permute(0, 2, 1, 3).contiguous() + + +def _assert_metadata_matches(actual: dict, reference: dict) -> None: + for key in ("query_i8", "key_i8", "lut", "valid_block_num", "block_map", "raw_block_map", "tile_block_map"): + assert torch.equal(actual[key], reference[key]), key + for key in ("sim_qblocks", "sim_kblocks"): + assert torch.equal(actual[key], reference[key]), key + for key in ("qscale", "kscale"): + assert torch.allclose(actual[key], reference[key], atol=2.0e-5, rtol=0.0), key + assert actual["query_tile_tokens"] == reference["query_tile_tokens"] + assert actual["quant_block_size"] == reference["quant_block_size"] + assert actual["sparse_q_block_tokens"] == reference["sparse_q_block_tokens"] + assert actual["sparse_k_block_tokens"] == reference["sparse_k_block_tokens"] + assert actual["kernel_compatibility_added_blocks"] == reference["kernel_compatibility_added_blocks"] + assert actual["stats"] == reference["stats"] + + +def run_case( + head_dim: int, + *, + topk: float, + is_causal: bool, + tensor_layout: str, + query_tile_tokens: int | None = None, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, + seq_len_q: int = 256, + seq_len_kv: int = 256, + num_heads_q: int = 1, + num_heads_kv: int = 1, +) -> None: + device = torch.device("xpu") + batch = 1 + scale = 1.0 / math.sqrt(head_dim) + effective_query_tile_tokens = query_tile_tokens or (q_tile_override or None) + effective_q_tile_override = q_tile_override if q_tile_override != 0 else (query_tile_tokens or 0) + + seed = ( + 5200 + + head_dim + + int(topk * 100) + + (7 if is_causal else 0) + + (13 if tensor_layout == "NHD" else 0) + + seq_len_q + + (seq_len_kv * 3) + + (num_heads_q * 5) + + (num_heads_kv * 11) + ) + torch.manual_seed(seed) + query_hnd = torch.randn((batch, num_heads_q, seq_len_q, head_dim), dtype=torch.float16, device=device) + key_hnd = torch.randn((batch, num_heads_kv, seq_len_kv, head_dim), dtype=torch.float16, device=device) + value_hnd = torch.randn((batch, num_heads_kv, seq_len_kv, head_dim), dtype=torch.float16, device=device) + + query = _to_layout(query_hnd, tensor_layout) + key = _to_layout(key_hnd, tensor_layout) + value = _to_layout(value_hnd, tensor_layout) + + preprocess_meta = ark.sparge_preprocess_topk( + query, + key, + is_causal=is_causal, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=tensor_layout, + query_tile_tokens=effective_query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + torch_meta = ark._sparge_preprocess_topk_torch( + query, + key, + is_causal=is_causal, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=tensor_layout, + query_tile_tokens=effective_query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + _assert_metadata_matches(preprocess_meta, torch_meta) + + sparse_out, meta = ark.sparge_sage2_attn_meansim_topk_xpu( + query, + key, + value, + is_causal=is_causal, + scale=scale, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=tensor_layout, + query_tile_tokens=effective_query_tile_tokens, + q_tile_override=q_tile_override, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + return_metadata=True, + ) + + assert meta["backend"] in {"torch", "triton_xpu"} + _assert_metadata_matches(meta, preprocess_meta) + assert tuple(meta["query_i8"].shape) == tuple(query.shape) + assert tuple(meta["key_i8"].shape) == tuple(key.shape) + assert meta["valid_block_num"].dtype == torch.int32 + assert meta["lut"].dtype == torch.int32 + + direct_sparse = ark.sage_sparse( + meta["query_i8"], + meta["key_i8"], + value, + meta["lut"], + meta["valid_block_num"], + is_causal=is_causal, + scale=scale, + quant_block_size=meta["quant_block_size"], + qscale=meta["qscale"], + kscale=meta["kscale"], + q_tile_override=effective_q_tile_override, + sparse_q_block_tokens=meta["sparse_q_block_tokens"], + sparse_k_block_tokens=meta["sparse_k_block_tokens"], + tensor_layout=tensor_layout, + ) + torch_sparse = ark.sage_sparse( + torch_meta["query_i8"], + torch_meta["key_i8"], + value, + torch_meta["lut"], + torch_meta["valid_block_num"], + is_causal=is_causal, + scale=scale, + quant_block_size=torch_meta["quant_block_size"], + qscale=torch_meta["qscale"], + kscale=torch_meta["kscale"], + q_tile_override=effective_q_tile_override, + sparse_q_block_tokens=torch_meta["sparse_q_block_tokens"], + sparse_k_block_tokens=torch_meta["sparse_k_block_tokens"], + tensor_layout=tensor_layout, + ) + torch.xpu.synchronize() + + case_name = ( + f"topk_{topk:.2f}_{'causal' if is_causal else 'noncausal'}_{tensor_layout.lower()}" + f"_qtile{query_tile_tokens or 'default'}_kqtile{effective_q_tile_override or 'default'}" + ) + direct_diff = (direct_sparse.float() - sparse_out.float()).abs() + direct_max_diff = float(direct_diff.max().cpu()) + direct_mean_diff = float(direct_diff.mean().cpu()) + print( + f"[sparge_preprocess][{case_name}] D={head_dim} " + f"wrapper_max_diff={direct_max_diff:.6f} wrapper_mean_diff={direct_mean_diff:.6f} backend={meta['backend']}" + ) + if direct_max_diff > 5e-3 or direct_mean_diff > 5e-4: + raise RuntimeError(f"sparge preprocess wrapper mismatch for {case_name}, D={head_dim}") + + torch_diff = (torch_sparse.float() - sparse_out.float()).abs() + torch_max_diff = float(torch_diff.max().cpu()) + torch_mean_diff = float(torch_diff.mean().cpu()) + print( + f"[sparge_preprocess][{case_name}] D={head_dim} " + f"torch_replay_max_diff={torch_max_diff:.6f} torch_replay_mean_diff={torch_mean_diff:.6f}" + ) + if torch_max_diff > 5e-3 or torch_mean_diff > 5e-4: + raise RuntimeError(f"sparge preprocess torch replay mismatch for {case_name}, D={head_dim}") + + if topk == 1.0: + dense_mask = ark.sparge_block_map_to_mask( + meta["block_map"], + quant_block_size=meta["quant_block_size"], + q_block_tokens=meta["sparse_q_block_tokens"], + k_block_tokens=meta["sparse_k_block_tokens"], + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + is_causal=is_causal, + ) + dense_out = ark.sage( + meta["query_i8"], + meta["key_i8"], + value, + attn_mask=dense_mask, + is_causal=False, + scale=scale, + quant_block_size=meta["quant_block_size"], + qscale=meta["qscale"], + kscale=meta["kscale"], + tensor_layout=tensor_layout, + ) + diff = (dense_out.float() - sparse_out.float()).abs() + max_diff = float(diff.max().cpu()) + mean_diff = float(diff.mean().cpu()) + print( + f"[sparge_preprocess][{case_name}] D={head_dim} dense_max_diff={max_diff:.6f} dense_mean_diff={mean_diff:.6f}" + ) + if max_diff > 5e-3 or mean_diff > 5e-4: + raise RuntimeError(f"sparge preprocess dense-reference mismatch for {case_name}, D={head_dim}") + + if topk < 1.0: + assert 0.0 <= meta["stats"]["selected_ratio"] <= 1.0 + if seq_len_q == seq_len_kv: + assert meta["stats"]["selected_ratio"] < 1.0 + assert meta["kernel_compatibility_added_blocks"] == 0 + assert torch.equal(meta["block_map"], meta["raw_block_map"]) + assert torch.isfinite(sparse_out).all() + + if topk < 1.0 and not is_causal: + compatibility_seed = seed + 101 + torch.manual_seed(compatibility_seed) + crafted_query_hnd = torch.randn((batch, num_heads_q, seq_len_q, head_dim), dtype=torch.float16, device=device) + crafted_key_hnd = torch.randn((batch, num_heads_kv, seq_len_kv, head_dim), dtype=torch.float16, device=device) + crafted_query = _to_layout(crafted_query_hnd, tensor_layout) + crafted_key = _to_layout(crafted_key_hnd, tensor_layout) + crafted_meta = ark._sparge_preprocess_topk_torch( + crafted_query, + crafted_key, + is_causal=False, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=tensor_layout, + query_tile_tokens=effective_query_tile_tokens, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + assert crafted_meta["kernel_compatibility_added_blocks"] == 0 + assert torch.equal(crafted_meta["block_map"], crafted_meta["raw_block_map"]) + + +def main() -> None: + ensure_sparse_binding() + if not torch.xpu.is_available(): + raise RuntimeError("XPU device is required") + for tensor_layout in ("HND", "NHD"): + run_case(64, topk=1.0, is_causal=False, tensor_layout=tensor_layout) + run_case(128, topk=1.0, is_causal=False, tensor_layout=tensor_layout) + run_case(64, topk=0.5, is_causal=False, tensor_layout=tensor_layout) + run_case(128, topk=0.5, is_causal=False, tensor_layout=tensor_layout) + run_case(64, topk=0.5, is_causal=True, tensor_layout=tensor_layout) + run_case(128, topk=0.5, is_causal=True, tensor_layout=tensor_layout) + run_case(64, topk=0.25, is_causal=False, tensor_layout=tensor_layout) + run_case(128, topk=0.25, is_causal=False, tensor_layout=tensor_layout) + run_case(128, topk=0.5, is_causal=False, tensor_layout=tensor_layout, query_tile_tokens=256) + run_case(128, topk=0.5, is_causal=False, tensor_layout=tensor_layout, q_tile_override=256) + run_case( + 128, + topk=0.5, + is_causal=False, + tensor_layout=tensor_layout, + q_tile_override=256, + sparse_q_block_tokens=256, + sparse_k_block_tokens=64, + seq_len_q=512, + seq_len_kv=512, + num_heads_q=2, + num_heads_kv=2, + ) + run_case( + 128, + topk=0.5, + is_causal=False, + tensor_layout=tensor_layout, + seq_len_q=256, + seq_len_kv=128, + num_heads_q=2, + num_heads_kv=1, + ) + + +if __name__ == "__main__": + main() diff --git a/auto_round_extension/ark/examples/FLUX_SWEEP.md b/auto_round_extension/ark/examples/FLUX_SWEEP.md new file mode 100644 index 000000000..5e0de4d9c --- /dev/null +++ b/auto_round_extension/ark/examples/FLUX_SWEEP.md @@ -0,0 +1,20 @@ +# FLUX Sweep + +Use the helper script below to sweep one dense baseline plus sparse `topk=1.0..0.1` +with the default FLUX generation settings (`1024x1024`, `50` steps, guidance `3.5`) +and the `q_tile=256`, `sparse_q_block_tokens=256`, `sparse_k_block_tokens=64` kernel. + +```bash +cd auto_round_extension/ark/examples +FLUX_SWEEP_PYTHON=/home/yiliu4/workspace/auto-round-py/.venv/bin/python \ +FLUX_MODEL=/home/yiliu4/workspace/models/black-forest-labs/FLUX.1-dev \ +FLUX_SWEEP_DEVICES=0,1,2,3,4,5,6,7 \ +bash run_flux_sweep.sh +``` + +Results are written into `auto_round_extension/ark/examples/results/flux_sweep_defaultsteps_/` +with: + +- `summary.csv` for the full sweep summary +- `commands.txt` for the exact per-run commands +- one `.png` and one `.log` per config diff --git a/auto_round_extension/ark/examples/__init__.py b/auto_round_extension/ark/examples/__init__.py new file mode 100644 index 000000000..0b3bf2eae --- /dev/null +++ b/auto_round_extension/ark/examples/__init__.py @@ -0,0 +1,2 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 diff --git a/auto_round_extension/ark/examples/flux_sparse_patch.py b/auto_round_extension/ark/examples/flux_sparse_patch.py new file mode 100644 index 000000000..6695bcbb9 --- /dev/null +++ b/auto_round_extension/ark/examples/flux_sparse_patch.py @@ -0,0 +1,313 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import contextlib +import os +import sys +import time +import warnings +from dataclasses import dataclass +from pathlib import Path + +import torch + +ARK_DIR = Path(__file__).resolve().parent.parent +if str(ARK_DIR) not in sys.path: + sys.path.insert(0, str(ARK_DIR)) + +import auto_round_kernel as ark +from diffusers.models.transformers.transformer_flux import FluxAttention, _get_qkv_projections, apply_rotary_emb +from wan_sparse_patch import _parse_bool_env, _parse_optional_int_env, ensure_ark_sparse_binding + + +def _normalize_attention_mask( + attn_mask: torch.Tensor | None, + batch: int, + seq_q: int, + seq_kv: int, + device: torch.device, +) -> torch.Tensor | None: + if attn_mask is None: + return None + + mask = attn_mask + if mask.dtype == torch.bool: + zero = torch.zeros((), dtype=torch.float32, device=device) + neg_inf = torch.full((), -1.0e9, dtype=torch.float32, device=device) + mask = torch.where(mask, zero, neg_inf) + else: + mask = mask.to(torch.float32) + + if mask.ndim == 2: + mask = mask.view(1, 1, seq_q, seq_kv) + elif mask.ndim == 3: + mask = mask.unsqueeze(1) + elif mask.ndim != 4: + raise ValueError(f"Unsupported Flux attention mask rank: {mask.ndim}") + + if mask.shape[0] == 1 and batch != 1: + mask = mask.expand(batch, -1, -1, -1) + return mask.contiguous() + + +@dataclass +class FluxSparseAttentionStats: + total_calls: int = 0 + single_stream_calls: int = 0 + joint_stream_calls: int = 0 + sparse_calls: int = 0 + unsupported_fallbacks: int = 0 + runtime_fallbacks: int = 0 + sparse_sparsity_sum: float = 0.0 + timed_calls: int = 0 + processor_time_ms_sum: float = 0.0 + processor_time_ms_min: float = float("inf") + processor_time_ms_max: float = 0.0 + + @property + def avg_sparsity(self) -> float: + if self.sparse_calls == 0: + return 0.0 + return self.sparse_sparsity_sum / self.sparse_calls + + @property + def avg_processor_time_ms(self) -> float: + if self.timed_calls == 0: + return 0.0 + return self.processor_time_ms_sum / self.timed_calls + + +class FluxSparseAttnProcessor: + def __init__( + self, + original_processor, + stats: FluxSparseAttentionStats, + *, + smooth_k: bool, + topk: float, + attention_sink: bool, + debug_timing: bool, + q_tile_override: int, + sparse_q_block_tokens: int | None, + sparse_k_block_tokens: int | None, + ): + self.original_processor = original_processor + self.stats = stats + self.smooth_k = smooth_k + self.topk = topk + self.attention_sink = attention_sink + self.debug_timing = debug_timing + self.q_tile_override = q_tile_override + self.sparse_q_block_tokens = sparse_q_block_tokens + self.sparse_k_block_tokens = sparse_k_block_tokens + self._warned_runtime_fallback = False + + @staticmethod + def _synchronize_for_timing(device: torch.device) -> None: + if device.type == "xpu" and hasattr(torch, "xpu") and hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + + def __call__( + self, + attn: "FluxAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + self.stats.total_calls += 1 + if encoder_hidden_states is None: + self.stats.single_stream_calls += 1 + else: + self.stats.joint_stream_calls += 1 + + call_kind = "single" if encoder_hidden_states is None else "joint" + call_status = "sparse" + start_time = None + if self.debug_timing: + self._synchronize_for_timing(hidden_states.device) + start_time = time.perf_counter() + + try: + if hidden_states.device.type != "xpu" or hidden_states.dtype not in (torch.float16, torch.bfloat16): + self.stats.unsupported_fallbacks += 1 + call_status = "fallback_unsupported_device_or_dtype" + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + image_rotary_emb, + **kwargs, + ) + + if encoder_hidden_states is not None and attn.added_kv_proj_dim is None: + self.stats.unsupported_fallbacks += 1 + call_status = "fallback_unsupported_joint_layout" + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + image_rotary_emb, + **kwargs, + ) + + query, key, value, encoder_query, encoder_key, encoder_value = _get_qkv_projections( + attn, + hidden_states, + encoder_hidden_states, + ) + + query = query.unflatten(-1, (attn.heads, -1)).contiguous() + key = key.unflatten(-1, (attn.heads, -1)).contiguous() + value = value.unflatten(-1, (attn.heads, -1)).contiguous() + + query = attn.norm_q(query) + key = attn.norm_k(key) + + encoder_seq_len = 0 + if encoder_hidden_states is not None: + encoder_seq_len = encoder_hidden_states.shape[1] + encoder_query = encoder_query.unflatten(-1, (attn.heads, -1)).contiguous() + encoder_key = encoder_key.unflatten(-1, (attn.heads, -1)).contiguous() + encoder_value = encoder_value.unflatten(-1, (attn.heads, -1)).contiguous() + + encoder_query = attn.norm_added_q(encoder_query) + encoder_key = attn.norm_added_k(encoder_key) + + query = torch.cat([encoder_query, query], dim=1) + key = torch.cat([encoder_key, key], dim=1) + value = torch.cat([encoder_value, value], dim=1) + + if image_rotary_emb is not None: + query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) + key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + + mask = _normalize_attention_mask( + attention_mask, + batch=query.shape[0], + seq_q=query.shape[1], + seq_kv=key.shape[1], + device=query.device, + ) + hidden_states, sparsity = ark.sparge_sage2_attn_meansim_topk_xpu( + query, + key, + value, + attn_mask=mask, + is_causal=False, + smooth_k=self.smooth_k, + simthreshd1=-1.0, + topk=self.topk, + attention_sink=self.attention_sink, + tensor_layout="NHD", + q_tile_override=self.q_tile_override, + sparse_q_block_tokens=self.sparse_q_block_tokens, + sparse_k_block_tokens=self.sparse_k_block_tokens, + return_sparsity=True, + ) + self.stats.sparse_calls += 1 + self.stats.sparse_sparsity_sum += float(sparsity) + + hidden_states = hidden_states.flatten(2, 3).to(query.dtype) + + if encoder_hidden_states is not None: + encoder_hidden_states_out, hidden_states_out = hidden_states.split_with_sizes( + [encoder_seq_len, hidden_states.shape[1] - encoder_seq_len], + dim=1, + ) + hidden_states_out = attn.to_out[0](hidden_states_out.contiguous()) + hidden_states_out = attn.to_out[1](hidden_states_out) + encoder_hidden_states_out = attn.to_add_out(encoder_hidden_states_out.contiguous()) + return hidden_states_out, encoder_hidden_states_out + + return hidden_states + except Exception as exc: # noqa: BLE001 + self.stats.runtime_fallbacks += 1 + call_status = "fallback_runtime" + if not self._warned_runtime_fallback: + warnings.warn( + f"Flux sparse attention fell back to the original processor after a runtime error: {exc}", + stacklevel=2, + ) + self._warned_runtime_fallback = True + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + image_rotary_emb, + **kwargs, + ) + finally: + if self.debug_timing and start_time is not None: + self._synchronize_for_timing(hidden_states.device) + elapsed_ms = (time.perf_counter() - start_time) * 1000.0 + self.stats.timed_calls += 1 + self.stats.processor_time_ms_sum += elapsed_ms + self.stats.processor_time_ms_min = min(self.stats.processor_time_ms_min, elapsed_ms) + self.stats.processor_time_ms_max = max(self.stats.processor_time_ms_max, elapsed_ms) + print( + "[flux_sparse][timing]" + f" call={self.stats.timed_calls}" + f" kind={call_kind}" + f" status={call_status}" + f" hidden_shape={tuple(hidden_states.shape)}" + f" elapsed_ms={elapsed_ms:.3f}" + ) + + +@contextlib.contextmanager +def patch_flux_sparse_attention( + transformer, + *, + smooth_k: bool = True, + topk: float = 0.75, + attention_sink: bool = False, + debug_timing: bool = False, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +): + ensure_ark_sparse_binding() + originals: list[tuple[FluxAttention, object]] = [] + stats = FluxSparseAttentionStats() + + for module in transformer.modules(): + if isinstance(module, FluxAttention): + original = module.processor + module.set_processor( + FluxSparseAttnProcessor( + original, + stats, + smooth_k=smooth_k, + topk=topk, + attention_sink=attention_sink, + debug_timing=debug_timing, + q_tile_override=q_tile_override, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + ) + originals.append((module, original)) + + try: + yield stats + finally: + for module, original in originals: + module.set_processor(original) + + +def patch_flux_sparse_attention_from_env(transformer): + return patch_flux_sparse_attention( + transformer, + smooth_k=_parse_bool_env("FLUX_SPARSE_SMOOTH_K", True), + topk=float(os.getenv("FLUX_SPARSE_TOPK", "0.75")), + attention_sink=_parse_bool_env("FLUX_SPARSE_ATTENTION_SINK", False), + debug_timing=_parse_bool_env("FLUX_SPARSE_DEBUG_TIMING", False), + q_tile_override=int(os.getenv("FLUX_SPARSE_Q_TILE_OVERRIDE", "0")), + sparse_q_block_tokens=_parse_optional_int_env("FLUX_SPARSE_Q_BLOCK_TOKENS"), + sparse_k_block_tokens=_parse_optional_int_env("FLUX_SPARSE_K_BLOCK_TOKENS"), + ) diff --git a/auto_round_extension/ark/examples/run_flux.py b/auto_round_extension/ark/examples/run_flux.py new file mode 100644 index 000000000..660a0d9c4 --- /dev/null +++ b/auto_round_extension/ark/examples/run_flux.py @@ -0,0 +1,764 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import contextlib +import gzip +import json +import os +import statistics +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +import torch +from diffusers import FluxPipeline +from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps +from flux_sparse_patch import patch_flux_sparse_attention_from_env + +dtype = torch.bfloat16 +device = "xpu" + + +def env_flag(name, default="0"): + return os.getenv(name, default).lower() not in {"0", "false", "no", "off"} + + +benchmark_enabled = env_flag("FLUX_BENCHMARK_ENABLE", "0") +profiler_enabled = env_flag("FLUX_PROFILER_ENABLE", "0") +benchmark_scope = os.getenv("FLUX_BENCHMARK_SCOPE", "full").strip().lower() + +if benchmark_scope not in {"full", "denoising", "block"}: + raise ValueError("FLUX_BENCHMARK_SCOPE must be one of: full, denoising, block") + +benchmark_block_kind = os.getenv("FLUX_BENCHMARK_BLOCK_KIND", "joint").strip().lower() +if benchmark_block_kind not in {"joint", "single"}: + raise ValueError("FLUX_BENCHMARK_BLOCK_KIND must be one of: joint, single") + +benchmark_block_index = int(os.getenv("FLUX_BENCHMARK_BLOCK_INDEX", "0")) +benchmark_block_timestep_index = int(os.getenv("FLUX_BENCHMARK_BLOCK_TIMESTEP_INDEX", "0")) +cpu_offload_enabled = env_flag("FLUX_ENABLE_CPU_OFFLOAD", "0" if benchmark_scope == "block" else "1") + +model_id = os.getenv("FLUX_MODEL", "black-forest-labs/FLUX.1-dev") +pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=dtype) +if benchmark_scope != "block": + pipe.to(device) +if cpu_offload_enabled: + pipe.enable_model_cpu_offload() + +height = int(os.getenv("FLUX_HEIGHT", "1024")) +width = int(os.getenv("FLUX_WIDTH", "1024")) +num_inference_steps = int(os.getenv("FLUX_STEPS", "50")) +guidance_scale = float(os.getenv("FLUX_GUIDANCE_SCALE", "3.5")) +max_sequence_length = int(os.getenv("FLUX_MAX_SEQUENCE_LENGTH", "512")) +seed = int(os.getenv("FLUX_SEED", "0")) +use_sparse = os.getenv("FLUX_USE_SPARSE", "1").lower() not in {"0", "false", "no", "off"} + +output_file = ( + f"flux_output_{height}x{width}_{num_inference_steps}steps_" + f"{guidance_scale}gs_sparse{os.getenv('FLUX_SPARSE_TOPK', '0.5')}.png" +) +output_path = os.getenv("FLUX_OUTPUT", output_file) + +prompt = os.getenv("FLUX_PROMPT", "A cat holding a sign that says hello world") + + +def create_profiler(enabled, run_tag): + if not enabled: + return contextlib.nullcontext(), None, None, None + + xpu_activity = getattr(torch.profiler.ProfilerActivity, "XPU", None) + if xpu_activity is None: + raise RuntimeError("FLUX_PROFILER_ENABLE=1 requires torch.profiler.ProfilerActivity.XPU") + + profile_dir = Path(os.getenv("FLUX_PROFILER_DIR", "profiles")) + profile_dir.mkdir(parents=True, exist_ok=True) + trace_stem = os.getenv( + "FLUX_PROFILER_TRACE_NAME", + f"flux_{run_tag}_{height}x{width}_{num_inference_steps}steps_seed{seed}", + ) + trace_path = profile_dir / f"{trace_stem}.json.gz" + summary_path = profile_dir / f"{trace_stem}.txt" + sort_key = "self_xpu_time_total" + activities = [torch.profiler.ProfilerActivity.CPU, xpu_activity] + profiler = torch.profiler.profile( + activities=activities, + record_shapes=env_flag("FLUX_PROFILER_RECORD_SHAPES", "1"), + profile_memory=env_flag("FLUX_PROFILER_PROFILE_MEMORY", "0"), + with_stack=env_flag("FLUX_PROFILER_WITH_STACK", "1"), + ) + print(f"[flux_profile] capturing torch trace at {trace_path}") + return profiler, trace_path, summary_path, sort_key + + +def save_profiler_results(profiler, trace_path, summary_path, sort_key): + uncompressed_trace_path = trace_path.with_suffix("") + profiler.export_chrome_trace(str(uncompressed_trace_path)) + with open(uncompressed_trace_path, "rb") as src, gzip.open(trace_path, "wb") as dst: + dst.writelines(src) + uncompressed_trace_path.unlink() + print(f"[flux_profile] trace saved to {trace_path}") + if env_flag("FLUX_PROFILER_WRITE_SUMMARY", "0"): + try: + summary = profiler.key_averages().table( + sort_by=sort_key, + row_limit=int(os.getenv("FLUX_PROFILER_ROW_LIMIT", "40")), + ) + summary_path.write_text(summary) + print(f"[flux_profile] summary saved to {summary_path}") + except UnicodeDecodeError as exc: + warning = ( + "Skipped profiler summary because PyTorch failed to decode an XPU event " + f"name while parsing Kineto results: {exc}\n" + "The compressed Chrome trace was still saved successfully.\n" + ) + summary_path.write_text(warning) + print(f"[flux_profile] warning: {warning.strip()}") + + +def maybe_synchronize_xpu(): + if device == "xpu" and hasattr(torch, "xpu") and hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + + +def log_sparse_stats(sparse_stats): + print( + "[flux_sparse] stats:" + f" total_calls={sparse_stats.total_calls}" + f" single_stream_calls={sparse_stats.single_stream_calls}" + f" joint_stream_calls={sparse_stats.joint_stream_calls}" + f" sparse_calls={sparse_stats.sparse_calls}" + f" unsupported_fallbacks={sparse_stats.unsupported_fallbacks}" + f" runtime_fallbacks={sparse_stats.runtime_fallbacks}" + f" avg_sparsity={sparse_stats.avg_sparsity:.4f}" + ) + if sparse_stats.timed_calls: + print( + "[flux_sparse][timing] summary:" + f" timed_calls={sparse_stats.timed_calls}" + f" avg_ms={sparse_stats.avg_processor_time_ms:.3f}" + f" min_ms={sparse_stats.processor_time_ms_min:.3f}" + f" max_ms={sparse_stats.processor_time_ms_max:.3f}" + ) + + +def common_generation_kwargs(): + return dict( + height=height, + width=width, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + max_sequence_length=max_sequence_length, + generator=torch.Generator("cpu").manual_seed(seed), + ) + + +@dataclass +class TransformerTimingStats: + calls: int = 0 + total_ms: float = 0.0 + min_ms: float = float("inf") + max_ms: float = 0.0 + + @property + def avg_ms(self) -> float: + if self.calls == 0: + return 0.0 + return self.total_ms / self.calls + + +@contextlib.contextmanager +def measure_transformer_forward_time(transformer): + original_forward = transformer.forward + stats = TransformerTimingStats() + + def wrapped_forward(*args, **kwargs): + maybe_synchronize_xpu() + start_time = time.perf_counter() + output = original_forward(*args, **kwargs) + maybe_synchronize_xpu() + elapsed_ms = (time.perf_counter() - start_time) * 1000.0 + stats.calls += 1 + stats.total_ms += elapsed_ms + stats.min_ms = min(stats.min_ms, elapsed_ms) + stats.max_ms = max(stats.max_ms, elapsed_ms) + return output + + transformer.forward = wrapped_forward + try: + yield stats + finally: + transformer.forward = original_forward + + +def sparse_patch_context(): + if not use_sparse: + return contextlib.nullcontext(None) + print( + f"[flux_sparse] enabled attention sparse patch: topk={os.getenv('FLUX_SPARSE_TOPK', '0.5')} " + f"smooth_k={os.getenv('FLUX_SPARSE_SMOOTH_K', '1')}" + f" q_tile={os.getenv('FLUX_SPARSE_Q_TILE_OVERRIDE', '0')}" + f" q_block={os.getenv('FLUX_SPARSE_Q_BLOCK_TOKENS', 'default')}" + f" k_block={os.getenv('FLUX_SPARSE_K_BLOCK_TOKENS', 'default')}" + ) + return patch_flux_sparse_attention_from_env(pipe.transformer) + + +def run_generation(output_type="pil"): + common_kwargs = common_generation_kwargs() + with sparse_patch_context() as sparse_stats: + output = pipe(prompt, output_type=output_type, **common_kwargs).images + if sparse_stats is not None: + log_sparse_stats(sparse_stats) + if output_type == "latent": + return output + return output[0] + + +def prepare_denoising_state(): + generation_kwargs = common_generation_kwargs() + batch_size = 1 + num_images_per_prompt = 1 + execution_device = pipe._execution_device + + prompt_embeds, pooled_prompt_embeds, text_ids = pipe.encode_prompt( + prompt=prompt, + prompt_2=None, + prompt_embeds=None, + pooled_prompt_embeds=None, + device=execution_device, + num_images_per_prompt=num_images_per_prompt, + max_sequence_length=max_sequence_length, + lora_scale=None, + ) + + num_channels_latents = pipe.transformer.config.in_channels // 4 + base_latents, latent_image_ids = pipe.prepare_latents( + batch_size * num_images_per_prompt, + num_channels_latents, + height, + width, + prompt_embeds.dtype, + execution_device, + generation_kwargs["generator"], + latents=None, + ) + + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) + if hasattr(pipe.scheduler.config, "use_flow_sigmas") and pipe.scheduler.config.use_flow_sigmas: + sigmas = None + + image_seq_len = base_latents.shape[1] + mu = calculate_shift( + image_seq_len, + pipe.scheduler.config.get("base_image_seq_len", 256), + pipe.scheduler.config.get("max_image_seq_len", 4096), + pipe.scheduler.config.get("base_shift", 0.5), + pipe.scheduler.config.get("max_shift", 1.15), + ) + timesteps, _ = retrieve_timesteps( + pipe.scheduler, + num_inference_steps, + execution_device, + sigmas=sigmas, + mu=mu, + ) + + if pipe.transformer.config.guidance_embeds: + guidance = torch.full([1], guidance_scale, device=execution_device, dtype=torch.float32) + guidance = guidance.expand(base_latents.shape[0]) + else: + guidance = None + + return { + "prompt_embeds": prompt_embeds, + "pooled_prompt_embeds": pooled_prompt_embeds, + "text_ids": text_ids, + "base_latents": base_latents, + "latent_image_ids": latent_image_ids, + "timesteps": timesteps, + "guidance": guidance, + "joint_attention_kwargs": {}, + } + + +def run_denoising_loop(prepared_state): + latents = prepared_state["base_latents"].clone() + prompt_embeds = prepared_state["prompt_embeds"] + pooled_prompt_embeds = prepared_state["pooled_prompt_embeds"] + text_ids = prepared_state["text_ids"] + latent_image_ids = prepared_state["latent_image_ids"] + timesteps = prepared_state["timesteps"] + guidance = prepared_state["guidance"] + joint_attention_kwargs = dict(prepared_state["joint_attention_kwargs"]) + + pipe.scheduler.set_begin_index(0) + for timestep_value in timesteps: + expanded_timestep = timestep_value.expand(latents.shape[0]).to(latents.dtype) + with pipe.transformer.cache_context("cond"): + noise_pred = pipe.transformer( + hidden_states=latents, + timestep=expanded_timestep / 1000, + guidance=guidance, + pooled_projections=pooled_prompt_embeds, + encoder_hidden_states=prompt_embeds, + txt_ids=text_ids, + img_ids=latent_image_ids, + joint_attention_kwargs=joint_attention_kwargs, + return_dict=False, + )[0] + + latents_dtype = latents.dtype + latents = pipe.scheduler.step(noise_pred, timestep_value, latents, return_dict=False)[0] + if latents.dtype != latents_dtype: + latents = latents.to(latents_dtype) + + return latents + + +def decode_latents_to_image(latents): + unpacked = pipe._unpack_latents(latents, height, width, pipe.vae_scale_factor) + unpacked = (unpacked / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor + image = pipe.vae.decode(unpacked, return_dict=False)[0].detach() + image = pipe.image_processor.postprocess(image, output_type="pil") + pipe.maybe_free_model_hooks() + return image[0] + + +def run_transformer_only_generation(): + latents = run_generation(output_type="latent") + return decode_latents_to_image(latents) + + +def prepare_block_benchmark_state(): + prepared_state = prepare_denoising_state() + transformer = pipe.transformer + joint_attention_kwargs = dict(prepared_state["joint_attention_kwargs"]) + + joint_blocks = transformer.transformer_blocks + single_blocks = transformer.single_transformer_blocks + if benchmark_block_kind == "joint": + if not 0 <= benchmark_block_index < len(joint_blocks): + raise ValueError( + f"FLUX_BENCHMARK_BLOCK_INDEX={benchmark_block_index} is out of range for joint blocks " + f"(0..{len(joint_blocks) - 1})" + ) + target_block = joint_blocks[benchmark_block_index] + else: + if not 0 <= benchmark_block_index < len(single_blocks): + raise ValueError( + f"FLUX_BENCHMARK_BLOCK_INDEX={benchmark_block_index} is out of range for single blocks " + f"(0..{len(single_blocks) - 1})" + ) + target_block = single_blocks[benchmark_block_index] + + timesteps = prepared_state["timesteps"] + if not 0 <= benchmark_block_timestep_index < len(timesteps): + raise ValueError( + f"FLUX_BENCHMARK_BLOCK_TIMESTEP_INDEX={benchmark_block_timestep_index} is out of range " + f"for {len(timesteps)} denoising steps" + ) + + prompt_embeds = prepared_state["prompt_embeds"].to(device) + pooled_prompt_embeds = prepared_state["pooled_prompt_embeds"].to(device) + text_ids = prepared_state["text_ids"].to(device) + latent_image_ids = prepared_state["latent_image_ids"].to(device) + guidance = prepared_state["guidance"] + if guidance is not None: + guidance = guidance.to(device) + + modules_to_move = [ + transformer.x_embedder, + transformer.context_embedder, + transformer.time_text_embed, + ] + if benchmark_block_kind == "joint": + modules_to_move.extend(joint_blocks[: benchmark_block_index + 1]) + else: + modules_to_move.extend(joint_blocks) + modules_to_move.extend(single_blocks[: benchmark_block_index + 1]) + for module in modules_to_move: + module.to(device) + + latents = prepared_state["base_latents"].clone().to(device) + timestep_value = timesteps[benchmark_block_timestep_index] + timestep = timestep_value.expand(latents.shape[0]).to(device=device, dtype=latents.dtype) / 1000 + hidden_states = transformer.x_embedder(latents) + timestep_embed = timestep.to(hidden_states.dtype) * 1000 + if guidance is None: + temb = transformer.time_text_embed(timestep_embed, pooled_prompt_embeds) + else: + temb = transformer.time_text_embed( + timestep_embed, + guidance.to(hidden_states.dtype) * 1000, + pooled_prompt_embeds, + ) + + encoder_hidden_states = transformer.context_embedder(prompt_embeds) + img_ids = latent_image_ids + if text_ids.ndim == 3: + text_ids = text_ids[0] + if img_ids.ndim == 3: + img_ids = img_ids[0] + image_rotary_emb = transformer.pos_embed(torch.cat((text_ids, img_ids), dim=0)) + + with torch.no_grad(): + with transformer.cache_context("cond"): + if benchmark_block_kind == "single": + for block in joint_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + preceding_blocks = ( + joint_blocks[:benchmark_block_index] + if benchmark_block_kind == "joint" + else single_blocks[:benchmark_block_index] + ) + for block in preceding_blocks: + encoder_hidden_states, hidden_states = block( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + temb=temb, + image_rotary_emb=image_rotary_emb, + joint_attention_kwargs=joint_attention_kwargs, + ) + + return { + "block": target_block, + "block_kind": benchmark_block_kind, + "block_index": benchmark_block_index, + "timestep_index": benchmark_block_timestep_index, + "call_kwargs": { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "temb": temb, + "image_rotary_emb": image_rotary_emb, + "joint_attention_kwargs": joint_attention_kwargs, + }, + "input_shapes": { + "hidden_states": tuple(hidden_states.shape), + "encoder_hidden_states": tuple(encoder_hidden_states.shape), + "temb": tuple(temb.shape), + "text_ids": tuple(text_ids.shape), + "img_ids": tuple(img_ids.shape), + }, + } + + +def run_block_call(block_state): + with pipe.transformer.cache_context("cond"): + return block_state["block"](**block_state["call_kwargs"]) + + +def run_block_benchmark(current_run_tag): + warmup = int(os.getenv("FLUX_BENCHMARK_WARMUP", "2")) + iters = int(os.getenv("FLUX_BENCHMARK_ITERS", "3")) + if warmup < 0 or iters <= 0: + raise ValueError("FLUX_BENCHMARK_WARMUP must be >= 0 and FLUX_BENCHMARK_ITERS must be > 0") + + block_state = prepare_block_benchmark_state() + benchmark_path = Path( + os.getenv( + "FLUX_BENCHMARK_OUTPUT", + ( + f"flux_benchmark_block_{block_state['block_kind']}{block_state['block_index']}_{current_run_tag}_" + f"{height}x{width}_{num_inference_steps}steps_seed{seed}_t{block_state['timestep_index']}.json" + ), + ) + ) + benchmark_path.parent.mkdir(parents=True, exist_ok=True) + + print( + f"[flux_bench] benchmarking single FLUX block: warmup={warmup} iters={iters}" + f" mode={current_run_tag} kind={block_state['block_kind']} index={block_state['block_index']}" + f" timestep_index={block_state['timestep_index']}" + ) + print( + "[flux_bench] block input shapes:" + f" hidden_states={block_state['input_shapes']['hidden_states']}" + f" encoder_hidden_states={block_state['input_shapes']['encoder_hidden_states']}" + f" temb={block_state['input_shapes']['temb']}" + ) + if cpu_offload_enabled: + print("[flux_bench] warning: CPU offload is enabled; block benchmark isolation will be weaker") + + with sparse_patch_context() as sparse_stats: + for idx in range(warmup): + start_time = time.perf_counter() + _ = run_block_call(block_state) + maybe_synchronize_xpu() + latency_ms = (time.perf_counter() - start_time) * 1000.0 + print(f"[flux_bench] warmup {idx + 1}/{warmup} complete, latency: {latency_ms:.3f} ms") + + latencies_ms = [] + for idx in range(iters): + start_time = time.perf_counter() + _ = run_block_call(block_state) + maybe_synchronize_xpu() + latency_ms = (time.perf_counter() - start_time) * 1000.0 + latencies_ms.append(latency_ms) + print(f"[flux_bench] iter {idx + 1}/{iters}: {latency_ms:.3f} ms") + + if sparse_stats is not None: + log_sparse_stats(sparse_stats) + + result = { + "scope": "block", + "mode": current_run_tag, + "warmup": warmup, + "iterations": iters, + "latencies_ms": latencies_ms, + "avg_ms": statistics.mean(latencies_ms), + "median_ms": statistics.median(latencies_ms), + "min_ms": min(latencies_ms), + "max_ms": max(latencies_ms), + "config": { + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "max_sequence_length": max_sequence_length, + "seed": seed, + "use_sparse": use_sparse, + "cpu_offload_enabled": cpu_offload_enabled, + "block_kind": block_state["block_kind"], + "block_index": block_state["block_index"], + "block_timestep_index": block_state["timestep_index"], + "input_shapes": block_state["input_shapes"], + "sparse_topk": os.getenv("FLUX_SPARSE_TOPK", "0.5"), + "sparse_smooth_k": os.getenv("FLUX_SPARSE_SMOOTH_K", "1"), + }, + } + benchmark_path.write_text(json.dumps(result, indent=2)) + print( + f"[flux_bench] avg={result['avg_ms']:.3f} ms median={result['median_ms']:.3f} ms " + f"min={result['min_ms']:.3f} ms max={result['max_ms']:.3f} ms" + ) + print(f"[flux_bench] results saved to {benchmark_path}") + pipe.maybe_free_model_hooks() + return result + + +run_tag = "sparse" if use_sparse else "dense" + + +def maybe_run_with_profiler(run_callable, current_run_tag, record_label="flux_generate"): + profiler_enabled = env_flag("FLUX_PROFILER_ENABLE", "0") + profiler, trace_path, summary_path, sort_key = create_profiler(profiler_enabled, current_run_tag) + with profiler: + with torch.profiler.record_function(record_label): + image = run_callable() + maybe_synchronize_xpu() + if profiler_enabled: + save_profiler_results(profiler, trace_path, summary_path, sort_key) + return image + + +def run_benchmark(run_callable, current_run_tag, benchmark_scope): + warmup = int(os.getenv("FLUX_BENCHMARK_WARMUP", "2")) + iters = int(os.getenv("FLUX_BENCHMARK_ITERS", "3")) + if warmup < 0 or iters <= 0: + raise ValueError("FLUX_BENCHMARK_WARMUP must be >= 0 and FLUX_BENCHMARK_ITERS must be > 0") + + benchmark_path = Path( + os.getenv( + "FLUX_BENCHMARK_OUTPUT", + ( + f"flux_benchmark_{current_run_tag}_{height}x{width}_{num_inference_steps}steps_seed{seed}.json" + if benchmark_scope == "full" + else f"flux_benchmark_{benchmark_scope}_{current_run_tag}_{height}x{width}_{num_inference_steps}steps_seed{seed}.json" + ), + ) + ) + benchmark_path.parent.mkdir(parents=True, exist_ok=True) + + print( + f"[flux_bench] benchmarking flux workload: warmup={warmup} iters={iters} " + f"scope={benchmark_scope} mode={current_run_tag} size={height}x{width} steps={num_inference_steps}" + ) + for idx in range(warmup): + start_time = time.perf_counter() + _ = run_callable() + maybe_synchronize_xpu() + latency_ms = (time.perf_counter() - start_time) * 1000.0 + print(f"[flux_bench] warmup {idx + 1}/{warmup} complete, latency: {latency_ms:.3f} ms") + + latencies_ms = [] + image = None + for idx in range(iters): + start_time = time.perf_counter() + image = run_callable() + maybe_synchronize_xpu() + latency_ms = (time.perf_counter() - start_time) * 1000.0 + latencies_ms.append(latency_ms) + print(f"[flux_bench] iter {idx + 1}/{iters}: {latency_ms:.3f} ms") + + result = { + "scope": benchmark_scope, + "mode": current_run_tag, + "warmup": warmup, + "iterations": iters, + "latencies_ms": latencies_ms, + "avg_ms": statistics.mean(latencies_ms), + "median_ms": statistics.median(latencies_ms), + "min_ms": min(latencies_ms), + "max_ms": max(latencies_ms), + "config": { + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "max_sequence_length": max_sequence_length, + "seed": seed, + "use_sparse": use_sparse, + "sparse_topk": os.getenv("FLUX_SPARSE_TOPK", "0.5"), + "sparse_smooth_k": os.getenv("FLUX_SPARSE_SMOOTH_K", "1"), + "output_path": output_path, + }, + } + benchmark_path.write_text(json.dumps(result, indent=2)) + print( + f"[flux_bench] avg={result['avg_ms']:.3f} ms median={result['median_ms']:.3f} ms " + f"min={result['min_ms']:.3f} ms max={result['max_ms']:.3f} ms" + ) + print(f"[flux_bench] results saved to {benchmark_path}") + return image + + +def run_denoising_benchmark(current_run_tag): + warmup = int(os.getenv("FLUX_BENCHMARK_WARMUP", "2")) + iters = int(os.getenv("FLUX_BENCHMARK_ITERS", "3")) + if warmup < 0 or iters <= 0: + raise ValueError("FLUX_BENCHMARK_WARMUP must be >= 0 and FLUX_BENCHMARK_ITERS must be > 0") + + benchmark_path = Path( + os.getenv( + "FLUX_BENCHMARK_OUTPUT", + f"flux_benchmark_denoising_{current_run_tag}_{height}x{width}_{num_inference_steps}steps_seed{seed}.json", + ) + ) + benchmark_path.parent.mkdir(parents=True, exist_ok=True) + + print( + f"[flux_bench] benchmarking transformer forwards via pipeline latent output: warmup={warmup} iters={iters} " + f"mode={current_run_tag} size={height}x{width} steps={num_inference_steps}" + ) + + common_kwargs = common_generation_kwargs() + with sparse_patch_context() as sparse_stats: + with measure_transformer_forward_time(pipe.transformer) as timing_stats: + for idx in range(warmup): + _ = pipe(prompt, output_type="latent", **common_kwargs).images + print( + f"[flux_bench] warmup {idx + 1}/{warmup} complete," + f" transformer_total_ms={timing_stats.total_ms:.3f}" + ) + + latencies_ms = [] + final_latents = None + prev_total_ms = timing_stats.total_ms + prev_calls = timing_stats.calls + for idx in range(iters): + final_latents = pipe(prompt, output_type="latent", **common_kwargs).images + iter_total_ms = timing_stats.total_ms - prev_total_ms + iter_calls = timing_stats.calls - prev_calls + latencies_ms.append(iter_total_ms) + prev_total_ms = timing_stats.total_ms + prev_calls = timing_stats.calls + print( + f"[flux_bench] iter {idx + 1}/{iters}: transformer_total_ms={iter_total_ms:.3f}" + f" transformer_calls={iter_calls}" + ) + + if sparse_stats is not None: + log_sparse_stats(sparse_stats) + + if timing_stats.calls: + print( + "[flux_bench] transformer timing summary:" + f" calls={timing_stats.calls}" + f" avg_call_ms={timing_stats.avg_ms:.3f}" + f" min_call_ms={timing_stats.min_ms:.3f}" + f" max_call_ms={timing_stats.max_ms:.3f}" + ) + + result = { + "scope": "denoising", + "mode": current_run_tag, + "warmup": warmup, + "iterations": iters, + "latencies_ms": latencies_ms, + "avg_ms": statistics.mean(latencies_ms), + "median_ms": statistics.median(latencies_ms), + "min_ms": min(latencies_ms), + "max_ms": max(latencies_ms), + "config": { + "height": height, + "width": width, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "max_sequence_length": max_sequence_length, + "seed": seed, + "use_sparse": use_sparse, + "sparse_topk": os.getenv("FLUX_SPARSE_TOPK", "0.5"), + "sparse_smooth_k": os.getenv("FLUX_SPARSE_SMOOTH_K", "1"), + "output_path": output_path, + }, + } + benchmark_path.write_text(json.dumps(result, indent=2)) + print( + f"[flux_bench] avg={result['avg_ms']:.3f} ms median={result['median_ms']:.3f} ms " + f"min={result['min_ms']:.3f} ms max={result['max_ms']:.3f} ms" + ) + print(f"[flux_bench] results saved to {benchmark_path}") + + if final_latents is None: + raise RuntimeError("Denoising benchmark did not produce final latents") + pipe.maybe_free_model_hooks() + return decode_latents_to_image(final_latents) + + +if benchmark_enabled: + print(f"=======start benchmark for {run_tag} mode=======") + if benchmark_scope == "full": + image = run_benchmark(run_generation, run_tag, benchmark_scope) + elif benchmark_scope == "denoising": + image = run_denoising_benchmark(run_tag) + else: + image = None + _ = run_block_benchmark(run_tag) + if profiler_enabled: + print("[flux_bench] capturing profiler trace in a separate run; benchmark latencies exclude profiler overhead") + if benchmark_scope == "full": + profile_callable = run_generation + profile_label = "flux_full" + elif benchmark_scope == "denoising": + profile_callable = run_transformer_only_generation + profile_label = "flux_denoising" + else: + profile_label = "flux_block" + with sparse_patch_context() as sparse_stats: + block_profile_state = prepare_block_benchmark_state() + _ = maybe_run_with_profiler( + lambda: run_block_call(block_profile_state), + f"{run_tag}_{benchmark_scope}_bench_profile", + record_label=profile_label, + ) + if sparse_stats is not None: + log_sparse_stats(sparse_stats) + if benchmark_scope != "block": + _ = maybe_run_with_profiler( + profile_callable, + f"{run_tag}_{benchmark_scope}_bench_profile", + record_label=profile_label, + ) +else: + image = maybe_run_with_profiler(run_generation, run_tag) + +if image is not None: + print(f"Saving output to {output_path}") + image.save(output_path) diff --git a/auto_round_extension/ark/examples/run_flux_sweep.sh b/auto_round_extension/ark/examples/run_flux_sweep.sh new file mode 100755 index 000000000..fabcdb6dd --- /dev/null +++ b/auto_round_extension/ark/examples/run_flux_sweep.sh @@ -0,0 +1,162 @@ +#!/usr/bin/env bash + +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +RESULTS_ROOT_DEFAULT="${SCRIPT_DIR}/results/flux_sweep_defaultsteps_$(date -u +%Y%m%dT%H%M%SZ)" +RESULTS_ROOT=${FLUX_SWEEP_OUTPUT_ROOT:-${RESULTS_ROOT_DEFAULT}} +PYTHON_BIN=${FLUX_SWEEP_PYTHON:-python} +MODEL_PATH=${FLUX_MODEL:-black-forest-labs/FLUX.1-dev} +RUN_FLUX_PY="${SCRIPT_DIR}/run_flux.py" +PROMPT_VALUE=${FLUX_PROMPT:-A cat holding a sign that says hello world} + +mkdir -p "${RESULTS_ROOT}" + +IFS=, read -r -a DEVICES <<< "${FLUX_SWEEP_DEVICES:-0,1,2,3,4,5,6,7}" +if [[ ${#DEVICES[@]} -eq 0 ]]; then + echo "FLUX_SWEEP_DEVICES must provide at least one device id" >&2 + exit 1 +fi + +if [[ -n "${FLUX_SWEEP_CONFIGS:-}" ]]; then + read -r -a CONFIGS <<< "${FLUX_SWEEP_CONFIGS}" +else + CONFIGS=("dense" "1.0" "0.9" "0.8" "0.7" "0.6" "0.5" "0.4" "0.3" "0.2" "0.1") +fi +SUMMARY_CSV="${RESULTS_ROOT}/summary.csv" +COMMANDS_TXT="${RESULTS_ROOT}/commands.txt" + +sanitize_topk_label() { + local topk="$1" + echo "topk${topk/./p}" +} + +write_status_file() { + local status_path="$1" + shift + : > "${status_path}" + while (($#)); do + local key="$1" + local value="$2" + printf '%s=%q\n' "${key}" "${value}" >> "${status_path}" + shift 2 + done +} + +active_jobs=0 +launch_index=0 + +for config in "${CONFIGS[@]}"; do + device="${DEVICES[$((launch_index % ${#DEVICES[@]}))]}" + if [[ "${config}" == "dense" ]]; then + label="dense" + mode="dense" + topk_value="dense" + sparse_env=( + FLUX_USE_SPARSE=0 + ) + else + label=$(sanitize_topk_label "${config}") + mode="sparse" + topk_value="${config}" + sparse_env=( + FLUX_USE_SPARSE=1 + FLUX_SPARSE_TOPK="${config}" + FLUX_SPARSE_Q_TILE_OVERRIDE=256 + FLUX_SPARSE_Q_BLOCK_TOKENS=256 + FLUX_SPARSE_K_BLOCK_TOKENS=64 + ) + fi + + png_path="${RESULTS_ROOT}/flux_${label}.png" + log_path="${RESULTS_ROOT}/flux_${label}.log" + status_path="${RESULTS_ROOT}/flux_${label}.status" + cmd_path="${RESULTS_ROOT}/flux_${label}.cmd" + + cat > "${cmd_path}" < "${log_path}" 2>&1 || exit_code=$? + end_iso=$(date -u +%Y-%m-%dT%H:%M:%SZ) + end_epoch=$(date +%s) + elapsed_sec=$((end_epoch - start_epoch)) + write_status_file "${status_path}" \ + label "${label}" \ + mode "${mode}" \ + topk "${topk_value}" \ + device "${device}" \ + exit_code "${exit_code}" \ + start_iso "${start_iso}" \ + end_iso "${end_iso}" \ + elapsed_sec "${elapsed_sec}" \ + image_path "${png_path}" \ + log_path "${log_path}" \ + cmd_path "${cmd_path}" + if [[ ${exit_code} -ne 0 ]]; then + exit "${exit_code}" + fi + } & + + launch_index=$((launch_index + 1)) + active_jobs=$((active_jobs + 1)) + if [[ ${active_jobs} -ge ${#DEVICES[@]} ]]; then + wait -n || true + active_jobs=$((active_jobs - 1)) + fi +done + +wait || true + +{ + echo "# Flux sweep commands" + echo "# Generated at $(date -u +%Y-%m-%dT%H:%M:%SZ)" + echo + for config in "${CONFIGS[@]}"; do + if [[ "${config}" == "dense" ]]; then + label="dense" + else + label=$(sanitize_topk_label "${config}") + fi + cat "${RESULTS_ROOT}/flux_${label}.cmd" + echo + done +} > "${COMMANDS_TXT}" + +{ + echo "label,mode,topk,device,exit_code,elapsed_sec,image_path,log_path,cmd_path" + for config in "${CONFIGS[@]}"; do + if [[ "${config}" == "dense" ]]; then + label="dense" + else + label=$(sanitize_topk_label "${config}") + fi + # shellcheck disable=SC1090 + source "${RESULTS_ROOT}/flux_${label}.status" + printf '%s,%s,%s,%s,%s,%s,%s,%s,%s\n' \ + "${label}" "${mode}" "${topk}" "${device}" "${exit_code}" "${elapsed_sec}" \ + "${image_path}" "${log_path}" "${cmd_path}" + done +} > "${SUMMARY_CSV}" + +echo "Saved FLUX sweep outputs to ${RESULTS_ROOT}" +echo "Summary: ${SUMMARY_CSV}" +echo "Commands: ${COMMANDS_TXT}" diff --git a/auto_round_extension/ark/examples/run_wan.py b/auto_round_extension/ark/examples/run_wan.py new file mode 100644 index 000000000..263aff254 --- /dev/null +++ b/auto_round_extension/ark/examples/run_wan.py @@ -0,0 +1,121 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import os + +import torch +from diffusers import AutoencoderKLWan, WanPipeline +from diffusers.utils import export_to_video +from wan_sparse_patch import patch_wan_sparse_attention_from_env + +dtype = torch.bfloat16 +device = "xpu" + + +def env_flag(name, default="0"): + return os.getenv(name, default).lower() not in {"0", "false", "no", "off"} + + +def env_value(name): + value = os.getenv(name, "").strip() + return value or None + + +model_id = os.getenv("WAN_MODEL", "Wan-AI/Wan2.2-T2V-A14B-Diffusers") +device_map = env_value("WAN_DEVICE_MAP") +vae = AutoencoderKLWan.from_pretrained(model_id, subfolder="vae", torch_dtype=torch.float32) +pipe_kwargs = {"vae": vae, "torch_dtype": dtype} +if device_map is not None: + pipe_kwargs["device_map"] = device_map +pipe = WanPipeline.from_pretrained(model_id, **pipe_kwargs) +cpu_offload_enabled = env_flag("WAN_ENABLE_CPU_OFFLOAD", "1") +cpu_offload_mode = os.getenv("WAN_CPU_OFFLOAD_MODE", "model").strip().lower() +if device_map is not None: + if cpu_offload_enabled: + raise ValueError("WAN_DEVICE_MAP cannot be combined with WAN_ENABLE_CPU_OFFLOAD=1") + print(f"[wan_sparse] using device_map={device_map} hf_device_map={pipe.hf_device_map}") +elif cpu_offload_enabled: + if cpu_offload_mode == "model": + pipe.enable_model_cpu_offload(device=device) + elif cpu_offload_mode == "sequential": + pipe.enable_sequential_cpu_offload(device=device) + else: + raise ValueError("WAN_CPU_OFFLOAD_MODE must be one of: model, sequential") +else: + pipe.to(device) +# Match the Wan2.2 T2V-A14B diffusers model-card example by default. +height = int(os.getenv("WAN_HEIGHT", "720")) +width = int(os.getenv("WAN_WIDTH", "1280")) +num_frames = int(os.getenv("WAN_NUM_FRAMES", "81")) +num_inference_steps = int(os.getenv("WAN_STEPS", "40")) +guidance_scale = float(os.getenv("WAN_GUIDANCE_SCALE", "4.0")) +guidance_scale_2 = float(os.getenv("WAN_GUIDANCE_SCALE_2", "3.0")) +use_sparse = os.getenv("WAN_USE_SPARSE", "0").lower() not in {"0", "false", "no", "off"} +if use_sparse: + output_path_file = ( + f"wan_sparse_topk{os.getenv('WAN_SPARSE_TOPK', '0.75')}_" + f"{num_frames}f_{num_inference_steps}steps_{guidance_scale}gs.mp4" + ) +else: + output_path_file = f"wan_dense_{num_frames}f_{num_inference_steps}steps_{guidance_scale}gs.mp4" +output_path = os.getenv("WAN_OUTPUT", output_path_file) + + +prompt = os.getenv( + "WAN_PROMPT", + "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.", +) +negative_prompt = os.getenv( + "WAN_NEGATIVE_PROMPT", + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走", +) +fps = int(os.getenv("WAN_FPS", "16")) + +if use_sparse: + print( + f"[wan_sparse] enabled sparse patch: topk={os.getenv('WAN_SPARSE_TOPK', '0.75')} " + f"smooth_k={os.getenv('WAN_SPARSE_SMOOTH_K', '1')}" + f" cross={os.getenv('WAN_SPARSE_ENABLE_CROSS_ATTN', '0')}" + f" q_tile={os.getenv('WAN_SPARSE_Q_TILE_OVERRIDE', '0')}" + f" q_block={os.getenv('WAN_SPARSE_Q_BLOCK_TOKENS', 'default')}" + f" k_block={os.getenv('WAN_SPARSE_K_BLOCK_TOKENS', 'default')}" + ) + with patch_wan_sparse_attention_from_env(pipe.transformer) as sparse_stats: + output = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + num_inference_steps=num_inference_steps, + ).frames[0] + print( + "[wan_sparse] stats:" + f" self_attn_total={sparse_stats.self_attn_total}" + f" cross_attn_total={sparse_stats.cross_attn_total}" + f" sparse_self_attn_calls={sparse_stats.sparse_self_attn_calls}" + f" sparse_cross_attn_calls={sparse_stats.sparse_cross_attn_calls}" + f" cross_attn_fallbacks={sparse_stats.cross_attn_fallbacks}" + f" cross_attn_policy_fallbacks={sparse_stats.cross_attn_policy_fallbacks}" + f" cross_attn_unsupported_fallbacks={sparse_stats.cross_attn_unsupported_fallbacks}" + f" cross_attn_runtime_fallbacks={sparse_stats.cross_attn_runtime_fallbacks}" + f" unsupported_fallbacks={sparse_stats.unsupported_fallbacks}" + f" runtime_fallbacks={sparse_stats.runtime_fallbacks}" + f" avg_sparsity={sparse_stats.avg_sparsity:.4f}" + ) +else: + output = pipe( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + num_frames=num_frames, + guidance_scale=guidance_scale, + guidance_scale_2=guidance_scale_2, + num_inference_steps=num_inference_steps, + ).frames[0] + +export_to_video(output, output_path, fps=fps) +print(f"[wan_sparse] wrote {output_path}") diff --git a/auto_round_extension/ark/examples/wan_sparse_patch.py b/auto_round_extension/ark/examples/wan_sparse_patch.py new file mode 100644 index 000000000..8475d3971 --- /dev/null +++ b/auto_round_extension/ark/examples/wan_sparse_patch.py @@ -0,0 +1,333 @@ +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import contextlib +import importlib.util +import os +import sys +import sysconfig +import warnings +from dataclasses import dataclass +from pathlib import Path + +import torch + +EXAMPLES_DIR = Path(__file__).resolve().parent +ARK_DIR = EXAMPLES_DIR.parent +KERNEL_DIR = ARK_DIR / "auto_round_kernel" +if str(ARK_DIR) not in sys.path: + sys.path.insert(0, str(ARK_DIR)) + +import auto_round_kernel as ark +from diffusers.models.transformers.transformer_wan import WanAttention, _get_qkv_projections + + +def _parse_bool_env(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.lower() not in {"0", "false", "no", "off"} + + +def _parse_optional_int_env(name: str) -> int | None: + value = os.getenv(name) + if value is None or value.strip() == "": + return None + parsed = int(value) + return None if parsed == 0 else parsed + + +def ensure_ark_sparse_binding() -> None: + if getattr(ark, "xpu_lib", None) is not None and hasattr(ark.xpu_lib, "sage_sparse"): + return + + ext_suffix = sysconfig.get_config_var("EXT_SUFFIX") + if not ext_suffix: + raise RuntimeError("Unable to determine Python extension suffix for the current interpreter") + + search_roots = [ + KERNEL_DIR / "xbuild", + KERNEL_DIR / "xbuild_diffuser", + KERNEL_DIR, + ] + candidates = [] + for root in search_roots: + if root.exists(): + candidates.extend(sorted(root.glob(f"auto_round_kernel_xpu*{ext_suffix}"))) + if not candidates: + raise RuntimeError( + "Unable to locate a built XPU extension matching the current Python ABI. " + f"Expected suffix {ext_suffix!r} under one of {[str(p) for p in search_roots]}." + ) + + ext_path = candidates[-1] + spec = importlib.util.spec_from_file_location("auto_round_kernel_xpu", ext_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load extension spec from {ext_path}") + module = importlib.util.module_from_spec(spec) + sys.modules["auto_round_kernel_xpu"] = module + spec.loader.exec_module(module) + required = ("sage_sparse", "sage_dynamic_quant_layout") + missing = [name for name in required if not hasattr(module, name)] + if missing: + raise RuntimeError(f"Loaded extension is missing required XPU bindings {missing}: {ext_path}") + ark.xpu_lib = module + + +def _apply_rotary_emb(hidden_states: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor: + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + +def _normalize_attention_mask( + attn_mask: torch.Tensor | None, + batch: int, + seq_q: int, + seq_kv: int, + device: torch.device, +) -> torch.Tensor | None: + if attn_mask is None: + return None + + mask = attn_mask + if mask.dtype == torch.bool: + zero = torch.zeros((), dtype=torch.float32, device=device) + neg_inf = torch.full((), -1.0e9, dtype=torch.float32, device=device) + mask = torch.where(mask, zero, neg_inf) + else: + mask = mask.to(torch.float32) + + if mask.ndim == 2: + mask = mask.view(1, 1, seq_q, seq_kv) + elif mask.ndim == 3: + mask = mask.unsqueeze(1) + elif mask.ndim != 4: + raise ValueError(f"Unsupported Wan attention mask rank: {mask.ndim}") + + if mask.shape[0] == 1 and batch != 1: + mask = mask.expand(batch, -1, -1, -1) + return mask.contiguous() + + +@dataclass +class WanSparseAttentionStats: + self_attn_total: int = 0 + cross_attn_total: int = 0 + sparse_self_attn_calls: int = 0 + sparse_cross_attn_calls: int = 0 + cross_attn_fallbacks: int = 0 + cross_attn_policy_fallbacks: int = 0 + cross_attn_unsupported_fallbacks: int = 0 + cross_attn_runtime_fallbacks: int = 0 + unsupported_fallbacks: int = 0 + runtime_fallbacks: int = 0 + sparse_sparsity_sum: float = 0.0 + + @property + def avg_sparsity(self) -> float: + sparse_calls = self.sparse_self_attn_calls + self.sparse_cross_attn_calls + if sparse_calls == 0: + return 0.0 + return self.sparse_sparsity_sum / sparse_calls + + +class WanSparseAttnProcessor: + def __init__( + self, + original_processor, + stats: WanSparseAttentionStats, + *, + smooth_k: bool, + topk: float, + attention_sink: bool, + enable_cross_attention: bool, + q_tile_override: int, + sparse_q_block_tokens: int | None, + sparse_k_block_tokens: int | None, + ): + self.original_processor = original_processor + self.stats = stats + self.smooth_k = smooth_k + self.topk = topk + self.attention_sink = attention_sink + self.enable_cross_attention = enable_cross_attention + self.q_tile_override = q_tile_override + self.sparse_q_block_tokens = sparse_q_block_tokens + self.sparse_k_block_tokens = sparse_k_block_tokens + self._warned_runtime_fallback = False + + def __call__( + self, + attn: "WanAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + **kwargs, + ) -> torch.Tensor: + if attn.is_cross_attention: + self.stats.cross_attn_total += 1 + if not self.enable_cross_attention: + self.stats.cross_attn_fallbacks += 1 + self.stats.cross_attn_policy_fallbacks += 1 + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + rotary_emb, + **kwargs, + ) + else: + self.stats.self_attn_total += 1 + + if encoder_hidden_states is not None and not attn.is_cross_attention: + self.stats.unsupported_fallbacks += 1 + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + rotary_emb, + **kwargs, + ) + + if hidden_states.device.type != "xpu" or hidden_states.dtype not in (torch.float16, torch.bfloat16): + if attn.is_cross_attention: + self.stats.cross_attn_fallbacks += 1 + self.stats.cross_attn_unsupported_fallbacks += 1 + self.stats.unsupported_fallbacks += 1 + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + rotary_emb, + **kwargs, + ) + + try: + query, key, value = _get_qkv_projections(attn, hidden_states, encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)).contiguous() + key = key.unflatten(2, (attn.heads, -1)).contiguous() + value = value.unflatten(2, (attn.heads, -1)).contiguous() + + if rotary_emb is not None: + query = _apply_rotary_emb(query, *rotary_emb) + key = _apply_rotary_emb(key, *rotary_emb) + + mask = _normalize_attention_mask( + attention_mask, + batch=query.shape[0], + seq_q=query.shape[1], + seq_kv=key.shape[1], + device=query.device, + ) + + attn_out, sparsity = ark.sparge_sage2_attn_meansim_topk_xpu( + query, + key, + value, + attn_mask=mask, + is_causal=False, + smooth_k=self.smooth_k, + simthreshd1=-1.0, + topk=self.topk, + attention_sink=self.attention_sink, + tensor_layout="NHD", + q_tile_override=self.q_tile_override, + sparse_q_block_tokens=self.sparse_q_block_tokens, + sparse_k_block_tokens=self.sparse_k_block_tokens, + return_sparsity=True, + ) + if attn.is_cross_attention: + self.stats.sparse_cross_attn_calls += 1 + else: + self.stats.sparse_self_attn_calls += 1 + self.stats.sparse_sparsity_sum += float(sparsity) + attn_out = attn_out.flatten(2, 3).type_as(query) + attn_out = attn.to_out[0](attn_out) + attn_out = attn.to_out[1](attn_out) + return attn_out + except Exception as exc: # noqa: BLE001 + if attn.is_cross_attention: + self.stats.cross_attn_fallbacks += 1 + self.stats.cross_attn_runtime_fallbacks += 1 + self.stats.runtime_fallbacks += 1 + if not self._warned_runtime_fallback: + warnings.warn( + f"Wan sparse attention fell back to the original processor after a runtime error: {exc}", + stacklevel=2, + ) + self._warned_runtime_fallback = True + return self.original_processor( + attn, + hidden_states, + encoder_hidden_states, + attention_mask, + rotary_emb, + **kwargs, + ) + + +@contextlib.contextmanager +def patch_wan_sparse_attention( + transformer, + *, + smooth_k: bool = True, + topk: float = 0.75, + attention_sink: bool = False, + enable_cross_attention: bool = False, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +): + ensure_ark_sparse_binding() + originals: list[tuple[WanAttention, object]] = [] + stats = WanSparseAttentionStats() + + for module in transformer.modules(): + if isinstance(module, WanAttention): + original = module.processor + module.set_processor( + WanSparseAttnProcessor( + original, + stats, + smooth_k=smooth_k, + topk=topk, + attention_sink=attention_sink, + enable_cross_attention=enable_cross_attention, + q_tile_override=q_tile_override, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + ) + ) + originals.append((module, original)) + + try: + yield stats + finally: + for module, original in originals: + module.set_processor(original) + + +def patch_wan_sparse_attention_from_env(transformer): + return patch_wan_sparse_attention( + transformer, + smooth_k=_parse_bool_env("WAN_SPARSE_SMOOTH_K", True), + topk=float(os.getenv("WAN_SPARSE_TOPK", "0.75")), + attention_sink=_parse_bool_env("WAN_SPARSE_ATTENTION_SINK", False), + enable_cross_attention=_parse_bool_env("WAN_SPARSE_ENABLE_CROSS_ATTN", False), + q_tile_override=int(os.getenv("WAN_SPARSE_Q_TILE_OVERRIDE", "0")), + sparse_q_block_tokens=_parse_optional_int_env("WAN_SPARSE_Q_BLOCK_TOKENS"), + sparse_k_block_tokens=_parse_optional_int_env("WAN_SPARSE_K_BLOCK_TOKENS"), + ) diff --git a/auto_round_extension/ark/test/bench_sparse_topk.py b/auto_round_extension/ark/test/bench_sparse_topk.py new file mode 100644 index 000000000..a11857601 --- /dev/null +++ b/auto_round_extension/ark/test/bench_sparse_topk.py @@ -0,0 +1,943 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +import argparse +import csv +import importlib.util +import math +import sys +import time +from pathlib import Path + +import torch +import torch.nn.functional as F + +REPO_ROOT = Path(__file__).resolve().parents[3] +LOCAL_ARK_PARENT = REPO_ROOT / "auto_round_extension" / "ark" +if str(LOCAL_ARK_PARENT) not in sys.path: + sys.path.insert(0, str(LOCAL_ARK_PARENT)) + +import auto_round_kernel as ark + + +def _load_sparse_binding(ext_path: Path) -> None: + ext_path = ext_path.resolve() + if not ext_path.is_file(): + raise RuntimeError(f"Unable to locate XPU extension: {ext_path}") + module_name = "auto_round_kernel._bench.auto_round_kernel_xpu" + print(f"Loading XPU extension from {ext_path} as {module_name}") + spec = importlib.util.spec_from_file_location(module_name, ext_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load extension spec from {ext_path}") + sys.modules.pop(module_name, None) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + if not hasattr(module, "sage_sparse"): + raise RuntimeError(f"Loaded extension does not expose sage_sparse: {ext_path}") + ark.xpu_lib = module + + +def _resolve_sparse_binding(args: argparse.Namespace) -> Path | None: + local_kernel_dir = REPO_ROOT / "auto_round_extension" / "ark" / "auto_round_kernel" + current_file = getattr(getattr(ark, "xpu_lib", None), "__file__", None) + if args.xpu_so is not None: + return args.xpu_so.resolve() + if args.xbuild_dir is not None: + candidates = sorted(args.xbuild_dir.resolve().glob("auto_round_kernel_xpu*.so")) + if not candidates: + raise RuntimeError(f"Unable to locate built XPU extension with sage_sparse in {args.xbuild_dir}") + return candidates[-1] + if current_file is not None and hasattr(ark.xpu_lib, "sage_sparse"): + try: + current_path = Path(current_file).resolve() + if current_path.is_relative_to(local_kernel_dir.resolve()): + print(f"Using already-loaded local XPU extension from {current_path}") + return None + except Exception: + pass + candidates = sorted((local_kernel_dir / "xbuild").glob("auto_round_kernel_xpu*.so")) + if not candidates: + raise RuntimeError("Unable to locate built XPU extension with sage_sparse in auto_round_kernel/xbuild") + return candidates[-1] + + +def ensure_sparse_binding(args: argparse.Namespace) -> None: + ext_path = _resolve_sparse_binding(args) + if ext_path is not None: + _load_sparse_binding(ext_path) + + +def is_xpu_available() -> bool: + return hasattr(torch, "xpu") and torch.xpu.is_available() + + +def bench(fn, warmup: int, iters: int) -> float: + for _ in range(warmup): + out = fn() + del out + torch.xpu.synchronize() + start = time.perf_counter() + for _ in range(iters): + out = fn() + del out + torch.xpu.synchronize() + return (time.perf_counter() - start) * 1000.0 / float(iters) + + +def bench_with_output(fn, warmup: int, iters: int) -> tuple[object, float]: + out = None + for _ in range(warmup): + out = fn() + del out + torch.xpu.synchronize() + start = time.perf_counter() + for i in range(iters): + out = fn() + if i != iters - 1: + del out + torch.xpu.synchronize() + return out, (time.perf_counter() - start) * 1000.0 / float(iters) + + +def classify_exception(exc: Exception) -> tuple[str, str]: + text = f"{type(exc).__name__}: {exc}" + lowered = text.lower() + if "out of memory" in lowered or "oom" in lowered: + return "oom", text + return "error", text + + +def empty_xpu_cache() -> None: + if is_xpu_available(): + torch.xpu.synchronize() + torch.xpu.empty_cache() + + +def build_inputs( + batch: int, + num_heads_q: int, + num_heads_kv: int, + seq_len: int, + head_dim: int, + dtype: torch.dtype, + device: torch.device, +): + torch.manual_seed(20260611) + q = torch.randn((batch, num_heads_q, seq_len, head_dim), dtype=dtype, device=device) + k = torch.randn((batch, num_heads_kv, seq_len, head_dim), dtype=dtype, device=device) + v = torch.randn((batch, num_heads_kv, seq_len, head_dim), dtype=dtype, device=device) + return q, k, v + + +def attention_dense_flops(batch: int, num_heads_q: int, seq_len_q: int, seq_len_kv: int, head_dim: int) -> float: + return float(4.0 * batch * num_heads_q * seq_len_q * seq_len_kv * head_dim) + + +def flops_to_tflops(flops: float, latency_ms: float) -> float: + return flops / (latency_ms * 1.0e-3) / 1.0e12 + + +def make_row( + *, + mode: str, + batch: int, + num_heads_q: int, + num_heads_kv: int, + seq_len: int, + head_dim: int, + dtype: torch.dtype, + is_causal: bool, + warmup: int, + iters: int, + requested_topk: float | None, + selected_ratio: float | None, + selected_blocks_per_row: float | None, + latency_ms: float | None, + status: str, + note: str = "", +) -> dict[str, object]: + return { + "mode": mode, + "batch": batch, + "num_heads_q": num_heads_q, + "num_heads_kv": num_heads_kv, + "seq_len": seq_len, + "head_dim": head_dim, + "dtype": str(dtype).replace("torch.", ""), + "causal": is_causal, + "requested_topk": requested_topk, + "selected_ratio": selected_ratio, + "selected_blocks_per_row": selected_blocks_per_row, + "warmup": warmup, + "iters": iters, + "latency_ms": latency_ms, + "status": status, + "note": note, + } + + +def try_benchmark( + mode: str, + fn, + *, + batch: int, + num_heads_q: int, + num_heads_kv: int, + seq_len: int, + head_dim: int, + dtype: torch.dtype, + is_causal: bool, + warmup: int, + iters: int, + requested_topk: float | None = None, + selected_ratio: float | None = None, + selected_blocks_per_row: float | None = None, +) -> dict[str, object]: + try: + latency_ms = bench(fn, warmup, iters) + return make_row( + mode=mode, + batch=batch, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + seq_len=seq_len, + head_dim=head_dim, + dtype=dtype, + is_causal=is_causal, + warmup=warmup, + iters=iters, + requested_topk=requested_topk, + selected_ratio=selected_ratio, + selected_blocks_per_row=selected_blocks_per_row, + latency_ms=latency_ms, + status="ok", + ) + except Exception as exc: + status, note = classify_exception(exc) + empty_xpu_cache() + return make_row( + mode=mode, + batch=batch, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + seq_len=seq_len, + head_dim=head_dim, + dtype=dtype, + is_causal=is_causal, + warmup=warmup, + iters=iters, + requested_topk=requested_topk, + selected_ratio=selected_ratio, + selected_blocks_per_row=selected_blocks_per_row, + latency_ms=None, + status=status, + note=note, + ) + + +def benchmark_preprocess_stages( + q: torch.Tensor, + k: torch.Tensor, + *, + topk: float, + is_causal: bool, + quant_block_size: int, + tensor_layout: str, + warmup: int, + iters: int, +) -> tuple[dict[str, float], dict[str, object]]: + from auto_round_kernel.sparge_preprocess_triton import ( + _block_map_lut_triton, + _fill_block_map_triton, + _get_pool_sim_triton_simmean_fuse_quant, + ) + from auto_round_kernel.sparge_preprocess_triton import _safe_softmax as _triton_safe_softmax + + stage_names = [ + "key_mean", + "pool_q_block64", + "pool_k_block64", + "pool_q_tile", + "pooled_score_softmax_sort", + "fill_block_map", + "tile_to_qblock_index", + "block_map_to_lut", + ] + sums = {name: 0.0 for name in stage_names} + last_meta = None + + for _ in range(warmup + iters): + ctx = ark._build_sparge_preprocess_context( + q, + k, + is_causal=is_causal, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + quant_block_size=quant_block_size, + tensor_layout=tensor_layout, + ) + stage_ms: dict[str, float] = {} + + key_mean, stage_ms["key_mean"] = bench_with_output( + lambda: ark._sequence_mean_native_layout(ctx.key, ctx.tensor_layout), + warmup=0, + iters=1, + ) + (pooled_q, sim_qblocks, q_int8_hnd, q_scale), stage_ms["pool_q_block64"] = bench_with_output( + lambda: _get_pool_sim_triton_simmean_fuse_quant( + ctx.query, + None, + ctx.quant_block_size, + ctx.simthreshd1, + ctx.tensor_layout, + ), + warmup=0, + iters=1, + ) + (pooled_k, sim_kblocks, k_int8_hnd, k_scale), stage_ms["pool_k_block64"] = bench_with_output( + lambda: _get_pool_sim_triton_simmean_fuse_quant( + ctx.key, + key_mean, + ctx.quant_block_size, + ctx.simthreshd1[: ctx.num_heads_kv], + ctx.tensor_layout, + ), + warmup=0, + iters=1, + ) + + def qtile_stage(): + if ctx.blocks_per_qtile <= 1: + return pooled_q, sim_qblocks + tile_pooled_q = [] + tile_sim_q = [] + for qtile in range(ctx.num_q_tiles): + qblk_start = qtile * ctx.blocks_per_qtile + qblk_end = min(qblk_start + ctx.blocks_per_qtile, pooled_q.size(2)) + tile_tokens = ark._slice_sequence_native_layout( + ctx.query, + ctx.tensor_layout, + qblk_start * ctx.quant_block_size, + min((qblk_end * ctx.quant_block_size), ctx.seq_len_q), + ) + pooled_tile, sim_tile, _, _ = _get_pool_sim_triton_simmean_fuse_quant( + tile_tokens, + None, + ctx.query_tile_tokens, + ctx.simthreshd1, + ctx.tensor_layout, + ) + tile_pooled_q.append(pooled_tile[:, :, 0, :]) + tile_sim_q.append(sim_tile[:, :, 0]) + return torch.stack(tile_pooled_q, dim=2), torch.stack(tile_sim_q, dim=2) + + (pooled_q_for_routing, sim_q_for_routing), stage_ms["pool_q_tile"] = bench_with_output( + qtile_stage, + warmup=0, + iters=1, + ) + + kv_head_index = torch.arange(ctx.num_heads_q, device=ctx.query.device, dtype=torch.int64) // ( + ctx.num_heads_q // ctx.num_heads_kv + ) + + def routing_stage(): + pooled_k_for_q = pooled_k[:, kv_head_index] + sim_k_for_q = sim_kblocks[:, kv_head_index] + sim_k_expand = sim_k_for_q.unsqueeze(-2).expand(-1, -1, ctx.num_q_tiles, -1) + sim_q_expand = sim_q_for_routing.unsqueeze(-1).expand(-1, -1, -1, pooled_k.size(2)) + pooled_score = torch.matmul( + pooled_q_for_routing.to(torch.float32), + pooled_k_for_q.transpose(-1, -2).to(torch.float32), + ) + pooled_score *= ctx.head_dim**-0.5 + pooled_score = pooled_score.masked_fill(~sim_k_expand, -torch.inf) + if ctx.is_causal: + causal_mask = ark._build_block_causal_mask( + ctx.num_q_tiles, pooled_k.size(2), ctx.blocks_per_qtile, ctx.query.device + ) + pooled_score = pooled_score.masked_fill( + ~causal_mask.view(1, 1, ctx.num_q_tiles, pooled_k.size(2)), + -torch.inf, + ) + else: + causal_mask = None + pooled_prob = _triton_safe_softmax(pooled_score) + sorted_prob = torch.sort(pooled_prob, dim=-1, descending=True) + return pooled_prob, sorted_prob, sim_k_expand, sim_q_expand, causal_mask + + ( + pooled_prob, + sorted_prob, + sim_k_expand, + sim_q_expand, + causal_mask, + ), stage_ms[ + "pooled_score_softmax_sort" + ] = bench_with_output(routing_stage, warmup=0, iters=1) + + _, _, _, num_k_blocks = pooled_prob.shape + num_to_select = ( + (ctx.topk.view(1, ctx.num_heads_q, 1) * num_k_blocks) + .to(torch.int64) + .expand(ctx.batch, -1, ctx.num_q_tiles) + .contiguous() + ) + + def fill_stage(): + final_tile_map = torch.zeros_like(pooled_prob, dtype=torch.bool) + final_tile_map[~sim_k_expand] = True + final_tile_map[~sim_q_expand] = True + final_tile_map = _fill_block_map_triton(final_tile_map, num_to_select, sorted_prob.indices) + if causal_mask is not None: + final_tile_map &= causal_mask.view(1, 1, ctx.num_q_tiles, num_k_blocks) + return final_tile_map + + final_tile_map, stage_ms["fill_block_map"] = bench_with_output(fill_stage, warmup=0, iters=1) + + q_block_to_tile = ( + torch.arange((ctx.seq_len_q + ctx.quant_block_size - 1) // ctx.quant_block_size, device=ctx.query.device) + * ctx.quant_block_size + ) // ctx.query_tile_tokens + q_block_to_tile = q_block_to_tile.clamp_max(ctx.num_q_tiles - 1) + raw_block_map, stage_ms["tile_to_qblock_index"] = bench_with_output( + lambda: final_tile_map.index_select(2, q_block_to_tile).contiguous(), + warmup=0, + iters=1, + ) + + (lut, valid_block_num), stage_ms["block_map_to_lut"] = bench_with_output( + lambda: _block_map_lut_triton(raw_block_map), + warmup=0, + iters=1, + ) + + if _ >= warmup: + for name in stage_names: + sums[name] += stage_ms[name] + + last_meta = ark._finalize_sparge_preprocess_outputs( + ctx, + { + "query_i8": q_int8_hnd, + "key_i8": k_int8_hnd, + "qscale": q_scale, + "kscale": k_scale, + "raw_block_map": raw_block_map, + "tile_block_map": final_tile_map.contiguous(), + "sim_qblocks": sim_q_for_routing.contiguous(), + "sim_kblocks": sim_kblocks.contiguous(), + "backend": "torch", + }, + ) + del key_mean, pooled_q, sim_qblocks, q_int8_hnd, q_scale + del pooled_k, sim_kblocks, k_int8_hnd, k_scale + del pooled_q_for_routing, sim_q_for_routing, pooled_prob, sorted_prob + del sim_k_expand, sim_q_expand, final_tile_map, raw_block_map, lut, valid_block_num + if causal_mask is not None: + del causal_mask + + avg = {name: sums[name] / float(iters) for name in stage_names} + avg["preprocess_total"] = sum(avg.values()) + return avg, last_meta + + +def nhd_to_hnd(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 1, 3).contiguous() + + +def hnd_to_nhd(x: torch.Tensor) -> torch.Tensor: + return x.permute(0, 2, 1, 3).contiguous() + + +def summarize_speedups(rows: list[dict[str, object]]) -> list[dict[str, object]]: + torch_row = next((row for row in rows if row["mode"] == "dense_torch_sdpa" and row["status"] == "ok"), None) + sage_row = next((row for row in rows if row["mode"] == "dense_sagev1" and row["status"] == "ok"), None) + torch_ms = None if torch_row is None else float(torch_row["latency_ms"]) + sage_ms = None if sage_row is None else float(sage_row["latency_ms"]) + for row in rows: + latency_ms = row["latency_ms"] + if latency_ms is None: + row["speedup_vs_torch"] = None + row["speedup_vs_sagev1"] = None + row["baseline_tflops"] = None + row["effective_tflops"] = None + continue + row["speedup_vs_torch"] = (torch_ms / latency_ms) if torch_ms is not None else None + row["speedup_vs_sagev1"] = (sage_ms / latency_ms) if sage_ms is not None else None + batch = int(row["batch"]) + num_heads_q = int(row["num_heads_q"]) + seq_len = int(row["seq_len"]) + head_dim = int(row["head_dim"]) + dense_flops = attention_dense_flops(batch, num_heads_q, seq_len, seq_len, head_dim) + mode = str(row["mode"]) + row["baseline_tflops"] = None + row["effective_tflops"] = None + if mode in { + "dense_torch_sdpa", + "dense_sagev1", + "sparse_kernel_only", + "sparse_e2e", + "sparse_qtile256_row64k_kernel_only", + "sparse_qtile256_row64k_e2e", + }: + row["baseline_tflops"] = flops_to_tflops(dense_flops, float(latency_ms)) + work_ratio = 1.0 + if mode.startswith("sparse_") and row["selected_ratio"] is not None: + work_ratio = float(row["selected_ratio"]) + row["effective_tflops"] = flops_to_tflops(dense_flops * work_ratio, float(latency_ms)) + return rows + + +def print_summary(rows: list[dict[str, object]]) -> None: + print( + "| mode | topk | selected_ratio | blocks/row | latency (ms) | baseline_tflops | effective_tflops | status | speedup_vs_torch | speedup_vs_sagev1 |" + ) + print("|---|---|---|---|---|---|---|---|---|---|") + for row in rows: + topk = "-" if row["requested_topk"] is None else f"{float(row['requested_topk']):.3f}" + ratio = "-" if row["selected_ratio"] is None else f"{float(row['selected_ratio']):.6f}" + blocks = "-" if row["selected_blocks_per_row"] is None else f"{float(row['selected_blocks_per_row']):.3f}" + latency = "-" if row["latency_ms"] is None else f"{float(row['latency_ms']):.3f}" + baseline_tflops = "-" if row.get("baseline_tflops") is None else f"{float(row['baseline_tflops']):.3f}" + effective_tflops = "-" if row.get("effective_tflops") is None else f"{float(row['effective_tflops']):.3f}" + sp_torch = "-" if row.get("speedup_vs_torch") is None else f"{float(row['speedup_vs_torch']):.3f}" + sp_sage = "-" if row.get("speedup_vs_sagev1") is None else f"{float(row['speedup_vs_sagev1']):.3f}" + print( + f"| {row['mode']} | {topk} | {ratio} | {blocks} | {latency} | {baseline_tflops} | {effective_tflops} | {row['status']} | {sp_torch} | {sp_sage} |" + ) + if row["note"]: + print(f"note[{row['mode']}]: {row['note']}") + + +def write_csv(rows: list[dict[str, object]], output_csv: Path) -> None: + if not rows: + return + fieldnames = list(rows[0].keys()) + output_csv.parent.mkdir(parents=True, exist_ok=True) + with output_csv.open("w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def run_benchmark(args: argparse.Namespace) -> list[dict[str, object]]: + ensure_sparse_binding(args) + if not is_xpu_available(): + raise RuntimeError("XPU device is required") + + device = torch.device("xpu") + dtype = torch.float16 + scale = 1.0 / math.sqrt(args.head_dim) + enable_gqa = args.num_heads_q // args.num_heads_kv > 1 + q_hnd_src, k_hnd_src, v_hnd_src = build_inputs( + args.batch, + args.num_heads_q, + args.num_heads_kv, + args.seq_len, + args.head_dim, + dtype, + device, + ) + q_nhd_src = hnd_to_nhd(q_hnd_src) + k_nhd_src = hnd_to_nhd(k_hnd_src) + v_nhd_src = hnd_to_nhd(v_hnd_src) + if args.tensor_layout == "NHD": + q, k, v = q_nhd_src, k_nhd_src, v_nhd_src + q_hnd, k_hnd, v_hnd = q_hnd_src, k_hnd_src, v_hnd_src + else: + q, k, v = q_hnd_src, k_hnd_src, v_hnd_src + q_hnd, k_hnd, v_hnd = q_hnd_src, k_hnd_src, v_hnd_src + + rows: list[dict[str, object]] = [] + rows.append( + try_benchmark( + "dense_torch_sdpa", + lambda: ( + hnd_to_nhd( + F.scaled_dot_product_attention( + nhd_to_hnd(q_nhd_src), + nhd_to_hnd(k_nhd_src), + nhd_to_hnd(v_nhd_src), + dropout_p=0.0, + is_causal=args.causal, + scale=scale, + enable_gqa=enable_gqa, + ) + ) + if args.tensor_layout == "HND" + else F.scaled_dot_product_attention( + q_hnd, k_hnd, v_hnd, dropout_p=0.0, is_causal=args.causal, scale=scale, enable_gqa=enable_gqa + ) + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + ) + ) + rows.append( + try_benchmark( + "dense_sagev1", + lambda: ( + hnd_to_nhd( + ark.sagev1( + nhd_to_hnd(q_nhd_src), + nhd_to_hnd(k_nhd_src), + nhd_to_hnd(v_nhd_src), + scale=scale, + is_causal=args.causal, + quant_block_size=args.quant_block_size, + tensor_layout="HND", + ) + ) + if args.tensor_layout == "HND" + else ark.sagev1( + q, + k, + v, + scale=scale, + is_causal=args.causal, + quant_block_size=args.quant_block_size, + tensor_layout=args.tensor_layout, + ) + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + ) + ) + + for topk in args.topk: + preprocess = None + selected_ratio = None + selected_blocks_per_row = None + preprocess_note = "" + try: + preprocess = ark.sparge_preprocess_topk( + q, + k, + is_causal=args.causal, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + quant_block_size=args.quant_block_size, + tensor_layout=args.tensor_layout, + query_tile_tokens=args.q_tile_override or None, + sparse_q_block_tokens=args.sparse_q_block_tokens, + sparse_k_block_tokens=args.sparse_k_block_tokens, + ) + stats = preprocess.get("stats", {}) + selected_ratio = float(stats.get("selected_ratio", 0.0)) + selected_blocks_per_row = float(stats.get("selected_blocks_per_row", 0.0)) + except Exception as exc: + status, preprocess_note = classify_exception(exc) + empty_xpu_cache() + rows.append( + make_row( + mode="sparse_kernel_only", + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + requested_topk=topk, + selected_ratio=None, + selected_blocks_per_row=None, + latency_ms=None, + status=status, + note=f"preprocess failed before kernel benchmark: {preprocess_note}", + ) + ) + rows.append( + make_row( + mode="sparse_e2e", + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + requested_topk=topk, + selected_ratio=None, + selected_blocks_per_row=None, + latency_ms=None, + status=status, + note=f"preprocess failed before e2e benchmark: {preprocess_note}", + ) + ) + continue + + # try: + # preprocess_stage_ms, preprocess_profile_meta = benchmark_preprocess_stages( + # q, + # k, + # topk=topk, + # is_causal=args.causal, + # quant_block_size=args.quant_block_size, + # tensor_layout=args.tensor_layout, + # warmup=args.warmup, + # iters=args.iters, + # ) + # stage_selected_ratio = float(preprocess_profile_meta["stats"].get("selected_ratio", 0.0)) + # stage_selected_blocks_per_row = float(preprocess_profile_meta["stats"].get("selected_blocks_per_row", 0.0)) + # for stage_name, latency_ms in preprocess_stage_ms.items(): + # mode = stage_name if stage_name == "preprocess_total" else f"preprocess_{stage_name}" + # rows.append( + # make_row( + # mode=mode, + # batch=args.batch, + # num_heads_q=args.num_heads_q, + # num_heads_kv=args.num_heads_kv, + # seq_len=args.seq_len, + # head_dim=args.head_dim, + # dtype=dtype, + # is_causal=args.causal, + # warmup=args.warmup, + # iters=args.iters, + # requested_topk=topk, + # selected_ratio=stage_selected_ratio, + # selected_blocks_per_row=stage_selected_blocks_per_row, + # latency_ms=latency_ms, + # status="ok", + # ) + # ) + # except Exception as exc: + # status, note = classify_exception(exc) + # rows.append( + # make_row( + # mode="preprocess_total", + # batch=args.batch, + # num_heads_q=args.num_heads_q, + # num_heads_kv=args.num_heads_kv, + # seq_len=args.seq_len, + # head_dim=args.head_dim, + # dtype=dtype, + # is_causal=args.causal, + # warmup=args.warmup, + # iters=args.iters, + # requested_topk=topk, + # selected_ratio=selected_ratio, + # selected_blocks_per_row=selected_blocks_per_row, + # latency_ms=None, + # status=status, + # note=f"preprocess stage profiling failed: {note}", + # ) + # ) + + sparse_kernel_mode = ( + "sparse_qtile256_row64k_kernel_only" + if preprocess["sparse_q_block_tokens"] == 256 + and preprocess["sparse_k_block_tokens"] == 64 + and args.q_tile_override == 256 + else "sparse_kernel_only" + ) + rows.append( + try_benchmark( + sparse_kernel_mode, + lambda preprocess=preprocess: ( + hnd_to_nhd( + ark.sage_sparse( + preprocess["query_i8"], + preprocess["key_i8"], + nhd_to_hnd(v_nhd_src), + preprocess["lut"], + preprocess["valid_block_num"], + is_causal=args.causal, + scale=scale, + quant_block_size=preprocess["quant_block_size"], + qscale=preprocess["qscale"], + kscale=preprocess["kscale"], + q_tile_override=args.q_tile_override, + sparse_q_block_tokens=preprocess["sparse_q_block_tokens"], + sparse_k_block_tokens=preprocess["sparse_k_block_tokens"], + tensor_layout="HND", + ) + ) + if args.tensor_layout == "HND" + else ark.sage_sparse( + preprocess["query_i8"], + preprocess["key_i8"], + v, + preprocess["lut"], + preprocess["valid_block_num"], + is_causal=args.causal, + scale=scale, + quant_block_size=preprocess["quant_block_size"], + qscale=preprocess["qscale"], + kscale=preprocess["kscale"], + q_tile_override=args.q_tile_override, + sparse_q_block_tokens=preprocess["sparse_q_block_tokens"], + sparse_k_block_tokens=preprocess["sparse_k_block_tokens"], + tensor_layout=args.tensor_layout, + ) + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + requested_topk=topk, + selected_ratio=selected_ratio, + selected_blocks_per_row=selected_blocks_per_row, + ) + ) + sparse_e2e_mode = ( + "sparse_qtile256_row64k_e2e" + if preprocess["sparse_q_block_tokens"] == 256 + and preprocess["sparse_k_block_tokens"] == 64 + and args.q_tile_override == 256 + else "sparse_e2e" + ) + rows.append( + try_benchmark( + sparse_e2e_mode, + lambda topk=topk: ( + hnd_to_nhd( + ark.sparge_sage2_attn_meansim_topk_xpu( + nhd_to_hnd(q_nhd_src), + nhd_to_hnd(k_nhd_src), + nhd_to_hnd(v_nhd_src), + is_causal=args.causal, + scale=scale, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout="HND", + q_tile_override=args.q_tile_override, + sparse_q_block_tokens=args.sparse_q_block_tokens, + sparse_k_block_tokens=args.sparse_k_block_tokens, + ) + ) + if args.tensor_layout == "HND" + else ark.sparge_sage2_attn_meansim_topk_xpu( + q, + k, + v, + is_causal=args.causal, + scale=scale, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=args.tensor_layout, + q_tile_override=args.q_tile_override, + sparse_q_block_tokens=args.sparse_q_block_tokens, + sparse_k_block_tokens=args.sparse_k_block_tokens, + ) + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=args.seq_len, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + requested_topk=topk, + selected_ratio=selected_ratio, + selected_blocks_per_row=selected_blocks_per_row, + ) + ) + del preprocess + empty_xpu_cache() + + return summarize_speedups(rows) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark torch SDPA, sagev1, and sparse attention across top-k.") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--num-heads-q", type=int, default=40) + parser.add_argument("--num-heads-kv", type=int, default=40) + parser.add_argument("--seq-len", type=int, default=75600) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--topk", type=float, nargs="+", default=[1.0, 0.75, 0.5, 0.25, 0.125]) + parser.add_argument("--quant-block-size", type=int, default=64) + parser.add_argument( + "--q-tile-override", + type=int, + default=0, + choices=(0, 64, 256), + help="Sparse kernel q_tile override. 0 keeps the default kernel choice.", + ) + parser.add_argument( + "--sparse-q-block-tokens", + type=int, + default=None, + help="Optional sparse Q-row granularity in tokens. Use 256 with --q-tile-override 256 for the decoupled path.", + ) + parser.add_argument( + "--sparse-k-block-tokens", + type=int, + default=None, + help="Optional sparse K logical-block granularity in tokens. Use 64 for the decoupled qtile256 path.", + ) + parser.add_argument("--tensor-layout", choices=("HND", "NHD"), default="HND") + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iters", type=int, default=3) + parser.add_argument( + "--causal", action="store_true", help="Run causal attention instead of the default non-causal mode." + ) + parser.add_argument( + "--xbuild-dir", + type=Path, + default=None, + help="Directory containing the benchmarked auto_round_kernel_xpu*.so build artifact.", + ) + parser.add_argument( + "--xpu-so", + type=Path, + default=None, + help="Explicit path to the auto_round_kernel_xpu*.so artifact to benchmark.", + ) + parser.add_argument( + "--output-csv", + type=Path, + default=Path("bench_sparse_topk_results.csv"), + help="Where to write the CSV result table.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + rows = run_benchmark(args) + write_csv(rows, args.output_csv) + print_summary(rows) + print(f"wrote csv: {args.output_csv}") + + +if __name__ == "__main__": + main() diff --git a/auto_round_extension/ark/test/test_sparge_preprocess_helpers.py b/auto_round_extension/ark/test/test_sparge_preprocess_helpers.py new file mode 100644 index 000000000..c4a27e08e --- /dev/null +++ b/auto_round_extension/ark/test/test_sparge_preprocess_helpers.py @@ -0,0 +1,172 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from auto_round_kernel.sparge_preprocess_triton import _fill_block_map_triton +from auto_round_kernel.sparse_attention import ( + _build_block_causal_mask, + _build_sparge_preprocess_context, + _fill_block_map_torch, +) + + +def test_fill_block_map_torch_skips_preselected_entries(): + final_map = torch.tensor([[[[True, False, True, False, False]]]], dtype=torch.bool) + num_to_select = torch.tensor([[[2]]], dtype=torch.int64) + sorted_indices = torch.tensor([[[[0, 2, 4, 1, 3]]]], dtype=torch.int64) + + filled = _fill_block_map_torch(final_map, num_to_select, sorted_indices) + + expected = torch.tensor([[[[True, True, True, False, True]]]], dtype=torch.bool) + assert torch.equal(filled, expected) + + +def test_fill_block_map_torch_selects_one_new_entry_when_topk_is_zero(): + final_map = torch.tensor([[[[True, False, False, False]]]], dtype=torch.bool) + num_to_select = torch.tensor([[[0]]], dtype=torch.int64) + sorted_indices = torch.tensor([[[[0, 2, 1, 3]]]], dtype=torch.int64) + + filled = _fill_block_map_torch(final_map, num_to_select, sorted_indices) + + expected = torch.tensor([[[[True, False, True, False]]]], dtype=torch.bool) + assert torch.equal(filled, expected) + + +@pytest.mark.skipif( + not (hasattr(torch, "xpu") and torch.xpu.is_available()), + reason="XPU not available", +) +def test_fill_block_map_triton_matches_torch_when_rows_have_preselected_entries(): + final_map = torch.tensor([[[[True, False, True, False, False]]]], dtype=torch.bool, device="xpu") + num_to_select = torch.tensor([[[2]]], dtype=torch.int64, device="xpu") + sorted_indices = torch.tensor([[[[0, 2, 4, 1, 3]]]], dtype=torch.int64, device="xpu") + + triton_filled = _fill_block_map_triton(final_map, num_to_select, sorted_indices) + torch_filled = _fill_block_map_torch(final_map.cpu(), num_to_select.cpu(), sorted_indices.cpu()).to("xpu") + torch.xpu.synchronize() + + assert torch.equal(triton_filled, torch_filled) + + +def test_build_block_causal_mask_supports_q64_k128_geometry(): + mask = _build_block_causal_mask( + num_q_tiles=4, + num_k_blocks=4, + q_route_block_tokens=64, + k_route_block_tokens=128, + device=torch.device("cpu"), + ) + + expected = torch.tensor( + [ + [True, False, False, False], + [True, False, False, False], + [True, True, False, False], + [True, True, False, False], + ], + dtype=torch.bool, + ) + assert torch.equal(mask, expected) + + +@pytest.mark.skipif( + not (hasattr(torch, "xpu") and torch.xpu.is_available()), + reason="XPU not available", +) +def test_preprocess_context_uses_wan_head_dim_128_routing_geometry(): + q = torch.randn((1, 2, 256, 128), device="xpu", dtype=torch.float16) + k = torch.randn((1, 2, 256, 128), device="xpu", dtype=torch.float16) + + ctx = _build_sparge_preprocess_context( + q, + k, + is_causal=False, + smooth_k=True, + simthreshd1=-0.1, + topk=0.5, + attention_sink=False, + quant_block_size=64, + tensor_layout="HND", + ) + + assert ctx.q_route_block_tokens == 64 + assert ctx.k_route_block_tokens == 128 + assert ctx.q_blocks_per_tile == 1 + assert ctx.k_blocks_per_tile == 2 + assert ctx.query_tile_tokens == 64 + + +@pytest.mark.skipif( + not (hasattr(torch, "xpu") and torch.xpu.is_available()), + reason="XPU not available", +) +def test_preprocess_context_accepts_explicit_query_tile_tokens_256(): + q = torch.randn((1, 2, 512, 128), device="xpu", dtype=torch.float16) + k = torch.randn((1, 2, 512, 128), device="xpu", dtype=torch.float16) + + ctx = _build_sparge_preprocess_context( + q, + k, + is_causal=False, + smooth_k=True, + simthreshd1=-0.1, + topk=0.5, + attention_sink=False, + quant_block_size=64, + tensor_layout="HND", + query_tile_tokens=256, + ) + + assert ctx.q_route_block_tokens == 256 + assert ctx.q_blocks_per_tile == 4 + assert ctx.num_q_tiles == 2 + assert ctx.query_tile_tokens == 256 + + +@pytest.mark.skipif( + not (hasattr(torch, "xpu") and torch.xpu.is_available()), + reason="XPU not available", +) +@pytest.mark.parametrize("tensor_layout", ("HND", "NHD")) +def test_preprocess_context_accepts_decoupled_sparse_qtile256_row64k(tensor_layout: str): + if tensor_layout == "HND": + q = torch.randn((1, 2, 512, 128), device="xpu", dtype=torch.float16) + k = torch.randn((1, 2, 512, 128), device="xpu", dtype=torch.float16) + else: + q = torch.randn((1, 512, 2, 128), device="xpu", dtype=torch.float16) + k = torch.randn((1, 512, 2, 128), device="xpu", dtype=torch.float16) + + ctx = _build_sparge_preprocess_context( + q, + k, + is_causal=False, + smooth_k=True, + simthreshd1=-0.1, + topk=0.5, + attention_sink=False, + quant_block_size=64, + tensor_layout=tensor_layout, + query_tile_tokens=256, + sparse_q_block_tokens=256, + sparse_k_block_tokens=64, + ) + + assert ctx.q_route_block_tokens == 256 + assert ctx.k_route_block_tokens == 64 + assert ctx.sparse_q_block_tokens == 256 + assert ctx.sparse_k_block_tokens == 64 + assert ctx.q_blocks_per_tile == 4 + assert ctx.k_blocks_per_tile == 1 + assert ctx.q_sparse_blocks_per_tile == 1 + assert ctx.k_sparse_blocks_per_tile == 1 + assert ctx.num_q_blocks == 8 + assert ctx.num_k_blocks == 8 + assert ctx.num_sparse_q_blocks == 2 + assert ctx.num_sparse_k_blocks == 8 diff --git a/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4.png b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4.png new file mode 100644 index 000000000..f3bb4c2fb Binary files /dev/null and b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4.png differ diff --git a/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels.png b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels.png new file mode 100644 index 000000000..5c6b0a6a8 Binary files /dev/null and b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels.png differ diff --git a/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels_with_0p05.png b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels_with_0p05.png new file mode 100644 index 000000000..041e1e166 Binary files /dev/null and b/docs/pr_artifacts/flux_sweep_defaultsteps_20260707T122523Z_grid_3x4_large_labels_with_0p05.png differ diff --git a/docs/pr_artifacts/flux_topk0p05.png b/docs/pr_artifacts/flux_topk0p05.png new file mode 100644 index 000000000..5727fd809 Binary files /dev/null and b/docs/pr_artifacts/flux_topk0p05.png differ