ark: add BF16/FP16 Sparge variants

auto_round_extension/ark/.gitignore

@@ -1,5 +1,9 @@
 build
 xbuild
+ark-xbuild*
+ark-xbuild-builtin*
+xbuild_bf16*
 *.csv
 *.so
 *.pyc
+benchmarks/results/

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

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,
 )
 
 

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"),