diff --git a/auto_round_extension/ark/.gitignore b/auto_round_extension/ark/.gitignore index d68455ecb5..8e16b488bb 100644 --- a/auto_round_extension/ark/.gitignore +++ b/auto_round_extension/ark/.gitignore @@ -1,5 +1,9 @@ build xbuild +ark-xbuild* +ark-xbuild-builtin* +xbuild_bf16* *.csv *.so *.pyc +benchmarks/results/ diff --git a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt index 0d915d7f3d..df38ae3795 100755 --- a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt +++ b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt @@ -182,6 +182,18 @@ 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) + # sycl-tla may provision oneMKL through an ExternalProject. Wire its staged + # headers into this module directly so parallel builds do not compile before + # oneMKL has been downloaded and installed. + if(TARGET onemkl_project) + add_dependencies(${PY_NAME} onemkl_project) + target_include_directories(${PY_NAME} SYSTEM PRIVATE ${CMAKE_BINARY_DIR}/deps/oneMKL/include) + elseif(TARGET MKL::MKL) + target_link_libraries(${PY_NAME} PRIVATE MKL::MKL) + if(DEFINED ENV{MKLROOT} AND EXISTS "$ENV{MKLROOT}/include") + target_include_directories(${PY_NAME} SYSTEM PRIVATE "$ENV{MKLROOT}/include") + endif() + endif() # 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) @@ -208,6 +220,15 @@ 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) + if(TARGET onemkl_project) + add_dependencies(${TEST_NAME} onemkl_project) + target_include_directories(${TEST_NAME} SYSTEM PRIVATE ${CMAKE_BINARY_DIR}/deps/oneMKL/include) + elseif(TARGET MKL::MKL) + target_link_libraries(${TEST_NAME} PRIVATE MKL::MKL) + if(DEFINED ENV{MKLROOT} AND EXISTS "$ENV{MKLROOT}/include") + target_include_directories(${TEST_NAME} SYSTEM PRIVATE "$ENV{MKLROOT}/include") + endif() + endif() 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 8d4e862edb..579212d19b 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -2146,10 +2146,12 @@ def sageattn( _sequence_mean_native_layout, _slice_sequence_native_layout, _to_hnd, + block_sparse_sdpa, sage_sparse, sparge_block_map_to_mask, sparge_preprocess_topk, sparge_sage2_attn_meansim_topk_xpu, + sparge_sage2_attn_meansim_topk_xpu_sdpa, ) diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index aa843020d5..feca265346 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -507,6 +507,98 @@ static void sage_sparse(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, throw std::invalid_argument("ark::sage_sparse: unsupported head_dim; supported values are 64 and 128"); } +static void block_sparse_sdpa(torch_ptr stream, torch_ptr Q, torch_ptr K, torch_ptr V, torch_ptr O, torch_ptr mask, + 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 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_dtype != (int)FlashAttnDtype::FP16 && q_dtype != (int)FlashAttnDtype::BF16) { + throw std::invalid_argument("ark::block_sparse_sdpa: q_dtype must be FP16 or BF16"); + } + if (mask && is_causal) { + throw std::invalid_argument("ark::block_sparse_sdpa: mask and is_causal cannot both be set"); + } + if (!lut || !valid_block_num) { + throw std::invalid_argument("ark::block_sparse_sdpa: lut and valid_block_num must be provided"); + } + auto matches_block_size = [](int seq_len, int num_blocks, int block_size) { + return block_size > 0 && num_blocks == ((seq_len + block_size - 1) / block_size); + }; + const bool key_block_is_64 = matches_block_size(seq_len_kv, num_k_blocks, 64); + if (!key_block_is_64) { + throw std::invalid_argument("ark::block_sparse_sdpa: only key block size 64 is supported"); + } + const bool is_bf16 = (q_dtype == (int)FlashAttnDtype::BF16); + if (head_dim == 64) { + if (!matches_block_size(seq_len_q, num_q_blocks, 64)) { + throw std::invalid_argument("ark::block_sparse_sdpa: head_dim=64 requires query block size 64"); + } + if (is_bf16) { + ark::sdpa_impl_bf16_sparse_sdpa_d64( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } else { + ark::sdpa_impl_fp16_sparse_sdpa_d64( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } + return; + } + if (head_dim == 128) { + if (matches_block_size(seq_len_q, num_q_blocks, 256)) { + if (q_tile_override != 256) { + throw std::invalid_argument( + "ark::block_sparse_sdpa: head_dim=128 query block size 256 requires q_tile_override=256"); + } + if (is_bf16) { + ark::sdpa_impl_bf16_sparse_sdpa_qtile256_row64k( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } else { + ark::sdpa_impl_fp16_sparse_sdpa_qtile256_row64k( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } + return; + } + if (matches_block_size(seq_len_q, num_q_blocks, 64)) { + if (is_bf16) { + ark::sdpa_impl_bf16_sparse_sdpa_row_linear( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } else { + ark::sdpa_impl_fp16_sparse_sdpa_row_linear( + (sycl::queue*)stream, (void*)Q, (void*)K, (void*)V, (void*)O, (void*)mask, (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); + } + return; + } + throw std::invalid_argument("ark::block_sparse_sdpa: head_dim=128 supports query block sizes 64 and 256 only"); + } + throw std::invalid_argument("ark::block_sparse_sdpa: unsupported head_dim; supported values are 64 and 128"); +} + 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) { @@ -1379,6 +1471,7 @@ PYBIND11_MODULE(PY_NAME, m) { 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("block_sparse_sdpa", &ark::block_sparse_sdpa); // 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"), diff --git a/auto_round_extension/ark/auto_round_kernel/sdpa_sparse_sdpa.cpp b/auto_round_extension/ark/auto_round_kernel/sdpa_sparse_sdpa.cpp new file mode 100644 index 0000000000..b569a49f2a --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/sdpa_sparse_sdpa.cpp @@ -0,0 +1,330 @@ +// 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. + +// Independent native-precision sparse SDPA path (BF16 + FP16). This dispatches to +// the sparse SDPA mainloop (SPARSESDPAFwdMainloop via SparseSDPAConfig), separate +// from the INT8-centric sparse SAGE path in sdpa_sparse.cpp. + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +#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_bf16_128_sparse_sdpa(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_128(options); +} + +int launch_prefill_kernel_bf16_128_sparse_sdpa_qtile64(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_128_qtile64< + cute::bfloat16_t, cute::bfloat16_t, cute::bfloat16_t>(options); +} + +int launch_prefill_kernel_bf16_64_sparse_sdpa(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_64(options); +} + +int launch_prefill_kernel_f16_128_sparse_sdpa(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_128(options); +} + +int launch_prefill_kernel_f16_128_sparse_sdpa_qtile64(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_128_qtile64(options); +} + +int launch_prefill_kernel_f16_64_sparse_sdpa(detail::Options const& options) { + return detail::launch_sparse_sdpa_prefill_kernel_64(options); +} + +KernelLauncher select_sparse_sdpa_prefill_launcher(BTLA_DTYPE q_dtype, int head_dim, int q_tile_override) { + switch (head_dim) { + case 128: + if (q_tile_override == 64) { + return q_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_128_sparse_sdpa_qtile64 + : launch_prefill_kernel_f16_128_sparse_sdpa_qtile64; + } + if (q_tile_override != 0 && q_tile_override != 256) return nullptr; + return q_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_128_sparse_sdpa + : launch_prefill_kernel_f16_128_sparse_sdpa; + case 64: + if (q_tile_override != 0 && q_tile_override != 64 && q_tile_override != 128) return nullptr; + return q_dtype == BTLA_DTYPE::BF16 ? launch_prefill_kernel_bf16_64_sparse_sdpa + : launch_prefill_kernel_f16_64_sparse_sdpa; + 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; +} + +void sparse_sdpa_prefill(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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 q_dtype, + int sparse_q_block_size = 0) { + const int effective_q_tile_override = + (head_dim == 128 && q_tile_override == 0) ? 64 : ((head_dim == 64 && 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); + // Native-precision path: no INT8 dequant scales, softmax scale applied directly. + options.scale_block_size = 0; + options.sparse_q_block_size = sparse_q_block_size; + options.q_tile_override = effective_q_tile_override; + options.qscale = nullptr; + options.kscale = nullptr; + options.vscale = nullptr; + 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_sdpa_prefill_launcher(q_dtype, head_dim, effective_q_tile_override); + if (launcher == nullptr) { + throw std::runtime_error("Unsupported sparse_sdpa_prefill config"); + } + + launcher(options); +} + +} // namespace + +void sdpa_impl_bf16_sparse_sdpa_d64( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_d64: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_d64: seq_len_q and seq_len_kv must be greater than 0"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_d64: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_d64: num_q_blocks and num_k_blocks must be greater than 0"); + } + if (head_dim != 64) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_d64: head_dim must be 64"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, 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::BF16, /*sparse_q_block_size=*/64); +} + +void sdpa_impl_fp16_sparse_sdpa_d64( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_d64: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_d64: seq_len_q and seq_len_kv must be greater than 0"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_d64: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_d64: num_q_blocks and num_k_blocks must be greater than 0"); + } + if (head_dim != 64) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_d64: head_dim must be 64"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, 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::F16, /*sparse_q_block_size=*/64); +} + +void sdpa_impl_bf16_sparse_sdpa_row_linear( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (q_tile_override != 0 && q_tile_override != 64) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_row_linear: q_tile_override must be 0 or 64"); + } + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_row_linear: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument( + "sdpa_impl_bf16_sparse_sdpa_row_linear: seq_len_q and seq_len_kv must be greater than 0"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_row_linear: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument( + "sdpa_impl_bf16_sparse_sdpa_row_linear: num_q_blocks and num_k_blocks must be greater than 0"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, valid_block_num, num_q_blocks, num_k_blocks, 64, + 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::BF16, /*sparse_q_block_size=*/64); +} + +void sdpa_impl_fp16_sparse_sdpa_row_linear( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (q_tile_override != 0 && q_tile_override != 64) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_row_linear: q_tile_override must be 0 or 64"); + } + if (mask && is_causal) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_row_linear: mask and is_causal cannot both be set"); + } + if (seq_len_q <= 0 || seq_len_kv <= 0) { + throw std::invalid_argument( + "sdpa_impl_fp16_sparse_sdpa_row_linear: seq_len_q and seq_len_kv must be greater than 0"); + } + if (lut == nullptr || valid_block_num == nullptr) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_row_linear: lut and valid_block_num must be provided"); + } + if (num_q_blocks <= 0 || num_k_blocks <= 0) { + throw std::invalid_argument( + "sdpa_impl_fp16_sparse_sdpa_row_linear: num_q_blocks and num_k_blocks must be greater than 0"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, valid_block_num, num_q_blocks, num_k_blocks, 64, + 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::F16, /*sparse_q_block_size=*/64); +} + +void sdpa_impl_bf16_sparse_sdpa_qtile256_row64k( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (q_tile_override != 256) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_qtile256_row64k: q_tile_override must be 256"); + } + if (head_dim != 128) { + throw std::invalid_argument("sdpa_impl_bf16_sparse_sdpa_qtile256_row64k: head_dim must be 128"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, 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::BF16, /*sparse_q_block_size=*/256); +} + +void sdpa_impl_fp16_sparse_sdpa_qtile256_row64k( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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) { + if (q_tile_override != 256) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_qtile256_row64k: q_tile_override must be 256"); + } + if (head_dim != 128) { + throw std::invalid_argument("sdpa_impl_fp16_sparse_sdpa_qtile256_row64k: head_dim must be 128"); + } + + sparse_sdpa_prefill(q, Q_ptr, K_ptr, V_ptr, O_ptr, mask, lut, 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::F16, /*sparse_q_block_size=*/256); +} + +} // 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 index d94263c47d..2fe1091f80 100644 --- a/auto_round_extension/ark/auto_round_kernel/sparge_preprocess_triton.py +++ b/auto_round_extension/ark/auto_round_kernel/sparge_preprocess_triton.py @@ -16,7 +16,11 @@ def _apply_xpu_triton_workarounds() -> None: - """Patch triton's Intel compiler so JIT kernels avoid the rejected SPIR-V extension.""" + """Patch triton's Intel compiler so JIT kernels avoid the rejected SPIR-V extension. + + Triton 3.7.x can emit SPV_INTEL_predicated_io for non-LTS Intel GPU drivers, + but some level-zero loaders reject that extension at kernel-load time. + """ global _APPLIED_TRITON_PREDICATED_CHECK if _APPLIED_TRITON_PREDICATED_CHECK: return @@ -416,8 +420,9 @@ def _run_triton_xpu_preprocess(ctx: Any) -> dict[str, Any]: .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 + # Boolean-mask assignment (`t[mask] = val`) is unreliable on XPU; use |= instead. + final_tile_map |= ~sim_k_expand + final_tile_map |= ~sim_q_expand 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) diff --git a/auto_round_extension/ark/auto_round_kernel/sparse_attention.py b/auto_round_extension/ark/auto_round_kernel/sparse_attention.py index 0d5932747c..e45febf887 100644 --- a/auto_round_extension/ark/auto_round_kernel/sparse_attention.py +++ b/auto_round_extension/ark/auto_round_kernel/sparse_attention.py @@ -199,6 +199,122 @@ def sage_sparse( return O +def block_sparse_sdpa( + 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, + 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 native-precision sparse attention (BF16/FP16) routed by sparse LUT metadata. + + Independent path built on the dense-SDPA sparse mainloop (SparseSDPAConfig). + Unlike the INT8 ``sage_sparse`` path it does not reuse the INT8-centric SAGE + mainloop; Q/K/V are native BF16/FP16 and the softmax scale is applied directly. + """ + del dropout_p, enable_gqa + if query.device.type != "xpu": + raise NotImplementedError("block_sparse_sdpa is only supported on XPU") + if query.dtype not in (torch.bfloat16, torch.float16): + raise ValueError(f"Q/K/V must be bfloat16 or float16, got {query.dtype}") + if key.dtype != query.dtype or value.dtype != query.dtype: + raise ValueError(f"Q/K/V dtypes must match, got Q={query.dtype}, K={key.dtype}, V={value.dtype}") + 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") + + 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") + _validate_gqa_head_config(Hq, Hkv, op_name="block_sparse_sdpa") + if D not in (64, 128): + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + effective_sparse_q_block_tokens = 64 if sparse_q_block_tokens is None else int(sparse_q_block_tokens) + effective_sparse_k_block_tokens = 64 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_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 + 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 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 D == 64 and q_tile_override not in (0, 64, 128): + raise ValueError( + f"q_tile_override must be one of {{0, 64, 128}} for block_sparse_sdpa with head_dim=64, got {q_tile_override}" + ) + + lib = get_lib(query) + if not hasattr(lib, "block_sparse_sdpa"): + raise RuntimeError("Loaded XPU extension does not expose block_sparse_sdpa") + effective_q_tile_override = _resolve_sparse_prefill_q_tile_override( + head_dim=D, + quant_block_size=64, + 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, + ) + stream = get_stream(query) + O = _empty_attention_output(B, Hq, Sq, D, dtype=value.dtype, device=query.device, tensor_layout=tensor_layout) + q_dtype = 1 if query.dtype == torch.bfloat16 else 0 # FlashAttnDtype: 0=FP16, 1=BF16 + lib.block_sparse_sdpa( + 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, + lut.data_ptr(), + valid_block_num.data_ptr(), + q_sparse_blocks, + kv_sparse_blocks, + effective_q_tile_override, + *_attention_strides_qko(query, tensor_layout), + *_attention_strides_qko(key, tensor_layout), + *_attention_strides_v(value, tensor_layout), + *_attention_strides_qko(O, tensor_layout), + q_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, @@ -399,8 +515,10 @@ def _normalize_per_head_hparam( def _query_tile_tokens_for_head_dim(head_dim: int) -> int: + # Keep the preprocess default aligned with the sparse e2e wrappers: for head_dim 64 + # the wrappers default to 64-token query tiles, for head_dim 128 to 64 as well. if head_dim == 64: - return 128 + return 64 if head_dim == 128: return 64 raise ValueError(f"Unsupported head_dim={head_dim}; supported: 64, 128") @@ -616,19 +734,25 @@ def _build_block_causal_mask( def _fill_block_map_torch( final_map: torch.Tensor, num_to_select: torch.Tensor, sorted_indices: torch.Tensor ) -> torch.Tensor: + """Fill ``final_map`` with the top ``num_to_select`` NEW blocks in sorted order. + + The naive implementation looped over every K block (``for rank in range(k_blocks)``), + launching a handful of elementwise kernels per iteration. At long sequences that is + thousands of sequential kernel launches and dominates the torch preprocess. This + version is fully vectorized: for a block at sorted ``rank r`` it is the + ``(r - #already_filled_before_r)``-th *new* block, and it is selected iff that index + is below the per-row quota. + """ 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 + already = torch.gather(final_map, -1, sorted_indices) # is the block at this rank already filled? + already_cum = torch.cumsum(already.to(torch.int64), dim=-1) - already.to(torch.int64) + rank = torch.arange(k_blocks, device=final_map.device).view(1, 1, 1, k_blocks) + new_idx = rank - already_cum # 0-based index among new blocks at this rank + add_new = (new_idx < target_new.unsqueeze(-1)) & ~already + chosen = torch.zeros_like(final_map) + chosen.scatter_(-1, sorted_indices, add_new) + return final_map | chosen def _block_map_lut_torch(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -637,8 +761,11 @@ def _block_map_lut_torch(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.T 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 + # NOTE: do not use `filled_matrix[~block_map] = 10_000_000` here — boolean-mask + # in-place assignment silently misses entries on the XPU backend (torch-xpu-ops), + # leaving 0s in the mask that become -1 in the LUT after sort and drive the + # ScatterGatherKernels out-of-bounds assert downstream. Use torch.where instead. + filled_matrix = torch.where(block_map, masked_cum_matrix, torch.full_like(masked_cum_matrix, 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( @@ -648,6 +775,27 @@ def _block_map_lut_torch(block_map: torch.Tensor) -> tuple[torch.Tensor, torch.T return lut.to(torch.int32).contiguous(), valid_block_num.to(torch.int32).contiguous() +def _lut_to_block_map(lut: torch.Tensor, valid_block_num: torch.Tensor) -> torch.Tensor: + 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.shape[:-1] != valid_block_num.shape: + raise ValueError( + f"lut and valid_block_num shape mismatch: expected lut.shape[:-1] == valid_block_num.shape, " + f"got {tuple(lut.shape)} and {tuple(valid_block_num.shape)}" + ) + # Run-length encoded LUT: cumsum recovers the 0-indexed selected block positions. + logical_blocks = torch.cumsum(lut.to(torch.int64), dim=-1) + # Scatter True at every recovered position. The invalid tail of `logical_blocks` + # plateaus at the last selected position, so those duplicate indices do not add + # spurious blocks. Do NOT scatter the tail with False (e.g. `scatter_(-1, idx, mask)`): + # masked-out positions all index slot 0, and their trailing False writes clobber a + # selected block at position 0, silently dropping it. Empty rows are cleared below. + block_map = torch.zeros_like(lut, dtype=torch.bool) + block_map.scatter_(-1, logical_blocks, True) + block_map &= valid_block_num.unsqueeze(-1) > 0 + return block_map.contiguous() + + def _prefix_keep_cross_attn_enabled() -> bool: return os.getenv("SPARGE_KEEP_PREFIX_ON_CROSS_ATTN", "0").strip().lower() in {"1", "true", "yes"} @@ -1088,8 +1236,8 @@ def _sparge_preprocess_topk_torch_impl(ctx: _SpargePreprocessContext) -> dict[st .contiguous() ) tail_map = torch.zeros_like(tail_prob, dtype=torch.bool) - tail_map[~tail_sim_k_expand] = True - tail_map[~tail_sim_q_expand] = True + tail_map |= ~tail_sim_k_expand + tail_map |= ~tail_sim_q_expand tail_map = _fill_block_map_torch(tail_map, num_to_select_tail, tail_sorted_prob.indices) if ctx.cdfthreshd is not None: prefix_mass = pooled_prob[..., :prefix_route_blocks].sum(dim=-1, keepdim=True) @@ -1111,8 +1259,9 @@ def _sparge_preprocess_topk_torch_impl(ctx: _SpargePreprocessContext) -> dict[st .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 + # Boolean-mask assignment (`t[mask] = val`) is unreliable on XPU; use |= instead. + final_tile_map |= ~sim_k_expand + final_tile_map |= ~sim_q_expand final_tile_map = _fill_block_map_torch(final_tile_map, num_to_select, sorted_prob.indices) if ctx.cdfthreshd is not None: final_tile_map |= _select_blocks_for_cdf( @@ -1383,17 +1532,31 @@ def sparge_sage2_attn_meansim_topk_xpu( effective_query_tile_tokens = query_tile_tokens effective_q_tile_override = q_tile_override + if D == 64 and effective_query_tile_tokens is None and q_tile_override == 0: + effective_query_tile_tokens = 64 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) + effective_q_tile_override = ( + 0 if D == 64 and int(effective_query_tile_tokens) == 64 else 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) + if D == 64: + supported_q_tiles = (0, 64, 128) + elif D == 128: + supported_q_tiles = (0, 64, 256) + else: + raise ValueError(f"Unsupported head_dim={D}; supported: 64, 128") + if effective_q_tile_override not in supported_q_tiles: + raise ValueError( + f"q_tile_override={effective_q_tile_override} is not supported for BF16 sparse head_dim={D}; " + f"supported values: {', '.join(str(v) for v in supported_q_tiles)}" + ) normalized_mask = _normalize_sparse_mask(attn_mask, B, Sq, Skv, query.device) metadata = sparge_preprocess_topk( @@ -1438,3 +1601,112 @@ def sparge_sage2_attn_meansim_topk_xpu( if return_sparsity: return out, sparsity_ratio return out + + +def sparge_sage2_attn_meansim_topk_xpu_sdpa( + 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, ...]: + """E2E sparse attention on the independent native-precision (BF16/FP16) SDPA path.""" + if query.device.type != "xpu": + raise NotImplementedError("sparge_sage2_attn_meansim_topk_xpu_sdpa is only supported on XPU") + if dropout_p != 0.0: + raise NotImplementedError("dropout_p must be 0.0 for sparge_sage2_attn_meansim_topk_xpu_sdpa") + if attn_mask is not None and is_causal: + raise ValueError("attn_mask and is_causal cannot both be set") + if query.dtype not in (torch.bfloat16, torch.float16): + raise ValueError( + "sparge_sage2_attn_meansim_topk_xpu_sdpa requires bfloat16 or float16 Q/K/V, " + f"got Q={query.dtype}, K={key.dtype}, V={value.dtype}" + ) + if key.dtype != query.dtype or value.dtype != query.dtype: + raise ValueError(f"Q/K/V dtypes must match, got Q={query.dtype}, K={key.dtype}, V={value.dtype}") + if output_dtype is not None and output_dtype != query.dtype: + raise ValueError(f"output_dtype must match Q dtype in the current implementation, got {output_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") + + 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, + cdfthreshd=cdfthreshd, + 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 = block_sparse_sdpa( + query, + key, + value, + metadata["lut"], + metadata["valid_block_num"], + attn_mask=normalized_mask, + is_causal=is_causal, + scale=scale, + 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_sparse_sdpa_fwd_mainloop.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sdpa_fwd_mainloop.hpp new file mode 100644 index 0000000000..c1b2d1fae2 --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sparse_sdpa_fwd_mainloop.hpp @@ -0,0 +1,636 @@ +/*************************************************************************************************** + * Copyright (C) 2026 Intel Corporation, All rights reserved. + * SPDX-License-Identifier: BSD-3-Clause + **************************************************************************************************/ + +#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 + +namespace cutlass::sdpa { + +template +class XeDefault {}; + +} // namespace cutlass::sdpa + +namespace cutlass::fmha::collective { + +using namespace cute; + +template +struct SPARSESDPAFwdMainloop { + static_assert(cutlass::detail::dependent_false, "Could not find a mainloop specialization."); +}; + +template +struct SPARSESDPAFwdMainloop, CausalMask_, FullMask_, CachedKV_, PagedKV_, + TiledMMAQK_, TiledMMAPV_, VTiles_, TensorQ_, TensorK_, TensorV_, TensorK_cache_, + TensorV_cache_, TiledCopyQ_, TiledCopyK_, TiledCopyV_, TiledCopyK_cache_, + TiledCopyV_cache_> { + 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 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_>; + + 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 FragSCol = decltype(reduce<0>(FragS{}, sycl::plus{})); + using ElementS = typename TiledMMAQK::ValTypeD; + using ElementM = typename TiledMMAQK::ValTypeA; + + using SingleFragA = FragC; + using FragA = expand_sg_fragment_t; + using FragARow = decltype(reduce<1>(FragA{}, sycl::plus{})); + using ElementA = typename TiledMMAPV::ValTypeD; + + static constexpr bool CausalMask = CausalMask_; + static constexpr bool CachedKV = CachedKV_; + static constexpr bool PagedKV = PagedKV_; + + struct Arguments { + float const scale; + float const* mask = nullptr; + int scale_block_size = 0; + float const* qscale = nullptr; + float const* kscale = nullptr; + float const* vscale = nullptr; + int const* lut = nullptr; + int const* valid_block_num = nullptr; + int num_q_blocks = 0; + int num_k_blocks = 0; + 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; + }; + + using Params = Arguments; + struct EmptySharedStorage {}; + using SharedStorage = EmptySharedStorage; + + Params params; + SharedStorage& shared_storage; + + SPARSESDPAFwdMainloop(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; + 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; } + + 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; + 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, TensorK2D const& K_2D, TensorV2D const& V_2D, FragA& tArA, + FragARow& tA_max, FragARow& tA_sum, QVCoord blk_qv, int blk_k0, int blk_k1, + int total_blk, int thr_id, int seq_len, int seq_len_kv_cache, int l_coord, + [[maybe_unused]] float* scaleQ, [[maybe_unused]] float* scaleK, + [[maybe_unused]] 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; + + auto tile_shape_v = make_shape(get<1>(TileShapePV{}) * C{}, get<2>(TileShapePV{})); + + Tensor cQ = make_identity_tensor(Q_2D.shape()); + Tensor cK = make_identity_tensor(K_2D.shape()); + Tensor cV = make_identity_tensor(V_2D.shape()); + Tensor cK_cache = make_identity_tensor(K_cache_2D.shape()); + Tensor cV_cache = make_identity_tensor(V_cache_2D.shape()); + Tensor cP = make_identity_tensor(take<0, 2>(TileShapeQK{})); + + Tensor gQ = local_tile(cQ, TileShapeQK{}, append(blk_qv, _), Step<_1, X, _1>{}); + Tensor gK = local_tile(cK, TileShapeQK{}, make_coord(_, _, _), Step{}); + Tensor gV = local_tile(cV, tile_shape_v, make_coord(get<1>(blk_qv), _)); + Tensor gV_split = local_tile(gV, TileShapePV{}, make_coord(_, _, 0), Step{}); + + Tensor gK_cache = local_tile(cK_cache, TileShapeQK{}, make_coord(_, _, _), Step{}); + Tensor gV_cache = local_tile(cV_cache, tile_shape_v, make_coord(get<1>(blk_qv), _)); + Tensor gV_cache_split = local_tile(gV_cache, TileShapePV{}, make_coord(_, _, 0), Step{}); + + TiledCopyQ copy_q{Q_2D}; + TiledCopyK copy_k{K_2D}; + TiledCopyV copy_v{V_2D}; + TiledCopyK_cache copy_k_cache{K_cache_2D}; + TiledCopyV_cache copy_v_cache{V_cache_2D}; + + TiledMMAQK mma_qk{}; + TiledMMAPV mma_pv{}; + + auto thr_copy_q = copy_q.get_slice(thr_id); + auto thr_copy_k = copy_k.get_slice(thr_id); + auto thr_copy_v = copy_v.get_slice(thr_id); + auto thr_copy_k_cache = copy_k_cache.get_slice(thr_id); + auto thr_copy_v_cache = copy_v_cache.get_slice(thr_id); + auto thr_mma_qk = mma_qk.get_slice(thr_id); + auto thr_mma_pv = mma_pv.get_slice(thr_id); + + auto tQgQ = thr_copy_q.partition_S(gQ); + auto tKgK = thr_copy_k.partition_S(gK); + auto tVgV = thr_copy_v.partition_S(gV_split); + auto tKgK_cache = thr_copy_k_cache.partition_S(gK_cache); + auto tVgV_cache = thr_copy_v_cache.partition_S(gV_cache_split); + + 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); + auto tArP = thr_mma_pv.partition_sg_fragment_A(cP); + auto tVrV = thr_copy_v.partition_sg_fragment_D(gV_split(_, _, 0, 0)); + auto tArV = thr_mma_pv.partition_sg_fragment_B(gV_split(_, _, 0, 0)); + + auto prefetch_q = make_block_2d_prefetch(copy_q); + auto prefetch_k = make_block_2d_prefetch(copy_k); + auto prefetch_v = make_block_2d_prefetch(copy_v); + auto prefetch_k_cache = make_block_2d_prefetch(copy_k_cache); + auto prefetch_v_cache = make_block_2d_prefetch(copy_v_cache); + + auto pQgQ = prefetch_q.get_slice(thr_id).partition_S(gQ); + auto pKgK = prefetch_k.get_slice(thr_id).partition_S(gK); + 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); + auto pVgV_cache = prefetch_v_cache.get_slice(thr_id).partition_S(gV_cache_split); + + 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); + } + + 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; + int q_blocks_per_wg_tile = + sparse_q_block_size > 0 ? cute::max(1, int(get<0>(TileShapeQK{})) / sparse_q_block_size) : 1; + int sg_rows_per_sparse_q_block = + sparse_q_block_size > 0 ? cute::max(1, sparse_q_block_size / q_sg_tile) : 1; + 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); + + int sparse_route_blocks = params.num_k_blocks > 0 ? params.num_k_blocks : total_blk; + // The causal scheduler shortens total_blk for each Q subgroup. Deriving + // this ratio from total_blk would then collapse a 64-token route block to + // one K32 tile near the beginning of the sequence and drop half its keys. + // Keep the route-to-physical-tile mapping tied to the full K shape instead. + int k_route_block_tokens = cute::ceil_div(int(size<0>(K_2D)), sparse_route_blocks); + int k_tiles_per_sparse_route_block = + cute::max(1, ceil_div(k_route_block_tokens, get<1>(TileShapeQK{}))); + auto route_block_to_k_tile = [&](int route_block, int micro_tile) { + return route_block * k_tiles_per_sparse_route_block + micro_tile; + }; + + auto prefetch_sparse_k_block = [&](int route_block) { + if (route_block < 0 || route_block >= sparse_route_blocks) return; + for (int micro_tile = 0; micro_tile < k_tiles_per_sparse_route_block; ++micro_tile) { + int logical_block = route_block_to_k_tile(route_block, micro_tile); + if (logical_block >= total_blk) break; + 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)); + } + } + } + }; + + 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) { + 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; + } + + 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); + } + + CUTLASS_PRAGMA_UNROLL + for (int VV = 0; VV < VTiles; VV++) { + prefetch(prefetch_v_cur, pVgV_cur(_, _, _, VV, k_idx)); + } + + if (subgroup_selected) { + if constexpr (!is_cache && CausalMask) { + // K32 splits each logical 64-token sparse route block into two + // physical K tiles. The diagonal can therefore occur in any + // physical tile, rather than only in the final K tile. + 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 < tSrS.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) { + tSrS(i) = ElementS(-INFINITY); + } + } + } else if constexpr (FullMask_) { + 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 < tSrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + tSrS(i) += ElementS(params.mask[col_idx + row_idx * seq_len]); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tSrS.size(); ++i) { + int row_idx = get<0>(cS_thread(i)); + int col_idx = get<1>(cS_thread(i)); + tSrS(i) += + (row_idx < seq_len && col_idx < seq_len) ? ElementS(params.mask[col_idx + row_idx * seq_len]) + : ElementS(-INFINITY); + } + } + } + + if constexpr (!is_cache) { + if (check_remainder_k && K == total_blk - 1) { + FragSCol 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) ? ElementS(sycl::nan(0u)) : ElementS(-INFINITY); + } + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tSrS.size(); i++) { + tSrS(i) = sycl::fmin(tSrS(i), broadcast<1>(k_rem_mask, tSrS, i)); + } + } + } + + auto rescale = softmax(first_block, tSrS, tA_max, tA_sum); + reorder(tSrS, tArP); + + 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)); + } + } + + 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 < sparse_route_blocks) { + prefetch_sparse_k_block(sparse_prefetch_block); + } + + barrier_wait(ScopeWorkgroup); + }; + + auto run_sparse_route_block = [&](int route_block, bool subgroup_selected, auto& subgroup_started, + int sparse_prefetch_block) { + if (route_block < 0 || route_block >= sparse_route_blocks) return; + for (int micro_tile = 0; micro_tile < k_tiles_per_sparse_route_block; ++micro_tile) { + int K = route_block_to_k_tile(route_block, micro_tile); + if (K >= total_blk) break; + bool first_selected_block = subgroup_selected ? !subgroup_started : false; + 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); + if (subgroup_selected) subgroup_started = true; + } + } else 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); + if (subgroup_selected) 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); + if (subgroup_selected) subgroup_started = true; + } + } + }; + + if (lut_rows_base != nullptr && valid_blocks_base != nullptr) { + bool subgroup_started = false; + if (sparse_q_rows_in_tile == 1) { + 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) : sparse_route_blocks; + 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 >= sparse_route_blocks) return; + frontier_pos += 1; + if (frontier_pos < row_valid) { + frontier_cur_block += row_ptr[frontier_pos]; + } else { + frontier_cur_block = sparse_route_blocks; + } + }; + + 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 >= sparse_route_blocks) break; + prefetch_sparse_k_block(sparse_prefetch_block); + } + + while (row_cur_block < sparse_route_blocks) { + int next_block = row_cur_block; + bool subgroup_selected = (subgroup_q_row_in_tile == 0); + int sparse_prefetch_block = pop_single_sparse_block(prefetch_pos, prefetch_cur_block); + run_sparse_route_block(next_block, subgroup_selected, subgroup_started, sparse_prefetch_block); + advance_single_sparse_block(row_pos, row_cur_block); + } + } else { + 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; + bool subgroup_started_rows[kMaxSparseRowsPerTile]; + 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; + subgroup_started_rows[row] = false; + 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) : sparse_route_blocks; + if (row_valid[row] > 0) { + active_rows[active_row_count++] = row; + } + } else { + row_valid[row] = 0; + row_pos[row] = 0; + row_cur_block[row] = sparse_route_blocks; + } + } + + auto find_sparse_block = [&](int* frontier_pos, int* frontier_cur_block) { + int block = sparse_route_blocks; + 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 >= sparse_route_blocks) 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] = sparse_route_blocks; + } + } + } + }; + + 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 >= sparse_route_blocks) break; + prefetch_sparse_k_block(sparse_prefetch_block); + } + + int next_block = find_sparse_block(row_pos, row_cur_block); + while (next_block < sparse_route_blocks) { + int selected_row = subgroup_q_row_in_tile; + bool subgroup_selected = subgroup_q_row_in_tile < sparse_q_rows_in_tile && + row_pos[selected_row] < row_valid[selected_row] && + row_cur_block[selected_row] == next_block; + int sparse_prefetch_block = pop_sparse_block(prefetch_pos, prefetch_cur_block); + run_sparse_route_block(next_block, subgroup_selected, subgroup_started_rows[selected_row], + sparse_prefetch_block); + + 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); + } + } + } + + CUTLASS_DEVICE + FragSRow softmax(bool first_block, FragS& tS, FragSRow& tS_max, FragSRow& tS_sum) { + auto tS_bmax = reduce<1>(tS, sycl::maximum{}); + FragSRow rescale; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_max.size(); i++) { + ElementS new_max = sycl::max(tS_max(i), params.scale * tS_bmax(i)); + rescale(i) = sycl::native::exp2(tS_max(i) - new_max); + tS_max(i) = new_max; + } + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS.size(); i++) { + tS(i) = sycl::native::exp2(params.scale * tS(i) - broadcast<0>(tS_max, tS, i)); + } + + if (!first_block) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < tS_sum.size(); i++) { + tS_sum(i) *= rescale(i); + } + } + + auto tS_bsum = reduce<1>(tS, sycl::plus{}); + 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 824053d1b0..379a05fcea 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 @@ -236,6 +236,54 @@ void sdpa_impl_qks8_sparse_qtile256_row64k_pvhalf( 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_bf16_sparse_sdpa_d64( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); + +void sdpa_impl_fp16_sparse_sdpa_d64( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); + +void sdpa_impl_bf16_sparse_sdpa_row_linear( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); + +void sdpa_impl_fp16_sparse_sdpa_row_linear( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); + +void sdpa_impl_bf16_sparse_sdpa_qtile256_row64k( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); + +void sdpa_impl_fp16_sparse_sdpa_qtile256_row64k( + sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, void* O_ptr, void* mask, 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); #endif // ARK_SYCL_TLA } // namespace ark diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp index 5a6dec850f..890c328341 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp @@ -486,4 +486,4 @@ inline void sycl_tla_igemm_s8s8_dequant(sycl::queue* q, int m, int n, int k, con #endif // ARK_XPU && ARK_SYCL_TLA -} // namespace ark \ No newline at end of file +} // 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 index a8b4d2d10c..32f49abbfd 100644 --- 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 @@ -35,6 +35,7 @@ #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_sdpa_fwd_mainloop.hpp" #include "stla/xe_sparse_sage_fwd_kernel.hpp" #endif #include "flash_attention_v2/collective/fmha_fusion.hpp" @@ -1038,6 +1039,122 @@ struct SparseSageConfig { return run(options); } }; + +// Independent native-precision sparse path built on the dense SDPA mainloop +// (SPARSESDPAFwdMainloop). Unlike SparseSageConfig it does not dequantize INT8 +// Q/K: Q/K/V are bf16/fp16 and the softmax applies `params.scale` directly. +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 SparseSDPAConfig { + static constexpr int SGTileQ = get<0>(shape_div(TileShapeQK{}, shape(SubgroupLayoutQK{})))(); + // Native-precision Q*K MMA: bf16 -> float/bf16 DPAS, fp16 -> float/half DPAS. + 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 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::sdpa::XeDefault; + if constexpr (Causal) { + using CollectiveMainloop = + cutlass::fmha::collective::SPARSESDPAFwdMainloop; + 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::SPARSESDPAFwdMainloop; + 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::SPARSESDPAFwdMainloop; + 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 SDPA 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 SDPA 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 SDPA cached decode requires block_K/block_V cache tensors"); + } + if (options.lut == nullptr || options.valid_block_num == nullptr) { + throw std::runtime_error("Sparse SDPA requires lut and valid_block_num"); + } + if (options.seq_len_kv_cache > 0) { + return run(options); + } + return run(options); + } +}; #endif // ARK_SDPA_ENABLE_SPARSE // ======================================================================== @@ -1178,6 +1295,55 @@ inline int launch_sparse_sage_prefill_kernel_64(Options const& options) { SubgroupLayoutPV, PipelineStages, false, ElementQ, ElementK, ElementV, ElementO>::run(options); } + +template +inline int launch_sparse_sdpa_prefill_kernel_128(Options const& options) { + constexpr int PipelineStages = 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 ? SparseSDPAConfig::run(options) + : SparseSDPAConfig::run(options); +} + +template +inline int launch_sparse_sdpa_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 ? SparseSDPAConfig::run(options) + : SparseSDPAConfig::run(options); +} + +template +inline int launch_sparse_sdpa_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 + ? SparseSDPAConfig::run(options) + : SparseSDPAConfig::run(options); +} #endif // ARK_SDPA_ENABLE_SPARSE #if defined(ARK_SDPA_ENABLE_DENSE) 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 index 27ab86385c..f258a5a053 100644 --- 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 @@ -15,8 +15,18 @@ from auto_round_kernel.xpu_loader import ensure_xpu_lib -def ensure_sparse_binding() -> None: - ensure_xpu_lib(required_symbols=("sage_sparse",)) +def ensure_sparse_binding(*, required_symbols: tuple[str, ...] = ("sage_sparse",)) -> None: + search_roots = ( + REPO_PARENT / "auto_round_kernel", + REPO_PARENT / "auto_round_kernel" / "xbuild_diffuser", + REPO_PARENT / "auto_round_kernel" / "xbuild", + REPO_PARENT / "auto_round_kernel" / "xbuild_bf16_v2", + REPO_PARENT / "auto_round_kernel" / "ark-xbuild", + ) + ensure_xpu_lib( + required_symbols=required_symbols, + search_roots=search_roots, + ) def quantize_qk(tensor: torch.Tensor, block_size: int) -> tuple[torch.Tensor, torch.Tensor]: @@ -43,7 +53,7 @@ def quantize_qk(tensor: torch.Tensor, block_size: int) -> tuple[torch.Tensor, to def build_sparse_metadata_and_mask( batch: int, heads: int, - seq_len: int, + seq_len_q: int, quant_block_size: int, query_tile_tokens: int, per_query_tile_selection: list[list[int]], @@ -51,17 +61,19 @@ def build_sparse_metadata_and_mask( is_causal: bool = False, sparse_q_block_tokens: int | None = None, sparse_k_block_tokens: int | None = None, + seq_len_kv: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + seq_len_kv = seq_len_q if seq_len_kv is None else seq_len_kv q_block_tokens = quant_block_size if sparse_q_block_tokens is None else sparse_q_block_tokens k_block_tokens = quant_block_size if sparse_k_block_tokens is None else sparse_k_block_tokens - q_blocks = (seq_len + q_block_tokens - 1) // q_block_tokens - kv_blocks = (seq_len + k_block_tokens - 1) // k_block_tokens - active_query_tiles = (seq_len + query_tile_tokens - 1) // query_tile_tokens + q_blocks = (seq_len_q + q_block_tokens - 1) // q_block_tokens + kv_blocks = (seq_len_kv + k_block_tokens - 1) // k_block_tokens + active_query_tiles = (seq_len_q + 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) + mask = torch.full((batch, 1, seq_len_q, seq_len_kv), -1.0e9, dtype=torch.float32, device=device) q_blocks_per_query_tile = max(1, query_tile_tokens // q_block_tokens) for qblk in range(q_blocks): @@ -75,11 +87,11 @@ def build_sparse_metadata_and_mask( 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) + q_end = min(q_start + query_tile_tokens, seq_len_q) for qt in range(q_start, q_end): for selected in selected_blocks: k_start = selected * k_block_tokens - k_end = min(k_start + k_block_tokens, seq_len) + k_end = min(k_start + k_block_tokens, seq_len_kv) if not is_causal: mask[:, :, qt : qt + 1, k_start:k_end] = 0.0 else: @@ -98,6 +110,33 @@ def _clamp_selected_blocks(kv_blocks: int, per_query_tile_selection: list[list[i return clamped +def bf16_sparse_reference( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None, + *, + scale: float, + enable_gqa: bool, +) -> torch.Tensor: + # Run the reference on CPU: torch-xpu's softmax is numerically broken on large + # masked tensors (sums to <1, drops visible entries), which corrupts the reference. + q = query.float().cpu() + k = key.float().cpu() + v = value.float().cpu() + mask = None if attn_mask is None else attn_mask.float().cpu() + if enable_gqa and q.shape[1] != k.shape[1]: + repeat = q.shape[1] // k.shape[1] + k = k.repeat_interleave(repeat, dim=1) + v = v.repeat_interleave(repeat, dim=1) + + scores = torch.matmul(q, k.transpose(-1, -2)) * scale + if mask is not None: + scores = scores + mask + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v).to(query.dtype).to(query.device) + + def run_case( head_dim: int, block_size: int = 64, @@ -201,6 +240,163 @@ def run_case( ) +def run_case_sdpa( + dtype: torch.dtype, + head_dim: int, + *, + seq_len_q: int = 256, + seq_len_kv: int = 256, + num_heads_q: int = 4, + num_heads_kv: int | None = None, + is_causal: bool = False, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +) -> None: + """Validate the independent native-precision sparse-SDPA path (bf16/fp16) + against a dense reference built from the same block selection.""" + ensure_sparse_binding(required_symbols=("sage_sparse", "block_sparse_sdpa")) + device = torch.device("xpu") + batch = 1 + num_heads_kv = num_heads_q if num_heads_kv is None else num_heads_kv + scale = 1.0 / math.sqrt(head_dim) + query_tile_tokens = q_tile_override or (64 if head_dim == 64 else 256) + k_block_tokens = 64 if sparse_k_block_tokens is None else sparse_k_block_tokens + kv_blocks = (seq_len_kv + k_block_tokens - 1) // k_block_tokens + active_query_tiles = (seq_len_q + query_tile_tokens - 1) // query_tile_tokens + if not is_causal: + per_query_tile_selection = [list(range(max(1, min(kv_blocks, 3)))) for _ in range(active_query_tiles)] + if active_query_tiles > 1 and kv_blocks > 3: + per_query_tile_selection[-1] = list(range(kv_blocks - 3, kv_blocks)) + case_name = f"{dtype}_prefill" + else: + per_query_tile_selection = [list(range(min(kv_blocks, idx + 1))) for idx in range(active_query_tiles)] + case_name = f"{dtype}_prefill_causal" + per_query_tile_selection = _clamp_selected_blocks(kv_blocks, per_query_tile_selection) + + torch.manual_seed(9026 + head_dim + num_heads_q + (num_heads_kv * 13) + seq_len_q + seq_len_kv) + query = torch.randn((batch, num_heads_q, seq_len_q, head_dim), dtype=dtype, device=device) + key = torch.randn((batch, num_heads_kv, seq_len_kv, head_dim), dtype=dtype, device=device) + value = torch.randn((batch, num_heads_kv, seq_len_kv, head_dim), dtype=dtype, device=device) + + effective_sparse_q_block_tokens = 64 if sparse_q_block_tokens is None else sparse_q_block_tokens + effective_sparse_k_block_tokens = 64 if sparse_k_block_tokens is None else sparse_k_block_tokens + lut, valid, dense_mask = build_sparse_metadata_and_mask( + batch, + num_heads_q, + seq_len_q, + 64, + query_tile_tokens, + per_query_tile_selection, + device, + is_causal=is_causal, + sparse_q_block_tokens=effective_sparse_q_block_tokens, + sparse_k_block_tokens=effective_sparse_k_block_tokens, + seq_len_kv=seq_len_kv, + ) + + dense_out = bf16_sparse_reference( + query, + key, + value, + dense_mask, + scale=scale, + enable_gqa=num_heads_q != num_heads_kv, + ) + sparse_out = ark.block_sparse_sdpa( + query, + key, + value, + lut, + valid, + is_causal=is_causal, + scale=scale, + q_tile_override=q_tile_override, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + 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"[block_sparse_sdpa][{case_name}] D={head_dim} Hq={num_heads_q} Hkv={num_heads_kv} " + f"Sq={seq_len_q} Skv={seq_len_kv} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}" + ) + if max_diff > 2e-2 or mean_diff > 2e-3: + raise RuntimeError( + f"block_sparse_sdpa mismatch for dtype={dtype}, D={head_dim}, Hq={num_heads_q}, Hkv={num_heads_kv}, " + f"Sq={seq_len_q}, Skv={seq_len_kv}, causal={is_causal}" + ) + + +def run_case_sdpa_full( + dtype: torch.dtype, + head_dim: int, + *, + seq_len_q: int = 256, + seq_len_kv: int = 256, + num_heads_q: int = 4, + q_tile_override: int = 0, + sparse_q_block_tokens: int | None = None, + sparse_k_block_tokens: int | None = None, +) -> None: + """Dense gate: selecting every KV block must make sparse-SDPA match the dense reference.""" + ensure_sparse_binding(required_symbols=("sage_sparse", "block_sparse_sdpa")) + device = torch.device("xpu") + batch = 1 + scale = 1.0 / math.sqrt(head_dim) + query_tile_tokens = q_tile_override or (64 if head_dim == 64 else 256) + k_block_tokens = 64 if sparse_k_block_tokens is None else sparse_k_block_tokens + kv_blocks = (seq_len_kv + k_block_tokens - 1) // k_block_tokens + active_query_tiles = (seq_len_q + query_tile_tokens - 1) // query_tile_tokens + per_query_tile_selection = [list(range(kv_blocks)) for _ in range(active_query_tiles)] + + torch.manual_seed(77 + head_dim + num_heads_q) + query = torch.randn((batch, num_heads_q, seq_len_q, head_dim), dtype=dtype, device=device) + key = torch.randn((batch, num_heads_q, seq_len_kv, head_dim), dtype=dtype, device=device) + value = torch.randn((batch, num_heads_q, seq_len_kv, head_dim), dtype=dtype, device=device) + + lut, valid, dense_mask = build_sparse_metadata_and_mask( + batch, + num_heads_q, + seq_len_q, + 64, + query_tile_tokens, + per_query_tile_selection, + device, + is_causal=False, + sparse_q_block_tokens=64 if sparse_q_block_tokens is None else sparse_q_block_tokens, + sparse_k_block_tokens=k_block_tokens, + seq_len_kv=seq_len_kv, + ) + + dense_out = bf16_sparse_reference(query, key, value, dense_mask, scale=scale, enable_gqa=False) + sparse_out = ark.block_sparse_sdpa( + query, + key, + value, + lut, + valid, + is_causal=False, + scale=scale, + q_tile_override=q_tile_override, + sparse_q_block_tokens=sparse_q_block_tokens, + sparse_k_block_tokens=sparse_k_block_tokens, + 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"[block_sparse_sdpa][{dtype}_all_selected] D={head_dim} max_diff={max_diff:.6f} mean_diff={mean_diff:.6f}") + if max_diff > 2e-2 or mean_diff > 2e-3: + raise RuntimeError(f"block_sparse_sdpa all-selected mismatch for dtype={dtype}, D={head_dim}") + + def run_multi_row_tile_case() -> None: device = torch.device("xpu") batch = 1 @@ -341,6 +537,37 @@ def main() -> None: ) run_case(64, is_causal=True) run_case(128, is_causal=True) + # Independent sparse-SDPA path (bf16 + fp16) + for dtype in (torch.bfloat16, torch.float16): + run_case_sdpa(dtype, 64) + run_case_sdpa(dtype, 128, q_tile_override=64, num_heads_q=32, num_heads_kv=8) + run_case_sdpa( + dtype, + 128, + seq_len_q=512, + seq_len_kv=768, + num_heads_q=32, + num_heads_kv=8, + q_tile_override=256, + sparse_q_block_tokens=256, + sparse_k_block_tokens=64, + ) + run_case_sdpa(dtype, 128, is_causal=True, q_tile_override=64) + # Regression coverage for causal masking across K32 qtile256 microtiles. + run_case_sdpa( + dtype, + 128, + seq_len_q=512, + seq_len_kv=512, + num_heads_q=8, + num_heads_kv=8, + is_causal=True, + q_tile_override=256, + sparse_q_block_tokens=256, + sparse_k_block_tokens=64, + ) + run_case_sdpa_full(dtype, 64) + run_case_sdpa_full(dtype, 128, q_tile_override=64) if __name__ == "__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 index 02426f3913..5476207284 100644 --- 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 @@ -16,7 +16,16 @@ def ensure_sparse_binding() -> None: - ensure_xpu_lib(required_symbols=("sage_sparse",)) + ensure_xpu_lib( + required_symbols=("sage_sparse",), + search_roots=( + REPO_PARENT / "auto_round_kernel", + REPO_PARENT / "auto_round_kernel" / "xbuild_diffuser", + REPO_PARENT / "auto_round_kernel" / "xbuild", + REPO_PARENT / "auto_round_kernel" / "xbuild_bf16_v2", + REPO_PARENT / "auto_round_kernel" / "ark-xbuild", + ), + ) def _to_layout(tensor: torch.Tensor, tensor_layout: str) -> torch.Tensor: @@ -264,6 +273,74 @@ def run_case( assert torch.equal(crafted_meta["block_map"], crafted_meta["raw_block_map"]) +def run_lut_roundtrip_case(*, num_heads_q: int = 40, q_blocks: int = 8, k_blocks: int = 32) -> None: + """Regression test for the XPU boolean-mask assignment bug. + + torch-xpu-ops' boolean-mask in-place assignment (`tensor[mask] = value`) silently + misses entries on 4-D int64/bool tensors of these sizes, which used to leave 0s in + ``_block_map_lut_torch``'s ``filled_matrix`` and produce **-1** LUT entries. Those + -1s then drove the ``scatter_`` in ``_lut_to_block_map`` out of bounds + (``ScatterGatherKernels.cpp:233``). The LUT must be non-negative and the + block_map -> lut -> block_map round-trip must be lossless. + """ + from auto_round_kernel.sparse_attention import _block_map_lut_torch, _lut_to_block_map + + device = torch.device("xpu") + torch.manual_seed(1234) + block_map = torch.rand(1, num_heads_q, q_blocks, k_blocks, device=device) < 0.5 + lut, valid_block_num = _block_map_lut_torch(block_map) + torch.xpu.synchronize() + + assert lut.dtype == torch.int32 + assert valid_block_num.dtype == torch.int32 + assert tuple(lut.shape) == tuple(block_map.shape) + # Regression: on XPU this was -1 because the masked assignment silently failed. + assert int(lut.min().item()) >= 0, f"lut must be non-negative, got min={int(lut.min().item())}" + assert int(valid_block_num.max().item()) <= k_blocks + + reconstructed = _lut_to_block_map(lut, valid_block_num) + torch.xpu.synchronize() + assert torch.equal(reconstructed, block_map), "block_map -> lut -> block_map round-trip must be lossless" + + +def run_fill_block_map_case() -> None: + """Regression test for the vectorized ``_fill_block_map_torch``. + + The naive implementation looped ``for rank in range(k_blocks)`` (thousands of + sequential kernel launches at long sequences); the vectorized rewrite must be + exactly equivalent. Runs on CPU against a reference loop implementation. + """ + from auto_round_kernel.sparse_attention import _fill_block_map_torch + + def reference_loop(final_map, num_to_select, sorted_indices): + 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 + + for shape in [(1, 4, 8, 32), (1, 40, 8, 32), (1, 40, 128, 512), (2, 16, 64, 128)]: + b, h, q, k = shape + for seed in range(4): + torch.manual_seed(seed) + probs = torch.rand(shape) + _, sorted_indices = torch.sort(probs, dim=-1, descending=True) + num_to_select = torch.randint(0, k + 1, (b, h, q), dtype=torch.int64) + final_map = torch.rand(shape) < 0.1 # some blocks already forced + assert torch.equal( + _fill_block_map_torch(final_map, num_to_select, sorted_indices), + reference_loop(final_map, num_to_select, sorted_indices), + ), f"vectorized _fill_block_map_torch differs from reference at shape={shape} seed={seed}" + + def run_causal_decoupled_wrapper_case(*, tensor_layout: str) -> None: device = torch.device("xpu") batch = 1 @@ -410,6 +487,8 @@ def main() -> None: num_heads_kv=1, ) run_causal_decoupled_wrapper_case(tensor_layout=tensor_layout) + run_lut_roundtrip_case() + run_fill_block_map_case() if __name__ == "__main__": diff --git a/auto_round_extension/ark/benchmarks/BF16_SPARSE_TILE_SHAPE_ANALYSIS_20260818.md b/auto_round_extension/ark/benchmarks/BF16_SPARSE_TILE_SHAPE_ANALYSIS_20260818.md new file mode 100644 index 0000000000..56086fcf3f --- /dev/null +++ b/auto_round_extension/ark/benchmarks/BF16_SPARSE_TILE_SHAPE_ANALYSIS_20260818.md @@ -0,0 +1,217 @@ +# BF16 Sparse SDPA Tile Shape Analysis + +Date: 2026-08-18 + +## Scope + +This note summarizes the tile-shape discussion for the native BF16 sparse SDPA kernel. The question is whether reducing the workgroup Q tile from 256 to 64 while increasing the GEMM head-dimension tile from 32 to 128 can reach current performance. + +The analysis is based on: + +- Current ARK sparse BF16 SDPA implementation: + - `auto_round_kernel/wrapper/include/sycl_tla_sdpa_sparse.hpp` + - `auto_round_kernel/wrapper/include/stla/xe_sparse_sdpa_fwd_mainloop.hpp` + - `auto_round_kernel/sdpa_sparse_sdpa.cpp` +- SYCL-TLA tuning guide: + - `/home/yiliu4/workspace/sycl-tla/media/docs/cpp/cute/12_intel_performance_guide.md` + +## Current BF16 Sparse Tile + +For head_dim=128 with `q_tile_override=256`, the current native BF16 sparse SDPA path uses: + +```cpp +using ShapeQK = Shape<_256, _32, _32>; +using ShapePV = Shape<_256, _32, _32>; +using ShapeOut = Shape<_256, _128>; +using SubgroupLayoutQK = Layout>; +``` + +Interpretation: + +| Item | Value | +|---|---:| +| Q rows per workgroup | 256 | +| K-token columns per QK tile | 32 | +| Head-dim chunk per QK loop | 32 | +| Head-dim loops for head_dim=128 | 4 | +| QK score tile per workgroup | 256 x 32 | +| Output tile per workgroup | 256 x 128 | +| Subgroups per workgroup | 16 | +| Work-items per workgroup | 256 | +| Q rows per subgroup | 16 | + +The sparse API variant `sdpa_impl_bf16_sparse_sdpa_qtile256_row64k` passes `sparse_q_block_size=256`, so one sparse LUT row covers the full 256-token Q workgroup tile. The logical sparse K route block is 64 tokens, but the physical QK K-token tile is 32, so each selected K64 route block maps to two K32 micro-tiles. + +## Proposed Shape + +The proposed direction is to reduce M and increase the head-dimension K chunk: + +```cpp +using ShapeQK = Shape<_64, _32, _128>; +``` + +Here, `K=128` means the GEMM inner dimension, i.e. the head-dimension chunk, not the token-axis K tile. + +To keep the per-subgroup Q work comparable to the current kernel, the matching subgroup layout would likely be: + +```cpp +using SubgroupLayoutQK = Layout>; +``` + +This keeps: + +```text +Q rows per subgroup = 64 / 4 = 16 +``` + +which matches the current: + +```text +Q rows per subgroup = 256 / 16 = 16 +``` + +Keeping 16 subgroups for `M=64` would give only 4 Q rows per subgroup and would also change the DPAS M shape, so it is not an apples-to-apples comparison. + +## Apple-To-Apple Comparison + +| Config | ShapeQK | Q rows/WG | K-token tile | Head-dim chunk | Head-dim loops | Subgroups/WG | Q rows/SG | Main benefit | Main risk | +|---|---|---:|---:|---:|---:|---:|---:|---|---| +| Current | `Shape<_256, _32, _32>` | 256 | 32 | 32 | 4 | 16 | 16 | Large Q tile, stable fragments | More head-dim loop iterations | +| Proposed | `Shape<_64, _32, _128>` | 64 | 32 | 128 | 1 | 4 | 16 | Fewer head-dim loops | Larger Q/K fragments and 4x more workgroups | + +Current pseudocode: + +```cpp +for q_tile in Q step 256: + for selected k_token_tile in sparse_route: // physical token tile = 32 + S[256, 32] = 0 + + for d in head_dim step 32: // 4 loops + q_frag = load Q[256, 32] + k_frag = load K[32, 32] + S += q_frag @ k_frag.T + + P = softmax(S) + O[256, 128] += P @ V[32, 128] +``` + +Proposed pseudocode: + +```cpp +for q_tile in Q step 64: + for selected k_token_tile in sparse_route: // physical token tile = 32 + S[64, 32] = 0 + + for d in head_dim step 128: // 1 loop + q_frag = load Q[64, 128] + k_frag = load K[32, 128] + S += q_frag @ k_frag.T + + P = softmax(S) + O[64, 128] += P @ V[32, 128] +``` + +For the same 256-query region, the proposed shape launches four workgroups instead of one. The QK math is similar, but the execution shape is different: fewer head-dim loops per workgroup, more workgroups, less per-workgroup Q reuse, and larger Q/K fragments. + +## Coexisting Fragment Estimate + +The table below estimates logical fragment elements per subgroup. It is not an exact GRF count, but it shows the direction of register pressure. Prefetch fragments are excluded because `XE_PREFETCH_2D` has no destination register fragment. + +Assumptions: + +- Current: `ShapeQK = Shape<_256, _32, _32>`, `ShapePV = Shape<_256, _32, _32>`, `ShapeOut = Shape<_256, _128>`, 16 subgroups/WG. +- Proposed: `ShapeQK = Shape<_64, _32, _128>`, `ShapePV = Shape<_64, _32, _32>`, `ShapeOut = Shape<_64, _128>`, 4 subgroups/WG. +- Both keep 16 Q rows per subgroup. +- Copy fragments and MMA fragments are counted separately because both can be live around `reorder(copy_frag, mma_frag)` and `cute::gemm(...)`. + +### QK Phase + +| Fragment | Role | Current per SG | Proposed per SG | Change | +|---|---:|---:|---:|---:| +| `tQrQ` | Q copy fragment | `16 x 32 = 512` | `16 x 128 = 2048` | 4.0x | +| `tSrQ` | Q MMA fragment | `16 x 32 = 512` | `16 x 128 = 2048` | 4.0x | +| `tKrK` | K copy fragment | `32 x 32 = 1024` | `32 x 128 = 4096` | 4.0x | +| `tSrK` | K MMA fragment | `32 x 32 = 1024` | `32 x 128 = 4096` | 4.0x | +| `tSrS` | QK score accumulator | `16 x 32 = 512` | `16 x 32 = 512` | 1.0x | +| `tArA` | PV/O accumulator | `16 x 128 = 2048` | `16 x 128 = 2048` | 1.0x | +| `tA_max + tA_sum` | Softmax row state | `16 + 16 = 32` | `16 + 16 = 32` | 1.0x | +| **Approximate live sum** | | **5664** | **14880** | **2.63x** | + +### PV Phase + +| Fragment | Role | Current per SG | Proposed per SG | Change | +|---|---:|---:|---:|---:| +| `tSrS / tArP` | Score/probability fragment | `512 + 512` | `512 + 512` | 1.0x | +| `tVrV` | V copy fragment, one V tile | `32 x 32 = 1024` | `32 x 32 = 1024` | 1.0x | +| `tArV` | V MMA fragment, one V tile | `32 x 32 = 1024` | `32 x 32 = 1024` | 1.0x | +| `tArA` | Full output accumulator | `16 x 128 = 2048` | `16 x 128 = 2048` | 1.0x | +| `tA_max + tA_sum` | Softmax row state | `32` | `32` | 1.0x | +| **Approximate live sum** | | **5152** | **5152** | **1.0x** | + +The pressure increase is concentrated in the QK phase. Reducing M from 256 to 64 does not reduce per-subgroup Q rows if the subgroup layout is adjusted from 16 subgroups to 4 subgroups. Increasing the head-dim chunk from 32 to 128 makes the Q/K copy and MMA fragments 4x wider. + +## Why Subgroups Per Workgroup Change + +The subgroup count changes only if we intentionally preserve the same per-subgroup Q work. + +Current: + +```cpp +ShapeQK = Shape<_256, _32, _32>; +SubgroupLayoutQK = Layout>; +``` + +```text +16 subgroups/WG * 16 lanes/subgroup = 256 work-items/WG +Q rows/subgroup = 256 / 16 = 16 +``` + +Proposed apples-to-apples layout: + +```cpp +ShapeQK = Shape<_64, _32, _128>; +SubgroupLayoutQK = Layout>; +``` + +```text +4 subgroups/WG * 16 lanes/subgroup = 64 work-items/WG +Q rows/subgroup = 64 / 4 = 16 +``` + +If we kept 16 subgroups with `M=64`, each subgroup would cover only 4 Q rows. That would make the proposed kernel a different experiment: smaller per-subgroup work, different DPAS M behavior, and more overhead relative to useful compute. + +## Expected Performance Direction + +The SYCL-TLA performance guide gives the key tradeoff: + +- Increasing K can reduce K-loop and 2D block load issue overhead. +- Increasing K can also increase GRF pressure because larger copy/MMA fragments are live. +- Too much GRF pressure can cause compiler spill and erase the benefit. + +For this sparse BF16 SDPA kernel, `M64, headK128` is not expected to be a guaranteed win because: + +1. It reduces head-dim loops from 4 to 1. +2. It launches 4x more workgroups for the same Q region. +3. It makes Q/K copy and MMA fragments 4x wider per subgroup. +4. Attention already keeps score, softmax, probability, V, and output accumulator fragments live. +5. The estimated QK-phase live fragment footprint increases by about 2.6x per subgroup. + +## Recommendation + +Do not assume `Shape<_64, _32, _128>` will reach current performance. It is worth testing only as an A/B if profiling shows the current kernel is dominated by head-dim loop or 2D load issue overhead and unitrace/compiler reporting shows no spill. + +A safer first experiment is: + +```cpp +ShapeQK = Shape<_128, _32, _64>; +``` + +This reduces the head-dim loop count from 4 to 2 while taking a smaller GRF-risk step than `headK=128`. + +Validation should include: + +1. Correctness against dense SDPA. +2. Kernel timing on the same topk, layout, dtype, and sequence length. +3. `unitrace -d -v` to check GRF count and spill/private memory per thread. +4. ComputeBasic and stall metrics if timing changes materially. + diff --git a/auto_round_extension/ark/benchmarks/bench_sparse_topk.py b/auto_round_extension/ark/benchmarks/bench_sparse_topk.py index 2813ed48b6..9efe91e070 100644 --- a/auto_round_extension/ark/benchmarks/bench_sparse_topk.py +++ b/auto_round_extension/ark/benchmarks/bench_sparse_topk.py @@ -23,13 +23,13 @@ from auto_round_kernel.xpu_loader import load_xpu_lib -def _load_sparse_binding(ext_path: Path) -> None: +def _load_sparse_binding(ext_path: Path, *, required_symbols: tuple[str, ...]) -> None: module_name = "auto_round_kernel._bench.auto_round_kernel_xpu" print(f"Loading XPU extension from {ext_path.resolve()} as {module_name}") - load_xpu_lib(ext_path, required_symbols=("sage_sparse",), module_name=module_name) + load_xpu_lib(ext_path, required_symbols=required_symbols, module_name=module_name) -def _resolve_sparse_binding(args: argparse.Namespace) -> Path | None: +def _resolve_sparse_binding(args: argparse.Namespace, *, required_symbols: tuple[str, ...]) -> 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: @@ -37,9 +37,11 @@ def _resolve_sparse_binding(args: argparse.Namespace) -> Path | None: 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}") + raise RuntimeError( + f"Unable to locate built XPU extension with required symbols {required_symbols} in {args.xbuild_dir}" + ) return candidates[-1] - if current_file is not None and hasattr(ark.xpu_lib, "sage_sparse"): + if current_file is not None and all(hasattr(ark.xpu_lib, symbol) for symbol in required_symbols): try: current_path = Path(current_file).resolve() if current_path.is_relative_to(local_kernel_dir.resolve()): @@ -47,22 +49,38 @@ def _resolve_sparse_binding(args: argparse.Namespace) -> Path | None: return None except Exception: pass - candidates = sorted((local_kernel_dir / "xbuild").glob("auto_round_kernel_xpu*.so")) + build_roots = [root for root in local_kernel_dir.iterdir() if root.is_dir() and "build" in root.name] + candidates: list[Path] = [] + for root in build_roots: + candidates.extend(sorted(root.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") + raise RuntimeError( + f"Unable to locate built XPU extension with required symbols {required_symbols} under {local_kernel_dir}" + ) + candidates.sort(key=lambda path: path.stat().st_mtime) return candidates[-1] def ensure_sparse_binding(args: argparse.Namespace) -> None: - ext_path = _resolve_sparse_binding(args) + required_symbols = ("sage_sparse", "block_sparse_sdpa") + ext_path = _resolve_sparse_binding(args, required_symbols=required_symbols) if ext_path is not None: - _load_sparse_binding(ext_path) + _load_sparse_binding(ext_path, required_symbols=required_symbols) def is_xpu_available() -> bool: return hasattr(torch, "xpu") and torch.xpu.is_available() +def resolve_dtype(name: str) -> torch.dtype: + normalized = name.strip().lower() + if normalized == "fp16": + return torch.float16 + if normalized == "bf16": + return torch.bfloat16 + raise ValueError(f"Unsupported dtype={name!r}; supported values: fp16, bf16") + + def bench(fn, warmup: int, iters: int) -> float: for _ in range(warmup): out = fn() @@ -476,18 +494,22 @@ def hnd_to_nhd(x: torch.Tensor) -> torch.Tensor: 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) + ark_row = next((row for row in rows if row["mode"] == "dense_ark_sdpa" 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"]) + ark_ms = None if ark_row is None else float(ark_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["speedup_vs_ark"] = 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 + row["speedup_vs_ark"] = (ark_ms / latency_ms) if ark_ms is not None else None batch = int(row["batch"]) num_heads_q = int(row["num_heads_q"]) seq_len = int(row["seq_len"]) @@ -499,10 +521,15 @@ def summarize_speedups(rows: list[dict[str, object]]) -> list[dict[str, object]] if mode in { "dense_torch_sdpa", "dense_sagev1", + "dense_ark_sdpa", "sparse_kernel_only", "sparse_e2e", "sparse_qtile256_row64k_kernel_only", "sparse_qtile256_row64k_e2e", + "sparse_sdpa_bf16_kernel_only", + "sparse_sdpa_bf16_e2e", + "sparse_sdpa_fp16_kernel_only", + "sparse_sdpa_fp16_e2e", }: row["baseline_tflops"] = flops_to_tflops(dense_flops, float(latency_ms)) work_ratio = 1.0 @@ -514,9 +541,9 @@ def summarize_speedups(rows: list[dict[str, object]]) -> list[dict[str, object]] def print_summary(rows: list[dict[str, object]]) -> None: print( - "| layout | seq_len | mode | pattern | gqa_group | topk | selected_ratio | blocks/row | latency (ms) | baseline_tflops | effective_tflops | status | speedup_vs_torch | speedup_vs_sagev1 |" + "| layout | seq_len | mode | pattern | gqa_group | topk | selected_ratio | blocks/row | latency (ms) | baseline_tflops | effective_tflops | status | speedup_vs_torch | speedup_vs_sagev1 | speedup_vs_ark |" ) - print("|---|---|---|---|---|---|---|---|---|---|---|---|---|---|") + 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}" @@ -526,10 +553,11 @@ def print_summary(rows: list[dict[str, object]]) -> None: 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}" + sp_ark = "-" if row.get("speedup_vs_ark") is None else f"{float(row['speedup_vs_ark']):.3f}" print( f"| {row['tensor_layout']} | {row['seq_len']} | {row['mode']} | {row['attention_pattern']} | {row['gqa_group_size']} | {topk} | {ratio} | " f"{blocks} | {latency} | {baseline_tflops} | {effective_tflops} | {row['status']} | {sp_torch} | " - f"{sp_sage} |" + f"{sp_sage} | {sp_ark} |" ) if row["note"]: print(f"note[{row['mode']}]: {row['note']}") @@ -548,7 +576,7 @@ def write_csv(rows: list[dict[str, object]], output_csv: Path) -> None: def run_single_benchmark(args: argparse.Namespace, *, seq_len: int, tensor_layout: str) -> list[dict[str, object]]: device = torch.device("xpu") - dtype = torch.float16 + dtype = resolve_dtype(args.dtype) 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( @@ -641,6 +669,29 @@ def run_single_benchmark(args: argparse.Namespace, *, seq_len: int, tensor_layou iters=args.iters, ) ) + rows.append( + try_benchmark( + "dense_ark_sdpa", + lambda: ark.sdpa( + q, + k, + v, + is_causal=args.causal, + scale=scale, + tensor_layout=tensor_layout, + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=seq_len, + tensor_layout=tensor_layout, + head_dim=args.head_dim, + dtype=dtype, + is_causal=args.causal, + warmup=args.warmup, + iters=args.iters, + ) + ) for topk in args.topk: preprocess = None @@ -891,6 +942,108 @@ def run_single_benchmark(args: argparse.Namespace, *, seq_len: int, tensor_layou selected_blocks_per_row=selected_blocks_per_row, ) ) + rows.append( + try_benchmark( + "sparse_sdpa_fp16_kernel_only" if dtype == torch.float16 else "sparse_sdpa_bf16_kernel_only", + lambda preprocess=preprocess: ( + hnd_to_nhd( + ark.block_sparse_sdpa( + nhd_to_hnd(q_nhd_src), + nhd_to_hnd(k_nhd_src), + nhd_to_hnd(v_nhd_src), + preprocess["lut"], + preprocess["valid_block_num"], + is_causal=args.causal, + scale=scale, + 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 tensor_layout == "HND" + else ark.block_sparse_sdpa( + q, + k, + v, + preprocess["lut"], + preprocess["valid_block_num"], + is_causal=args.causal, + scale=scale, + 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=tensor_layout, + ) + ), + batch=args.batch, + num_heads_q=args.num_heads_q, + num_heads_kv=args.num_heads_kv, + seq_len=seq_len, + tensor_layout=tensor_layout, + 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, + ) + ) + rows.append( + try_benchmark( + "sparse_sdpa_fp16_e2e" if dtype == torch.float16 else "sparse_sdpa_bf16_e2e", + lambda topk=topk: ( + hnd_to_nhd( + ark.sparge_sage2_attn_meansim_topk_xpu_sdpa( + 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 tensor_layout == "HND" + else ark.sparge_sage2_attn_meansim_topk_xpu_sdpa( + q, + k, + v, + is_causal=args.causal, + scale=scale, + smooth_k=True, + simthreshd1=-1.0, + topk=topk, + attention_sink=False, + tensor_layout=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=seq_len, + tensor_layout=tensor_layout, + 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() @@ -921,6 +1074,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--num-heads-kv", type=int, default=40) parser.add_argument("--seq-len", type=int, nargs="+", default=[32768, 75600]) parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument( + "--dtype", + choices=("fp16", "bf16"), + default="fp16", + help="Input/output dtype used for dense SDPA, dense SAGE, and sparse attention benchmarks.", + ) parser.add_argument("--topk", type=float, nargs="+", default=[0.5, 0.25, 0.125]) parser.add_argument("--quant-block-size", type=int, default=64) parser.add_argument( diff --git a/auto_round_extension/ark/examples/flux_gen_bf16_sweep.py b/auto_round_extension/ark/examples/flux_gen_bf16_sweep.py new file mode 100644 index 0000000000..904c8e7e76 --- /dev/null +++ b/auto_round_extension/ark/examples/flux_gen_bf16_sweep.py @@ -0,0 +1,182 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import contextlib +import csv +import os +import sys +import time +from pathlib import Path + +import torch +from diffusers import FluxPipeline +from flux_sparse_patch import patch_flux_sparse_attention_from_env + +dtype = torch.bfloat16 + +MODEL_DEFAULT = "/mnt/disk3/models/black-forest-labs/FLUX.1-dev" +PROMPT_DEFAULT = "A cat holding a sign that says hello world" + + +def env_flag(name, default="1"): + return os.getenv(name, default).lower() not in {"0", "false", "no", "off"} + + +def sync_xpu(): + if hasattr(torch, "xpu") and hasattr(torch.xpu, "synchronize"): + torch.xpu.synchronize() + + +def main(): + model_id = os.getenv("FLUX_MODEL", MODEL_DEFAULT) + topks = os.getenv("FLUX_SPARSE_TOPKS", "0.5 0.25 0.125").split() + run_dense = env_flag("FLUX_RUN_DENSE", "1") + prompt = os.getenv("FLUX_PROMPT", PROMPT_DEFAULT) + height = int(os.getenv("FLUX_HEIGHT", "512")) + width = int(os.getenv("FLUX_WIDTH", "512")) + steps = int(os.getenv("FLUX_STEPS", "50")) + seed = int(os.getenv("FLUX_SEED", "0")) + guidance_scale = float(os.getenv("FLUX_GUIDANCE_SCALE", "3.5")) + max_sequence_length = int(os.getenv("FLUX_MAX_SEQUENCE_LENGTH", "512")) + + default_out = Path("benchmarks/results") / f"flux_bf16_{time.strftime('%Y%m%d_%H%M%S')}" + out_dir = Path(os.getenv("FLUX_OUTPUT_DIR", str(default_out))) + out_dir.mkdir(parents=True, exist_ok=True) + + device_id = os.getenv("ZE_AFFINITY_MASK_VALUE", "") + dev_suffix = f"_dev{device_id}" if device_id else "" + + print(f"[flux_sweep] model={model_id}", flush=True) + print( + f"[flux_sweep] topks={topks} run_dense={run_dense} size={height}x{width} steps={steps} seed={seed}", flush=True + ) + print(f"[flux_sweep] out_dir={out_dir}", flush=True) + print( + f"[flux_sweep] kernel={os.getenv('FLUX_SPARSE_KERNEL', '?')} " + 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')}", + flush=True, + ) + + pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=dtype) + # Offload-only: pipe.to(device) would load the whole ~54 GB model onto the + # 24.4 GB device and OOM. + # + # Use sequential (block-level) offload, NOT enable_model_cpu_offload(): + # model-level offload pulls the whole ~46 GB transformer onto the device, + # peaking at ~24 GB (the full 24.4 GB device) during every transformer + # forward. The sparse BF16 preprocess (triton_xpu) then has <0.4 GB of + # headroom, and its kernel launch flakily OOMs / resets the GPU + # (UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY / OUT_OF_RESOURCES / DEVICE_LOST). + # Sequential offload keeps only a few blocks resident (peak ~0.16 GB), so + # the triton_xpu sparse path always has room. + pipe.enable_sequential_cpu_offload() + + common_kwargs = dict( + height=height, + width=width, + guidance_scale=guidance_scale, + num_inference_steps=steps, + max_sequence_length=max_sequence_length, + generator=torch.Generator("cpu").manual_seed(seed), + ) + + rows = [] + + def do_run(tag, topk, png_name): + if topk is None: + ctx = contextlib.nullcontext(None) + os.environ.pop("FLUX_SPARSE_TOPK", None) + else: + os.environ["FLUX_SPARSE_TOPK"] = str(topk) + ctx = patch_flux_sparse_attention_from_env(pipe.transformer) + + sync_xpu() + t0 = time.perf_counter() + with ctx as stats: + image = pipe(prompt, output_type="pil", **common_kwargs).images[0] + sync_xpu() + wall_s = time.perf_counter() - t0 + + png_path = out_dir / png_name + image.save(png_path) + + row = { + "tag": tag, + "topk": "" if topk is None else topk, + "wall_s": round(wall_s, 3), + "sparsity": "", + "calls": "", + "sparse_calls": "", + "runtime_fallbacks": "", + "unsupported_fallbacks": "", + "png": str(png_path), + } + if topk is not None: + if stats is None or stats.sparse_calls == 0: + raise RuntimeError( + f"topk={topk}: sparse path did not run (sparse_calls=0) — check FLUX_SPARSE_KERNEL / qtile256 config" + ) + if stats.runtime_fallbacks or stats.unsupported_fallbacks: + raise RuntimeError( + f"topk={topk}: silent fallback detected — runtime={stats.runtime_fallbacks}, " + f"unsupported={stats.unsupported_fallbacks}" + ) + row["sparsity"] = round(float(stats.avg_sparsity), 4) + row["calls"] = stats.total_calls + row["sparse_calls"] = stats.sparse_calls + row["runtime_fallbacks"] = stats.runtime_fallbacks + row["unsupported_fallbacks"] = stats.unsupported_fallbacks + print( + f"[flux_sweep] DONE tag={tag} topk={topk} wall={wall_s:.3f}s " + f"sparsity={stats.avg_sparsity:.4f} sparse_calls={stats.sparse_calls} " + f"runtime_fallbacks={stats.runtime_fallbacks}", + flush=True, + ) + else: + print(f"[flux_sweep] DONE tag={tag} wall={wall_s:.3f}s", flush=True) + rows.append(row) + return row + + if run_dense: + do_run("dense", None, f"flux_dense_512{dev_suffix}.png") + + for topk in topks: + do_run("sparse_bf16_qtile256", topk, f"flux_bf16_qtile256_topk{topk}_512{dev_suffix}.png") + + # CSV + markdown summary (device-suffixed so parallel instances sharing + # FLUX_OUTPUT_DIR do not clobber each other). + csv_path = out_dir / f"sweep_summary{dev_suffix}.csv" + fieldnames = [ + "tag", + "topk", + "wall_s", + "sparsity", + "calls", + "sparse_calls", + "runtime_fallbacks", + "unsupported_fallbacks", + "png", + ] + with open(csv_path, "w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + md_path = out_dir / f"SWEEP_SUMMARY{dev_suffix}.md" + with open(md_path, "w", encoding="utf-8") as fh: + fh.write("| tag | topk | wall_s | sparsity | calls | sparse_calls | runtime_fallbacks | png |\n") + fh.write("|---|---|---|---|---|---|---|---|\n") + fh.writelines( + f"| {r['tag']} | {r['topk']} | {r['wall_s']} | {r['sparsity']} | {r['calls']} " + f"| {r['sparse_calls']} | {r['runtime_fallbacks']} | `{r['png']}` |\n" + for r in rows + ) + + print(f"[flux_sweep] summary: {csv_path}", flush=True) + print(f"[flux_sweep] summary: {md_path}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/auto_round_extension/ark/examples/flux_sparse_patch.py b/auto_round_extension/ark/examples/flux_sparse_patch.py index 6695bcbb92..ab512c9aae 100644 --- a/auto_round_extension/ark/examples/flux_sparse_patch.py +++ b/auto_round_extension/ark/examples/flux_sparse_patch.py @@ -20,6 +20,19 @@ from wan_sparse_patch import _parse_bool_env, _parse_optional_int_env, ensure_ark_sparse_binding +def _parse_sparse_kernel_env(name: str = "FLUX_SPARSE_KERNEL", default: str = "int8") -> str: + value = os.getenv(name, default).strip().lower() + aliases = { + "int8": "int8", + "qks8": "int8", + "bf16": "bf16", + } + normalized = aliases.get(value) + if normalized is None: + raise ValueError(f"{name} must be one of: int8, bf16") + return normalized + + def _normalize_attention_mask( attn_mask: torch.Tensor | None, batch: int, @@ -90,6 +103,7 @@ def __init__( q_tile_override: int, sparse_q_block_tokens: int | None, sparse_k_block_tokens: int | None, + sparse_kernel: str, ): self.original_processor = original_processor self.stats = stats @@ -100,6 +114,7 @@ def __init__( 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.sparse_kernel = sparse_kernel self._warned_runtime_fallback = False @staticmethod @@ -192,22 +207,40 @@ def __call__( 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, - ) + if self.sparse_kernel == "bf16": + hidden_states, sparsity = ark.sparge_sage2_attn_meansim_topk_xpu_sdpa( + 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, + ) + else: + 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) @@ -270,8 +303,15 @@ def patch_flux_sparse_attention( q_tile_override: int = 0, sparse_q_block_tokens: int | None = None, sparse_k_block_tokens: int | None = None, + sparse_kernel: str = "int8", ): - ensure_ark_sparse_binding() + ensure_ark_sparse_binding( + required_symbols=( + ("block_sparse_sdpa", "sage_dynamic_quant_layout") + if sparse_kernel == "bf16" + else ("sage_sparse", "sage_dynamic_quant_layout") + ) + ) originals: list[tuple[FluxAttention, object]] = [] stats = FluxSparseAttentionStats() @@ -289,6 +329,7 @@ def patch_flux_sparse_attention( q_tile_override=q_tile_override, sparse_q_block_tokens=sparse_q_block_tokens, sparse_k_block_tokens=sparse_k_block_tokens, + sparse_kernel=sparse_kernel, ) ) originals.append((module, original)) @@ -310,4 +351,5 @@ def patch_flux_sparse_attention_from_env(transformer): 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"), + sparse_kernel=_parse_sparse_kernel_env(), ) diff --git a/auto_round_extension/ark/examples/run_flux.py b/auto_round_extension/ark/examples/run_flux.py index 0c1bd3544c..57fcf902d6 100644 --- a/auto_round_extension/ark/examples/run_flux.py +++ b/auto_round_extension/ark/examples/run_flux.py @@ -24,6 +24,19 @@ def env_flag(name, default="0"): return os.getenv(name, default).lower() not in {"0", "false", "no", "off"} +def env_sparse_kernel(name="FLUX_SPARSE_KERNEL", default="int8"): + value = os.getenv(name, default).strip().lower() + aliases = { + "int8": "int8", + "qks8": "int8", + "bf16": "bf16", + } + normalized = aliases.get(value) + if normalized is None: + raise ValueError(f"{name} must be one of: int8, bf16") + return normalized + + 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() @@ -41,10 +54,10 @@ def env_flag(name, default="0"): model_id = os.getenv("FLUX_MODEL", "~/workspace/models/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() + pipe.enable_model_cpu_offload(device=device) +elif benchmark_scope != "block": + pipe.to(device) height = int(os.getenv("FLUX_HEIGHT", "1024")) width = int(os.getenv("FLUX_WIDTH", "1024")) @@ -53,11 +66,15 @@ def env_flag(name, default="0"): 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"} +sparse_kernel = env_sparse_kernel() -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" -) +if use_sparse: + output_file = ( + f"flux_output_{height}x{width}_{num_inference_steps}steps_" + f"{guidance_scale}gs_sparse_{sparse_kernel}_topk{os.getenv('FLUX_SPARSE_TOPK', '0.5')}.png" + ) +else: + output_file = f"flux_output_{height}x{width}_{num_inference_steps}steps_{guidance_scale}gs_dense.png" output_path = os.getenv("FLUX_OUTPUT", output_file) prompt = os.getenv("FLUX_PROMPT", "A cat holding a sign that says hello world") @@ -195,7 +212,8 @@ 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"[flux_sparse] enabled attention sparse patch: kernel={sparse_kernel}" + f" 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')}" @@ -524,6 +542,7 @@ def run_block_benchmark(current_run_tag): "max_sequence_length": max_sequence_length, "seed": seed, "use_sparse": use_sparse, + "sparse_kernel": sparse_kernel if use_sparse else None, "cpu_offload_enabled": cpu_offload_enabled, "block_kind": block_state["block_kind"], "block_index": block_state["block_index"], @@ -543,7 +562,7 @@ def run_block_benchmark(current_run_tag): return result -run_tag = "sparse" if use_sparse else "dense" +run_tag = f"sparse_{sparse_kernel}" if use_sparse else "dense" def maybe_run_with_profiler(run_callable, current_run_tag, record_label="flux_generate"): @@ -615,6 +634,7 @@ def run_benchmark(run_callable, current_run_tag, benchmark_scope): "max_sequence_length": max_sequence_length, "seed": seed, "use_sparse": use_sparse, + "sparse_kernel": sparse_kernel if use_sparse else None, "sparse_topk": os.getenv("FLUX_SPARSE_TOPK", "0.5"), "sparse_smooth_k": os.getenv("FLUX_SPARSE_SMOOTH_K", "1"), "output_path": output_path, @@ -704,6 +724,7 @@ def run_denoising_benchmark(current_run_tag): "max_sequence_length": max_sequence_length, "seed": seed, "use_sparse": use_sparse, + "sparse_kernel": sparse_kernel if use_sparse else None, "sparse_topk": os.getenv("FLUX_SPARSE_TOPK", "0.5"), "sparse_smooth_k": os.getenv("FLUX_SPARSE_SMOOTH_K", "1"), "output_path": output_path, diff --git a/auto_round_extension/ark/examples/wan_sparse_patch.py b/auto_round_extension/ark/examples/wan_sparse_patch.py index 2e61ca6cdc..0d82071f84 100644 --- a/auto_round_extension/ark/examples/wan_sparse_patch.py +++ b/auto_round_extension/ark/examples/wan_sparse_patch.py @@ -36,10 +36,18 @@ def _parse_optional_int_env(name: str) -> int | None: return None if parsed == 0 else parsed -def ensure_ark_sparse_binding() -> None: +def ensure_ark_sparse_binding( + *, required_symbols: tuple[str, ...] = ("sage_sparse", "sage_dynamic_quant_layout") +) -> None: ensure_xpu_lib( - required_symbols=("sage_sparse", "sage_dynamic_quant_layout"), - search_roots=(KERNEL_DIR, KERNEL_DIR / "xbuild", KERNEL_DIR / "xbuild_diffuser"), + required_symbols=required_symbols, + search_roots=( + KERNEL_DIR, + KERNEL_DIR / "xbuild_diffuser", + KERNEL_DIR / "xbuild", + KERNEL_DIR / "xbuild_bf16_v2", + KERNEL_DIR / "ark-xbuild", + ), ) diff --git a/auto_round_extension/ark/tools/bench_sparse_bf16_fp16_sweep.sh b/auto_round_extension/ark/tools/bench_sparse_bf16_fp16_sweep.sh new file mode 100755 index 0000000000..29159a2eac --- /dev/null +++ b/auto_round_extension/ark/tools/bench_sparse_bf16_fp16_sweep.sh @@ -0,0 +1,110 @@ +#!/usr/bin/env bash +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +# bench_sparse_bf16_fp16_sweep.sh +# =============================== +# Run the sparse BF16 and FP16 kernel benchmark sweep (both HND and NHD layouts) +# against the rebuilt ARK XPU extension, and save CSVs under benchmarks/results/. +# +# USAGE +# ./tools/bench_sparse_bf16_fp16_sweep.sh +# +# ENV OVERRIDES (all optional) +# DTYPES dtypes to sweep, space-separated (default: "bf16 fp16") +# SEQ_LENS seq lengths, space-separated (default: "75000") +# TOPKS topk values, space-separated (default: "0.5 0.25 0.125") +# LAYOUTS tensor layouts, space-separated (default: "HND NHD") +# HEADS num heads q (kv = q) (default: 40) +# HEAD_DIM head dim (default: 128) +# Q_TILE q_tile_override (default: 256) +# WARMUP warmup iters (default: 2) +# ITERS measured iters (default: 3) +# ZE_AFFINITY_MASK_VALUE GPU device index (default: 6) +# XPU_SO explicit path to the built .so (default: newest under ark-xbuild) +# OUTPUT_DIR where CSVs are written (default: benchmarks/results) +# SPARGE_PREPROCESS_BACKEND sparse preprocess backend (default: triton_xpu); torch | triton_xpu | auto +# +# PREREQUISITES +# - oneAPI 2026.1 runtime (sourced from /opt/intel/oneapi/setvars.sh) +# - venv at .venv/bin/python with torch.xpu +# - rebuilt extension under auto_round_kernel/ark-xbuild/ +# - a free GPU (pick via ZE_AFFINITY_MASK_VALUE; check `xpu-smi ps` first) + +set -uo pipefail + +SCRIPT_PATH="$(readlink -f "$0")" +REPO_ROOT="$(cd "$(dirname "${SCRIPT_PATH}")/.." && pwd)" +PYTHON_BIN="${REPO_ROOT}/.venv/bin/python" +BENCH="${REPO_ROOT}/benchmarks/bench_sparse_topk.py" + +DTYPES="${DTYPES:-bf16 fp16}" +SEQ_LENS="${SEQ_LENS:-75000}" +TOPKS="${TOPKS:-0.5 0.25 0.125}" +LAYOUTS="${LAYOUTS:-HND NHD}" +HEADS="${HEADS:-40}" +HEAD_DIM="${HEAD_DIM:-128}" +Q_TILE="${Q_TILE:-256}" +WARMUP="${WARMUP:-2}" +ITERS="${ITERS:-3}" +ZE_AFFINITY_MASK_VALUE="${ZE_AFFINITY_MASK_VALUE:-6}" +XPU_SO="${XPU_SO:-}" +OUTPUT_DIR="${OUTPUT_DIR:-${REPO_ROOT}/benchmarks/results}" + +cd "${REPO_ROOT}" + +# Re-exec under the render/video groups if this shell does not already have them. +if [[ "${1:-}" == "--inner" ]]; then + shift +elif command -v sg >/dev/null 2>&1; then + current_groups="$(id -nG || true)" + if [[ ! " ${current_groups} " =~ [[:space:]]render[[:space:]] ]] || [[ ! " ${current_groups} " =~ [[:space:]]video[[:space:]] ]]; then + exec sg render -c "sg video -c '${SCRIPT_PATH} --inner'" + fi +fi + +set +u +source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 +set -u + +export MKLROOT=/opt/intel/oneapi/mkl/2026.1 +export ZE_AFFINITY_MASK="${ZE_AFFINITY_MASK_VALUE}" +export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND="${SPARGE_PREPROCESS_BACKEND:-triton_xpu}" # auto | torch | triton_xpu + +if [[ -z "${XPU_SO}" ]]; then + XPU_SO="$(find "${REPO_ROOT}/auto_round_kernel/ark-xbuild" -maxdepth 1 -name 'auto_round_kernel_xpu*.so' | sort | tail -n 1)" +fi +if [[ -z "${XPU_SO}" ]]; then + echo "error: no auto_round_kernel_xpu*.so found under auto_round_kernel/ark-xbuild" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" +stamp="$(date +%Y%m%d_%H%M%S)" + +echo "XPU process snapshot:" +xpu-smi ps || true +echo "Using XPU extension: ${XPU_SO}" +echo "dtypes=${DTYPES} seq=${SEQ_LENS} topk=${TOPKS} layouts=${LAYOUTS} heads=${HEADS} head_dim=${HEAD_DIM} q_tile=${Q_TILE}" +echo + +for dtype in ${DTYPES}; do + out_csv="${OUTPUT_DIR}/bench_sparse_${dtype}_${stamp}.csv" + echo "========== dtype=${dtype} -> ${out_csv} ==========" + "${PYTHON_BIN}" "${BENCH}" \ + --dtype "${dtype}" \ + --seq-len ${SEQ_LENS} \ + --topk ${TOPKS} \ + --tensor-layout ${LAYOUTS} \ + --head-dim "${HEAD_DIM}" \ + --num-heads-q "${HEADS}" --num-heads-kv "${HEADS}" \ + --q-tile-override "${Q_TILE}" \ + --sparse-q-block-tokens "${Q_TILE}" \ + --sparse-k-block-tokens 64 \ + --warmup "${WARMUP}" --iters "${ITERS}" \ + --xpu-so "${XPU_SO}" \ + --output-csv "${out_csv}" + echo +done + +echo "SWEEP_DONE: CSVs written under ${OUTPUT_DIR}" diff --git a/auto_round_extension/ark/tools/repro_sparse_bf16_sdpa_bench.sh b/auto_round_extension/ark/tools/repro_sparse_bf16_sdpa_bench.sh new file mode 100755 index 0000000000..fb5f372ce1 --- /dev/null +++ b/auto_round_extension/ark/tools/repro_sparse_bf16_sdpa_bench.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash + +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +SCRIPT_PATH="$(readlink -f "$0")" +REPO_ROOT="/home/yiliu4/workspace/auto-round/auto_round_extension/ark" +PYTHON_BIN="${REPO_ROOT}/.venv/bin/python" +OUTPUT_CSV="${OUTPUT_CSV:-/tmp/bench_sparse_topk_bf16_dev6.csv}" +ZE_AFFINITY_MASK_VALUE="${ZE_AFFINITY_MASK_VALUE:-6}" +XPU_SO="${XPU_SO:-}" + +cd "${REPO_ROOT}" + +if [[ "${1:-}" == "--inner" ]]; then + shift +elif command -v sg >/dev/null 2>&1; then + current_groups="$(id -nG || true)" + if [[ ! " ${current_groups} " =~ [[:space:]]render[[:space:]] ]] || [[ ! " ${current_groups} " =~ [[:space:]]video[[:space:]] ]]; then + exec sg render -c "sg video -c '${SCRIPT_PATH} --inner'" + fi +fi + +# oneAPI's setvars.sh is not compatible with `set -u` in this environment. +set +u +source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 +set -u + +export MKLROOT=/opt/intel/oneapi/mkl/2026.1 +export ZE_AFFINITY_MASK="${ZE_AFFINITY_MASK_VALUE}" +export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND=torch + +if [[ -z "${XPU_SO}" ]]; then + XPU_SO="$(find "${REPO_ROOT}/auto_round_kernel/ark-xbuild" -maxdepth 1 -name 'auto_round_kernel_xpu*.so' | sort | tail -n 1)" +fi + +echo "XPU process snapshot:" +xpu-smi ps || true +echo +echo "Running sparse BF16 SDPA benchmark on device ${ZE_AFFINITY_MASK}..." +echo "Using XPU extension: ${XPU_SO}" + +exec "${PYTHON_BIN}" benchmarks/bench_sparse_topk.py \ + --dtype bf16 \ + --xpu-so "${XPU_SO}" \ + --output-csv "${OUTPUT_CSV}" diff --git a/auto_round_extension/ark/tools/run_bf16_sparse_bench.sh b/auto_round_extension/ark/tools/run_bf16_sparse_bench.sh new file mode 100755 index 0000000000..93c4b65ca2 --- /dev/null +++ b/auto_round_extension/ark/tools/run_bf16_sparse_bench.sh @@ -0,0 +1,97 @@ +#!/usr/bin/env bash +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +# +# run_bf16_sparse_bench.sh +# ======================== +# Documented reproduce command for the BF16 sparse SDPA benchmark on Intel XPU. +# +# WHAT IT DOES +# Runs benchmarks/bench_sparse_topk.py --dtype bf16 against the rebuilt +# XPU extension (auto_round_kernel/ark-xbuild), pinning a GPU device, and +# saves the CSV + a readable log under benchmarks/results/. +# +# USAGE +# ./tools/run_bf16_sparse_bench.sh +# +# ENV OVERRIDES +# ZE_AFFINITY_MASK_VALUE GPU device index to run on (default: 6) +# OUTPUT_DIR where CSV + log are written (default: benchmarks/results) +# XPU_SO explicit path to the built auto_round_kernel_xpu*.so +# (default: newest under auto_round_kernel/ark-xbuild) +# SPARGE_PREPROCESS_BACKEND sparse preprocess backend (default: triton_xpu) +# torch | triton_xpu | auto +# +# PREREQUISITES +# - oneAPI 2026.1 runtime (sourced from /opt/intel/oneapi/setvars.sh) +# - Python venv at .venv/bin/python (PyTorch 2.13.0+xpu) +# - a rebuilt extension with BF16 sparse symbols under auto_round_kernel/ark-xbuild/ +# - a free GPU (pick via ZE_AFFINITY_MASK_VALUE; check `xpu-smi ps` first) +# +# BENCHMARK CONFIG (bench_sparse_topk.py defaults) +# batch=1 num_heads_q/kv=40 head_dim=128 dtype=bf16 qtile=256 +# seq_len=[32768, 75600] tensor_layout=[HND, NHD] topk=[0.5, 0.25, 0.125] +# warmup=2 iters=3 +# +# NOTE +# First call per process pays a ~30s SYCL JIT for the sparse kernels; the +# full sweep above takes roughly 15-25 minutes on one GPU. + +set -euo pipefail + +SCRIPT_PATH="$(readlink -f "$0")" +REPO_ROOT="$(cd "$(dirname "${SCRIPT_PATH}")/.." && pwd)" +PYTHON_BIN="${REPO_ROOT}/.venv/bin/python" +OUTPUT_DIR="${OUTPUT_DIR:-${REPO_ROOT}/benchmarks/results}" +ZE_AFFINITY_MASK_VALUE="${ZE_AFFINITY_MASK_VALUE:-6}" +XPU_SO="${XPU_SO:-}" +SPARGE_PREPROCESS_BACKEND="${SPARGE_PREPROCESS_BACKEND:-triton_xpu}" + +cd "${REPO_ROOT}" + +# Re-exec under the render/video groups if this shell does not already have them. +if [[ "${1:-}" == "--inner" ]]; then + shift +elif command -v sg >/dev/null 2>&1; then + current_groups="$(id -nG || true)" + if [[ ! " ${current_groups} " =~ [[:space:]]render[[:space:]] ]] || [[ ! " ${current_groups} " =~ [[:space:]]video[[:space:]] ]]; then + exec sg render -c "sg video -c '${SCRIPT_PATH} --inner'" + fi +fi + +# oneAPI's setvars.sh is not compatible with `set -u` in this environment. +set +u +source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 +set -u + +export MKLROOT=/opt/intel/oneapi/mkl/2026.1 +export ZE_AFFINITY_MASK="${ZE_AFFINITY_MASK_VALUE}" +export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND="${SPARGE_PREPROCESS_BACKEND}" + +if [[ -z "${XPU_SO}" ]]; then + XPU_SO="$(find "${REPO_ROOT}/auto_round_kernel/ark-xbuild" -maxdepth 1 -name 'auto_round_kernel_xpu*.so' | sort | tail -n 1)" +fi +if [[ -z "${XPU_SO}" ]]; then + echo "error: no auto_round_kernel_xpu*.so found under auto_round_kernel/ark-xbuild" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" +stamp="$(date +%Y%m%d_%H%M%S)" +out_csv="${OUTPUT_DIR}/bf16_sparse_${stamp}.csv" +out_log="${OUTPUT_DIR}/bf16_sparse_${stamp}.log" + +echo "XPU process snapshot:" +xpu-smi ps || true +echo +echo "Running BF16 sparse SDPA benchmark on device ${ZE_AFFINITY_MASK}..." +echo "Using XPU extension: ${XPU_SO}" +echo "CSV: ${out_csv}" +echo "Log: ${out_log}" +echo + +"${PYTHON_BIN}" benchmarks/bench_sparse_topk.py \ + --dtype bf16 \ + --xpu-so "${XPU_SO}" \ + --output-csv "${out_csv}" 2>&1 | tee "${out_log}" diff --git a/auto_round_extension/ark/tools/run_sparse_sagev1_bench.sh b/auto_round_extension/ark/tools/run_sparse_sagev1_bench.sh new file mode 100755 index 0000000000..1f3e7c0cd0 --- /dev/null +++ b/auto_round_extension/ark/tools/run_sparse_sagev1_bench.sh @@ -0,0 +1,114 @@ +#!/usr/bin/env bash +# # Copyright (C) 2026 Intel Corporation +# # SPDX-License-Identifier: Apache-2.0 + +# +# run_sparse_sagev1_bench.sh +# ========================== +# Documented reproduce command for the sparse SAGE v1 SDPA benchmark on Intel XPU. +# +# This runs the INT8 sparse path (fp16 Q/K/V inputs, int8 quantized routing) that +# lives on the `main` branch: dense torch SDPA + dense sagev1 + the sparse +# kernel-only and sparse e2e modes, across top-k selections. +# +# WHAT IT DOES +# Runs benchmarks/bench_sparse_topk.py (fp16 / INT8 sparse) against the rebuilt +# XPU extension, pinning a GPU device, and saves a timestamped CSV + log under +# benchmarks/results/. +# +# USAGE +# ./tools/run_sparse_sagev1_bench.sh +# +# ENV OVERRIDES +# ZE_AFFINITY_MASK_VALUE GPU device index to run on (default: 6) +# OUTPUT_DIR where CSV + log are written (default: benchmarks/results) +# XPU_SO explicit path to the built auto_round_kernel_xpu*.so +# (default: newest under auto_round_kernel/ark-xbuild) +# BENCH_ARGS extra args passed to bench_sparse_topk.py +# (e.g. "--seq-len 32768 --topk 0.5 --iters 3 --tensor-layout HND") +# +# PREREQUISITES +# - oneAPI 2026.1 runtime (sourced from /opt/intel/oneapi/setvars.sh) +# - Python venv at .venv/bin/python (PyTorch 2.13.0+xpu) +# - a rebuilt extension with the sage_sparse binding under +# auto_round_kernel/ark-xbuild/ (on `main`: rebuild after `git checkout main`) +# - a free GPU (pick via ZE_AFFINITY_MASK_VALUE; check `xpu-smi ps` first) +# +# IMPORTANT +# The default preprocess backend is triton-xpu (via "auto"), which is faster +# than torch. On some level-zero drivers triton's JIT emitted SPIR-V that the +# runtime rejected ("InvalidModule: ... unknown extension +# 'SPV_INTEL_predicated_io'); that is fixed by _apply_xpu_triton_workarounds() +# in auto_round_kernel/sparge_preprocess_triton.py when the triton preprocess loads. +# If you still hit the SPIR-V error, force the torch backend: +# export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND=torch +# +# BENCHMARK CONFIG (bench_sparse_topk.py defaults on `main`) +# dtype=fp16 batch=1 num_heads_q/kv=40 head_dim=128 qtile=256 +# seq_len=[32768, 75600] tensor_layout=[HND, NHD] topk=[0.5, 0.25, 0.125] +# warmup=2 iters=3 +# +# NOTE +# First call per process pays a ~30s SYCL JIT for the sparse kernels; the full +# sweep takes roughly 15-25 minutes on one GPU. + +set -euo pipefail + +SCRIPT_PATH="$(readlink -f "$0")" +REPO_ROOT="$(cd "$(dirname "${SCRIPT_PATH}")/.." && pwd)" +PYTHON_BIN="${REPO_ROOT}/.venv/bin/python" +OUTPUT_DIR="${OUTPUT_DIR:-${REPO_ROOT}/benchmarks/results}" +ZE_AFFINITY_MASK_VALUE="${ZE_AFFINITY_MASK_VALUE:-6}" +XPU_SO="${XPU_SO:-}" +BENCH_ARGS="${BENCH_ARGS:-}" + +cd "${REPO_ROOT}" + +# Re-exec under the render/video groups if this shell does not already have them. +if [[ "${1:-}" == "--inner" ]]; then + shift +elif command -v sg >/dev/null 2>&1; then + current_groups="$(id -nG || true)" + if [[ ! " ${current_groups} " =~ [[:space:]]render[[:space:]] ]] || [[ ! " ${current_groups} " =~ [[:space:]]video[[:space:]] ]]; then + exec sg render -c "sg video -c '${SCRIPT_PATH} --inner'" + fi +fi + +# oneAPI's setvars.sh is not compatible with `set -u` in this environment. +set +u +source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 +set -u + +export MKLROOT=/opt/intel/oneapi/mkl/2026.1 +export ZE_AFFINITY_MASK="${ZE_AFFINITY_MASK_VALUE}" +# Required: the triton-xpu preprocess backend is rejected by the runtime on this +# machine (SPV_INTEL_predicated_io); force the torch backend. +# export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND=torch + +if [[ -z "${XPU_SO}" ]]; then + XPU_SO="$(find "${REPO_ROOT}/auto_round_kernel/ark-xbuild" -maxdepth 1 -name 'auto_round_kernel_xpu*.so' | sort | tail -n 1)" +fi +if [[ -z "${XPU_SO}" ]]; then + echo "error: no auto_round_kernel_xpu*.so found under auto_round_kernel/ark-xbuild" >&2 + exit 1 +fi + +mkdir -p "${OUTPUT_DIR}" +stamp="$(date +%Y%m%d_%H%M%S)" +out_csv="${OUTPUT_DIR}/sparse_sagev1_${stamp}.csv" +out_log="${OUTPUT_DIR}/sparse_sagev1_${stamp}.log" + +echo "XPU process snapshot:" +xpu-smi ps || true +echo +echo "Running sparse SAGE v1 SDPA benchmark on device ${ZE_AFFINITY_MASK}..." +echo "Using XPU extension: ${XPU_SO}" +echo "CSV: ${out_csv}" +echo "Log: ${out_log}" +echo + +# shellcheck disable=SC2086 +"${PYTHON_BIN}" benchmarks/bench_sparse_topk.py \ + --xpu-so "${XPU_SO}" \ + --output-csv "${out_csv}" \ + ${BENCH_ARGS} 2>&1 | tee "${out_log}" diff --git a/auto_round_extension/ark/tools/sweep_flux_bf16_topk.sh b/auto_round_extension/ark/tools/sweep_flux_bf16_topk.sh new file mode 100755 index 0000000000..babc5083a4 --- /dev/null +++ b/auto_round_extension/ark/tools/sweep_flux_bf16_topk.sh @@ -0,0 +1,105 @@ +#!/usr/bin/env bash +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 +# +# sweep_flux_bf16_topk.sh +# ======================= +# One-GPU FLUX.1-dev BF16-sparse topk sweep. Runs examples/flux_gen_bf16_sweep.py +# with the perf-tuned qtile256 config (q_tile=256, q_block=256, k_block=64) and +# saves PNGs + a summary under a shared FLUX_OUTPUT_DIR. +# +# To use both XPU devices, launch two instances concurrently with disjoint +# topk subsets pointing at the SAME FLUX_OUTPUT_DIR, e.g.: +# STAMP=$(date +%Y%m%d_%H%M%S) +# OUT=benchmarks/results/flux_bf16_${STAMP} +# ZE_AFFINITY_MASK_VALUE=5 FLUX_RUN_DENSE=1 FLUX_SPARSE_TOPKS="0.5 0.25" \ +# FLUX_OUTPUT_DIR=$OUT ./tools/sweep_flux_bf16_topk.sh & +# ZE_AFFINITY_MASK_VALUE=6 FLUX_RUN_DENSE=0 FLUX_SPARSE_TOPKS="0.125" \ +# FLUX_OUTPUT_DIR=$OUT ./tools/sweep_flux_bf16_topk.sh & +# wait +# +# ENV OVERRIDES (all optional) +# ZE_AFFINITY_MASK_VALUE XPU device index (default: 6) +# FLUX_MODEL model dir (diffusers format) (default: /mnt/disk3/models/black-forest-labs/FLUX.1-dev) +# FLUX_SPARSE_TOPKS space-separated topk list (default: "0.5 0.25 0.125") +# FLUX_RUN_DENSE 1/0 include a dense baseline generation (default: 1) +# FLUX_OUTPUT_DIR results dir (default: benchmarks/results/flux_bf16_) +# FLUX_HEIGHT / FLUX_WIDTH generation size (keep 512) (default: 512) +# FLUX_STEPS / FLUX_SEED / FLUX_PROMPT (defaults: 50 / 0 / "A cat ...") +# SPARGE_PREPROCESS_BACKEND torch | triton_xpu | auto (default: triton_xpu) +# XPU_SO explicit path to the built auto_round_kernel_xpu*.so +# (default: newest under auto_round_kernel/ark-xbuild) +# +# PREREQUISITES +# - oneAPI 2026.1 runtime (sourced from /opt/intel/oneapi/setvars.sh) +# - venv at .venv/bin/python with torch.xpu + diffusers +# - rebuilt extension with BF16 sparse symbols under auto_round_kernel/ark-xbuild/ +# - a free GPU (pick via ZE_AFFINITY_MASK_VALUE; check `xpu-smi ps` first) + +set -euo pipefail + +SCRIPT_PATH="$(readlink -f "$0")" +REPO_ROOT="$(cd "$(dirname "${SCRIPT_PATH}")/.." && pwd)" +PYTHON_BIN="${REPO_ROOT}/.venv/bin/python" + +ZE_AFFINITY_MASK_VALUE="${ZE_AFFINITY_MASK_VALUE:-6}" +FLUX_MODEL="${FLUX_MODEL:-/mnt/disk3/models/black-forest-labs/FLUX.1-dev}" +FLUX_SPARSE_TOPKS="${FLUX_SPARSE_TOPKS:-0.5 0.25 0.125}" +FLUX_RUN_DENSE="${FLUX_RUN_DENSE:-1}" +FLUX_HEIGHT="${FLUX_HEIGHT:-512}" +FLUX_WIDTH="${FLUX_WIDTH:-512}" +FLUX_STEPS="${FLUX_STEPS:-50}" +FLUX_SEED="${FLUX_SEED:-0}" +SPARGE_PREPROCESS_BACKEND="${SPARGE_PREPROCESS_BACKEND:-triton_xpu}" +XPU_SO="${XPU_SO:-}" + +cd "${REPO_ROOT}" + +# Re-exec under the render/video groups if this shell does not already have them. +if [[ "${1:-}" == "--inner" ]]; then + shift +elif command -v sg >/dev/null 2>&1; then + current_groups="$(id -nG || true)" + if [[ ! " ${current_groups} " =~ [[:space:]]render[[:space:]] ]] || [[ ! " ${current_groups} " =~ [[:space:]]video[[:space:]] ]]; then + exec sg render -c "sg video -c '${SCRIPT_PATH} --inner'" + fi +fi + +# oneAPI's setvars.sh is not compatible with `set -u` in this environment. +set +u +source /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 +set -u + +export MKLROOT=/opt/intel/oneapi/mkl/2026.1 +export ZE_AFFINITY_MASK="${ZE_AFFINITY_MASK_VALUE}" +export SAGE_ATTN_SPARSE_PREPROCESS_BACKEND="${SPARGE_PREPROCESS_BACKEND}" + +# qtile256 perf config is fixed for this sweep; kernel is env-overridable (bf16 default, int8 for SAGE). +export FLUX_SPARSE_KERNEL="${FLUX_SPARSE_KERNEL:-bf16}" +export FLUX_SPARSE_Q_TILE_OVERRIDE=256 +export FLUX_SPARSE_Q_BLOCK_TOKENS=256 +export FLUX_SPARSE_K_BLOCK_TOKENS=64 + +if [[ -z "${XPU_SO}" ]]; then + XPU_SO="$(find "${REPO_ROOT}/auto_round_kernel/ark-xbuild" -maxdepth 1 -name 'auto_round_kernel_xpu*.so' | sort | tail -n 1)" +fi +if [[ -z "${XPU_SO}" ]]; then + echo "error: no auto_round_kernel_xpu*.so found under auto_round_kernel/ark-xbuild" >&2 + exit 1 +fi + +if [[ -z "${FLUX_OUTPUT_DIR:-}" ]]; then + FLUX_OUTPUT_DIR="${REPO_ROOT}/benchmarks/results/flux_bf16_$(date +%Y%m%d_%H%M%S)" +fi +export FLUX_OUTPUT_DIR="${FLUX_OUTPUT_DIR}" +mkdir -p "${FLUX_OUTPUT_DIR}" + +echo "XPU process snapshot:" +xpu-smi ps || true +echo "Using XPU extension: ${XPU_SO}" +echo "device=${ZE_AFFINITY_MASK_VALUE} topks='${FLUX_SPARSE_TOPKS}' run_dense=${FLUX_RUN_DENSE}" +echo "qtile256 config: FLUX_SPARSE_Q_TILE_OVERRIDE=256 FLUX_SPARSE_Q_BLOCK_TOKENS=256 FLUX_SPARSE_K_BLOCK_TOKENS=64" +echo "output dir: ${FLUX_OUTPUT_DIR}" +echo + +"${PYTHON_BIN}" examples/flux_gen_bf16_sweep.py