From b827f56c4cb41cf2f8ad0d1a88498f24b2ec9818 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Mon, 13 Jul 2026 11:45:25 +0800 Subject: [PATCH 1/8] support fp32 in sycltla Signed-off-by: Zhenzhong1 --- .../ark/auto_round_kernel/__init__.py | 4 +- .../wrapper/include/sycl_tla_dense_gemm.hpp | 60 ++++++++++++------- 2 files changed, 39 insertions(+), 25 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index c5af8ae74..ee37187ed 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -244,8 +244,8 @@ def matmul_sycl_tla(A: torch.Tensor, B: torch.Tensor, bias: Optional[torch.Tenso raise ValueError("A and B must be 2D tensors") if A.device != B.device: raise ValueError("A and B must be on the same device") - if A.dtype not in (torch.float16, torch.bfloat16): - raise ValueError("matmul_sycl_tla only supports torch.float16 and torch.bfloat16") + if A.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise ValueError("matmul_sycl_tla only supports torch.float32, torch.float16 and torch.bfloat16") if B.dtype != A.dtype: raise ValueError("A and B must have the same dtype") diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp index 383bce4a4..12903cb20 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_dense_gemm.hpp @@ -42,24 +42,27 @@ using namespace cute; template void gemm_device_impl(ATensor const& A, BTensor const& B, CTensor& C, TiledMMA const& mma, const BiasElement* bias_ptr) { + /* Get workgroup and local IDs */ auto item = sycl::ext::oneapi::this_work_item::get_nd_item<2>(); auto wg_m = int(item.get_group(1)); auto wg_n = int(item.get_group(0)); auto local_id = int(item.get_local_id(0)); - Tensor cA = make_identity_tensor(A.shape()); - Tensor cB = make_identity_tensor(B.shape()); - Tensor cC = make_identity_tensor(C.shape()); + /* Create proxy coordinate tensors for each global tensor */ + Tensor cA = make_identity_tensor(A.shape()); // (M,K) + Tensor cB = make_identity_tensor(B.shape()); // (N,K) + Tensor cC = make_identity_tensor(C.shape()); // (M,N) + /* Split GEMM into workgroup tiles, and identify our workgroup's tile (wg_coord) */ auto wg_tile = mma.tile_mnk(); auto wg_coord = make_coord(wg_m, wg_n, 0); - Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); - Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); - Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); // (BLK_M,BLK_K,k) + Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); // (BLK_N,BLK_K,k) + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); // (BLK_M,BLK_N) - auto copy_a = make_block_2d_copy_A(XE_LOAD_2D<16, 8, 32, 32>{}, mma, A); - auto copy_b = make_block_2d_copy_B(XE_LOAD_2D_TRANSPOSE<32, 16, 8>{}, mma, B); + auto copy_a = make_block_2d_copy_A(mma, A); + auto copy_b = make_block_2d_copy_B(mma, B); auto copy_c = make_block_2d_copy_D(mma, C); auto thr_mma = mma.get_slice(local_id); @@ -119,29 +122,34 @@ void gemm_device_impl(ATensor const& A, BTensor const& B, CTensor& C, TiledMMA c using ElementOutput = typename CTensor::element_type; if constexpr (is_same_v) { - copy(copy_c, tCrC, tCgC); - } else { - Tensor tCrD = make_tensor_like(tCrC); - - CUTE_UNROLL - for (int i = 0; i < size(tCrC); ++i) { - auto coord = tCgC(i); - int col = int(get<1>(coord)); - float value = static_cast(tCrC(i)); - if (bias_ptr != nullptr && col < int(shape<1>(C))) { - value += static_cast(bias_ptr[col]); - } - tCrD(i) = static_cast(value); + if (bias_ptr == nullptr) { + copy(copy_c, tCrC, tCgC); + return; } + } - copy(copy_c, tCrD, tCgC); + Tensor tCrD = make_tensor_like(tCrC); + + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + auto coord = tCgC(i); + int col = int(get<1>(coord)); + float value = static_cast(tCrC(i)); + if (bias_ptr != nullptr && col < int(shape<1>(C))) { + value += static_cast(bias_ptr[col]); + } + tCrD(i) = static_cast(value); } + + copy(copy_c, tCrD, tCgC); } template auto choose_mma_op() { if constexpr (is_complete_v>) { return XE_DPAS_TT<8, TC, TA, TB>{}; + } else if constexpr (is_same_v && is_same_v) { + return XE_DPAS_TT<8, float, cute::tfloat32_t>{}; } else if constexpr (is_same_v) { return XE_DPAS_TT<8, float, cute::bfloat16_t>{}; } else { @@ -256,6 +264,12 @@ inline void sycl_tla_dense_gemm(sycl::queue* q, int m, int n, int k, const void* } switch (at) { + case BTLA_DTYPE::F32: + dense_gemm_detail::run_dense_gemm( + q, m, n, k, static_cast(a), static_cast(b), + static_cast(c), static_cast(bias)); + return; + case BTLA_DTYPE::F16: dense_gemm_detail::run_dense_gemm( q, m, n, k, static_cast(a), static_cast(b), @@ -269,7 +283,7 @@ inline void sycl_tla_dense_gemm(sycl::queue* q, int m, int n, int k, const void* return; default: - throw std::invalid_argument("sycl_tla_dense_gemm: only FP16 and BF16 are supported"); + throw std::invalid_argument("sycl_tla_dense_gemm: only FP32, FP16 and BF16 are supported"); } } From 14bbb0387b1f9c34faa68a54b61fd49e1778ac7e Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Mon, 13 Jul 2026 16:26:28 +0800 Subject: [PATCH 2/8] update CMakeLists & setup.py Signed-off-by: Zhenzhong1 --- .../ark/auto_round_kernel/CMakeLists.txt | 76 ++++++++++++------- auto_round_extension/ark/setup.py | 10 +++ 2 files changed, 60 insertions(+), 26 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt index 5d40df7a4..370d4ea5e 100755 --- a/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt +++ b/auto_round_extension/ark/auto_round_kernel/CMakeLists.txt @@ -4,7 +4,11 @@ project(ark LANGUAGES CXX) option(ARK_XPU "Build XPU kernels" OFF) option(ARK_UT "Build XPU kernels UT" OFF) option(ARK_RESCALE "Experimental" OFF) -set(ARK_DNNL_BUILD_SOURCE ON CACHE BOOL "Build oneDNN from source" FORCE) +if(ARK_XPU) + option(ARK_DNNL "Build oneDNN-backed kernels" OFF) +else() + option(ARK_DNNL "Build oneDNN-backed kernels" ON) +endif() # ARK_SYCL_TLA defaults to ON when ARK_XPU is enabled if(ARK_XPU AND NOT DEFINED ARK_SYCL_TLA) @@ -13,6 +17,14 @@ else() option(ARK_SYCL_TLA "Build SYCL TLA" OFF) endif() +if(ARK_SYCL_TLA AND NOT ARK_XPU) + message(FATAL_ERROR "ARK_SYCL_TLA requires ARK_XPU=ON") +endif() + +if(ARK_XPU AND NOT ARK_DNNL AND NOT ARK_SYCL_TLA) + message(FATAL_ERROR "XPU build requires ARK_SYCL_TLA=ON when ARK_DNNL=OFF") +endif() + include(FetchContent) FetchContent_Declare( pybind11 @@ -25,40 +37,43 @@ include(FindOpenMP) set(BTLA_ENABLE_OPENMP ON CACHE BOOL "BesTLA enable compiling OpenMP threading") -FetchContent_Declare( - dnnl - GIT_REPOSITORY https://github.com/uxlfoundation/oneDNN.git - GIT_TAG v3.10.2 -) -set(libs dnnl) -set(DNNL_LIBRARY_TYPE "STATIC" CACHE INTERNAL "") -set( - DNNL_ENABLE_PRIMITIVE - "MATMUL;ELTWISE;REDUCTION;REORDER" - CACHE STRING "oneDNN primitives used by ARK production code" -) +set(libs) -set(DNNL_BUILD_EXAMPLES OFF CACHE INTERNAL "") -set(DNNL_BUILD_TESTS OFF CACHE INTERNAL "") -set(DNNL_BUILD_DOC OFF CACHE INTERNAL "") -set(ONEDNN_BUILD_GRAPH OFF CACHE INTERNAL "") -set(ONEDNN_ENABLE_WORKLOAD "INFERENCE" CACHE INTERNAL "") +if(ARK_DNNL) + set(ARK_DNNL_BUILD_SOURCE ON CACHE BOOL "Build oneDNN from source" FORCE) + FetchContent_Declare( + dnnl + GIT_REPOSITORY https://github.com/uxlfoundation/oneDNN.git + GIT_TAG v3.10.2 + ) + list(APPEND libs dnnl) + set(DNNL_LIBRARY_TYPE "STATIC" CACHE INTERNAL "") + set(DNNL_ENABLE_PRIMITIVE "MATMUL;ELTWISE;REDUCTION;REORDER" CACHE STRING "oneDNN primitives used by ARK production code") + set(DNNL_BUILD_EXAMPLES OFF CACHE INTERNAL "") + set(DNNL_BUILD_TESTS OFF CACHE INTERNAL "") + set(DNNL_BUILD_DOC OFF CACHE INTERNAL "") + set(ONEDNN_BUILD_GRAPH OFF CACHE INTERNAL "") + set(ONEDNN_ENABLE_WORKLOAD "INFERENCE" CACHE INTERNAL "") +endif() if(ARK_XPU) - # Skip oneDNN CPU engine compilation for the XPU module build. - set(DNNL_CPU_RUNTIME "NONE" CACHE INTERNAL "") - set(DNNL_GPU_RUNTIME "SYCL" CACHE INTERNAL "") - FetchContent_MakeAvailable(dnnl) + if(ARK_DNNL) + set(DNNL_CPU_RUNTIME "NONE" CACHE INTERNAL "") + set(DNNL_GPU_RUNTIME "SYCL" CACHE INTERNAL "") + FetchContent_MakeAvailable(dnnl) + endif() set(BTLA_SYCL ON CACHE BOOL "BesTLA with SYCL") set(PY_NAME auto_round_kernel_xpu) set(ARK_TYPE ARK_XPU) find_package(IntelSYCL REQUIRED) list(APPEND libs IntelSYCL::SYCL_CXX) else() - set(DNNL_CPU_RUNTIME "OMP" CACHE INTERNAL "") - set(DNNL_CPU_THREADING_RUNTIME "OMP") - set(DNNL_GPU_RUNTIME "NONE" CACHE INTERNAL "") - FetchContent_MakeAvailable(dnnl) + if(ARK_DNNL) + set(DNNL_CPU_RUNTIME "OMP" CACHE INTERNAL "") + set(DNNL_CPU_THREADING_RUNTIME "OMP") + set(DNNL_GPU_RUNTIME "NONE" CACHE INTERNAL "") + FetchContent_MakeAvailable(dnnl) + endif() set(BTLA_SYCL OFF CACHE BOOL "BesTLA without SYCL") set(PY_NAME auto_round_kernel_cpu) set(ARK_TYPE ARK_CPU) @@ -134,9 +149,15 @@ endif() pybind11_add_module(${PY_NAME} ${SRCS} ${HEADERS} ${SDPA_GENERATED_SRCS}) target_compile_features(${PY_NAME} PRIVATE cxx_std_17) target_compile_definitions(${PY_NAME} PRIVATE PY_NAME=${PY_NAME} ${ARK_TYPE}=1) + +if(ARK_DNNL) + target_compile_definitions(${PY_NAME} PRIVATE ARK_DNNL=1) +endif() + if(ARK_RESCALE) target_compile_definitions(${PY_NAME} PRIVATE ARK_RESCALE=1) endif() + 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) @@ -175,5 +196,8 @@ if(ARK_UT) target_compile_options(${TEST_NAME} PRIVATE -fsycl -fno-sycl-instrument-device-code) target_link_options(${TEST_NAME} PRIVATE ${SYCL_TLA_LINK_FLAGS}) endif() + if(ARK_DNNL) + target_compile_definitions(${TEST_NAME} PRIVATE ARK_DNNL=1) + endif() target_link_libraries(${TEST_NAME} PRIVATE ${libs}) endif() diff --git a/auto_round_extension/ark/setup.py b/auto_round_extension/ark/setup.py index fb957b4ac..7bc781f6d 100644 --- a/auto_round_extension/ark/setup.py +++ b/auto_round_extension/ark/setup.py @@ -89,8 +89,17 @@ def detect_oneapi_version(): "and that the 'icx' compiler is in your PATH." ) + +def env_flag(name, default=False): + value = os.environ.get(name) + if value is None: + return default + return value.strip().lower() in ("1", "on", "true", "yes") + + requirements = fetch_requirements("requirements.txt") enable_sycl_tla = parse_major_minor(oneapi_version) >= (2025, 3) +enable_dnnl = env_flag("ARK_DNNL", False) def get_system_memory_gb(): @@ -181,6 +190,7 @@ def run(self): "-DCMAKE_BUILD_TYPE=Release", "-DCMAKE_CXX_COMPILER=icx", "-DARK_XPU=ON", + f"-DARK_DNNL={'ON' if enable_dnnl else 'OFF'}", f"-DARK_SYCL_TLA={'ON' if enable_sycl_tla else 'OFF'}", ] if sys.platform == "win32": From 03f03075439f365ffa0b99fcb694390d44c9bdae Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Mon, 13 Jul 2026 16:33:23 +0800 Subject: [PATCH 3/8] decompling oneDNNL Signed-off-by: Zhenzhong1 --- .../ark/auto_round_kernel/ark.cpp | 19 +- .../ark/auto_round_kernel/sdpa.cpp | 8 +- .../wrapper/include/cpu_wrapper.hpp | 4 +- .../wrapper/include/sycl_s8_wrapper.hpp | 87 ++++++++ .../wrapper/include/utils.hpp | 192 ++++++++++++------ .../wrapper/include/xpu_wrapper.hpp | 32 ++- auto_round_extension/ark/test/test_matmul.py | 10 +- 7 files changed, 271 insertions(+), 81 deletions(-) create mode 100644 auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index df06dc2c0..3ff619f33 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -21,6 +21,7 @@ typedef uintptr_t torch_ptr; #if ARK_XPU #include #include "xpu_wrapper.hpp" +#include "sycl_s8_wrapper.hpp" #if ARK_SYCL_TLA // Only include declarations, implementations are in separate .cpp files #include "sycl_tla_common.hpp" @@ -32,17 +33,26 @@ typedef uintptr_t torch_ptr; #include "cpu_wrapper.hpp" #endif +#if ARK_DNNL #include "dnnl_wrapper.hpp" +#endif namespace ark { static void matmul(torch_ptr stream, int m, int n, int k, torch_ptr A, int Adt, torch_ptr B, int Bdt, torch_ptr C, int Cdt, torch_ptr bias, bool BT) { +#ifdef ARK_XPU +#if ARK_DNNL auto dt = ark::to_dt((BTLA_DTYPE)Adt); auto cdt = dt; if (Adt == (int)BTLA_DTYPE::S8) cdt = dnnl::memory::data_type::s32; -#ifdef ARK_XPU ark::DnnlWrapper::gemm((sycl::queue*)stream, m, n, k, (void*)A, dt, (void*)B, dt, BT, (void*)C, cdt, (void*)bias); +#elif ARK_SYCL_TLA + ark::sycl_tla_dense_gemm((sycl::queue*)stream, m, n, k, (void*)A, (BTLA_DTYPE)Adt, (void*)B, (BTLA_DTYPE)Bdt, + (void*)C, (BTLA_DTYPE)Cdt, (void*)bias, BT); +#else + throw std::runtime_error("ark::matmul on XPU requires ARK_DNNL=ON or ARK_SYCL_TLA=ON"); +#endif #else CpuWrapper::gemm(m, n, k, (void*)A, (BTLA_DTYPE)Adt, (void*)B, BT, (float*)C, (const float*)bias); #endif @@ -50,9 +60,16 @@ static void matmul(torch_ptr stream, int m, int n, int k, torch_ptr A, int Adt, static void woqgemm_s8(torch_ptr stream, int m, int n, int k, torch_ptr A, int ACdt, torch_ptr B, torch_ptr C, torch_ptr bias, bool BT, torch_ptr scaleb) { +#if ARK_XPU + ark::SyclS8Wrapper::woq_s8((sycl::queue*)stream, m, n, k, (void*)A, (void*)B, BT, (void*)C, + (BTLA_DTYPE)ACdt, (void*)scaleb, (void*)bias, k); +#elif ARK_DNNL auto dt = ark::to_dt((BTLA_DTYPE)ACdt); ark::DnnlWrapper::woq_s8((sycl::queue*)stream, m, n, k, (void*)A, (void*)B, BT, (void*)C, dt, (void*)scaleb, (void*)bias, k); +#else + throw std::runtime_error("ark::woqgemm_s8 requires ARK_XPU or ARK_DNNL"); +#endif } static void woqgemm(torch_ptr stream, int m, int n, int k, torch_ptr A, int ACdt, torch_ptr BlobB, torch_ptr C, diff --git a/auto_round_extension/ark/auto_round_kernel/sdpa.cpp b/auto_round_extension/ark/auto_round_kernel/sdpa.cpp index 10afcf646..2d9b5f424 100644 --- a/auto_round_extension/ark/auto_round_kernel/sdpa.cpp +++ b/auto_round_extension/ark/auto_round_kernel/sdpa.cpp @@ -449,9 +449,9 @@ void sage_prefill_varlen(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, compat::set_default_queue(*q); - // Zero-filled workspace for cu_seqlens_kv_cache via DnnlContext scratch pool. + // Zero-filled workspace for cu_seqlens_kv_cache via DeviceMemoryPool scratch pool. int* zero_cu_buf = static_cast( - DnnlContext::Instance()->get_scratch_mem((batch + 1) * sizeof(int), 3, q)); + DeviceMemoryPool::Instance()->get_scratch_mem((batch + 1) * sizeof(int), 3, q)); q->memset(zero_cu_buf, 0, (batch + 1) * sizeof(int)); options.cu_seqlens_kv_cache = zero_cu_buf; options.use_tensor_strides = true; @@ -512,10 +512,10 @@ void sdpa_varlen_impl(sycl::queue* q, void* Q_ptr, void* K_ptr, void* V_ptr, voi // When isVarLen=true, the kernel's apply_variable_length accesses // cumulative_length for ALL three fields. Even with max_seqlen_kv_cache=0, - // the pointer must be non-null and device-accessible. Use the DnnlContext + // the pointer must be non-null and device-accessible. Use the DeviceMemoryPool // scratch pool (reuses allocation across calls, only grows when needed). int* zero_cu_buf = static_cast( - DnnlContext::Instance()->get_scratch_mem((batch + 1) * sizeof(int), 4, q)); + DeviceMemoryPool::Instance()->get_scratch_mem((batch + 1) * sizeof(int), 4, q)); q->memset(zero_cu_buf, 0, (batch + 1) * sizeof(int)); options.cu_seqlens_kv_cache = zero_cu_buf; options.use_tensor_strides = true; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp index 91f842bb5..559722158 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/cpu_wrapper.hpp @@ -616,7 +616,7 @@ class CpuWrapper { auto bsize = bestla::BTLAGemmNPackBSize(n, k, dt); auto wssize = bsize; wssize += BT ? n * k * bestla::utils::bestla_dtype_bytes(dt) : 0; - auto ptr = DnnlContext::Instance()->get_scratch_mem(wssize, 0, nullptr); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_mem(wssize, 0, nullptr); auto packwptr = ptr; auto bntptr = (int8_t*)packwptr + bsize; auto bptr = B; @@ -631,7 +631,7 @@ class CpuWrapper { static void woq_gemm(int m, const void* a, const void* b, void* c, const void* bias, BTLA_DTYPE acdt, QuantParam* p, size_t blob_count = 0) { auto wssize = bestla::BTLAWOQGemmForwardWorkspace((void*)b, m, p->n, p->k, p->k, p->n, blob_count); - auto ptr = DnnlContext::Instance()->get_scratch_mem(wssize, 0, nullptr); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_mem(wssize, 0, nullptr); bestla::BTLAWOQGemmFp32Forward((void*)b, (const float*)a, (float*)c, (const float*)bias, m, p->n, p->k, p->k, p->n, ptr, get_threading(), blob_count); } diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp new file mode 100644 index 000000000..25e37655a --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp @@ -0,0 +1,87 @@ +#pragma once + +#include "utils.hpp" + +#if ARK_XPU + +namespace ark { + +class SyclS8Wrapper { + public: + static void dyn_quant_s8(sycl::queue* q, int m, int k, const void* a, BTLA_DTYPE adt, int8_t* qa, void* scalea, + int mask) { + if (adt == BTLA_DTYPE::F32) { + using T = float; + using Pro = bestla::sycl_prologue_a::ActivationBase; + Pro::template quant_s8(m, k, mask, {(T*)a, k}, qa, (T*)scalea, q); + } else if (adt == BTLA_DTYPE::F16) { + using T = sycl::half; + using Pro = bestla::sycl_prologue_a::ActivationBase; + Pro::template quant_s8(m, k, mask, {(T*)a, k}, qa, (T*)scalea, q); + } else if (adt == BTLA_DTYPE::BF16) { + using T = sycl::ext::oneapi::bfloat16; + using Pro = bestla::sycl_prologue_a::ActivationBase; + Pro::template quant_s8(m, k, mask, {(T*)a, k}, qa, (T*)scalea, q); + } else { + throw std::invalid_argument("SyclS8Wrapper::dyn_quant_s8: unsupported activation dtype"); + } + } + + static void igemm_s8s8(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE ct, void* scale_a, void* scale_b, void* bias, int blocksize) { + if (!BT) { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: only B as n x k is supported"); + } + + using namespace bestla::sycl_gemm; + + if (blocksize == k || blocksize == -1) { + if (ct == BTLA_DTYPE::F32) { + using T = float; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else if (ct == BTLA_DTYPE::F16) { + using T = sycl::half; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else if (ct == BTLA_DTYPE::BF16) { + using T = sycl::ext::oneapi::bfloat16; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: unsupported output dtype"); + } + return; + } + + if (ct == BTLA_DTYPE::F32) { + using T = float; + Launcher, xmx::IKblockGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b, blocksize}); + } else if (ct == BTLA_DTYPE::F16) { + using T = sycl::half; + Launcher, xmx::IKblockGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b, blocksize}); + } else { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: k-block path supports only F32/F16 output"); + } + } + + static void woq_s8(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE act, void* scale_b, void* bias, int blocksize) { + size_t qa_size = size_t(m) * size_t(k); + size_t scalea_offset = (qa_size + alignof(float) - 1) & ~(size_t(alignof(float)) - 1); + size_t tmp_size = scalea_offset + size_t(m) * sizeof(float); + + auto tmp_ptr = static_cast(DeviceMemoryPool::Instance()->get_scratch_mem(tmp_size, 1, q)); + auto qa_ptr = tmp_ptr; + auto scalea_ptr = tmp_ptr + scalea_offset; + + dyn_quant_s8(q, m, k, a, act, qa_ptr, scalea_ptr, 0); + igemm_s8s8(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); + } +}; + +} // namespace ark + +#endif // ARK_XPU \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp index bbf2f2bad..ac182fe2a 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/utils.hpp @@ -11,20 +11,34 @@ // #pragma once -#include + #include +#include +#include +#include +#include +#include +#include #include + +#if ARK_DNNL #include +#endif + #include "bestla/bestla_wrapper.h" + #if ARK_XPU #include "bestla/sycl/sycl_wrapper.h" -#include #else namespace sycl { typedef void queue; } #endif +#if ARK_XPU && ARK_DNNL +#include +#endif + #define LOG_LINE() printf("%s:L%d\n", __FUNCTION__, __LINE__); namespace ark { @@ -35,10 +49,12 @@ struct env_params { int sage_use_mean_bias = 1; int sage_print_kbias = 0; int sage_disable_packed_hnd_fast = 0; + static env_params* Instance() { static env_params instance; return &instance; } + env_params() { env_i("ARK_VERBOSE", verbose); env_i("ARK_AUTO_S8", auto_s8); @@ -46,22 +62,21 @@ struct env_params { env_i("ARK_SAGE_PRINT_KBIAS", sage_print_kbias); env_i("ARK_SAGE_DISABLE_PACKED_HND_FAST", sage_disable_packed_hnd_fast); } + static inline void env_i(const char* envstr, int& default_) { - const char* log_level_env = getenv(envstr); + const char* log_level_env = std::getenv(envstr); if (log_level_env != nullptr) default_ = std::stoi(log_level_env); } }; using UUIDArray = std::array; -// 高性能 UUID 哈希器 + struct UUIDHasher { size_t operator()(const UUIDArray& uuid) const { - // 将 16 字节视为两个 64 位整数进行处理，性能远高于逐字节遍历 const uint64_t* p = reinterpret_cast(uuid.data()); uint64_t h1 = p[0]; uint64_t h2 = p[1]; - // 使用类似 splitmix64 的混合逻辑 h1 ^= h1 >> 33; h1 *= 0xff51afd7ed558ccdLLU; h2 ^= h2 >> 33; @@ -71,20 +86,103 @@ struct UUIDHasher { } }; +class DeviceMemoryPool { + public: + static DeviceMemoryPool* Instance() { + static DeviceMemoryPool instance; + return &instance; + } + + size_t get_device_key(sycl::queue* q) { +#if ARK_XPU + if (q != nullptr) { + auto uuid = q->get_device().get_info(); + return UUIDHasher{}(uuid); + } +#endif + return 0; + } + + void* get_scratch_mem(size_t size, size_t buf_loc, sycl::queue* q) { + auto key = get_device_key(q); + return get_scratch_ptr(size, buf_loc, q, key); + } + + void* get_scratch_ptr(size_t size, size_t buf_loc, sycl::queue* q, size_t key) { + if (size == 0 || buf_loc >= MaxLocNum) return nullptr; + + auto it = dev_mem_ptr_map[buf_loc].find(key); + if (it == dev_mem_ptr_map[buf_loc].end()) { + auto newptr = allocate(size, q); + dev_mem_size_map[buf_loc][key] = size; + dev_mem_ptr_map[buf_loc][key] = newptr; + return newptr; + } + + auto old_size = dev_mem_size_map[buf_loc][key]; + if (old_size < size) { + release(it->second, q); + auto newptr = allocate(size, q); + dev_mem_size_map[buf_loc][key] = size; + dev_mem_ptr_map[buf_loc][key] = newptr; + return newptr; + } + + return it->second; + } + + private: + static constexpr int MaxLocNum = 8; + using SizeMap = std::unordered_map; + using PtrMap = std::unordered_map; + + int8_t* allocate(size_t size, sycl::queue* q) { +#if ARK_XPU + if (q == nullptr) { + throw std::invalid_argument("DeviceMemoryPool: XPU allocation requires a non-null SYCL queue"); + } + return sycl::aligned_alloc_device(128, size, *q); +#else + return static_cast(std::malloc(size)); +#endif + } + + void release(void* ptr, sycl::queue* q) { + if (ptr == nullptr) return; +#if ARK_XPU + if (q == nullptr) { + throw std::invalid_argument("DeviceMemoryPool: XPU free requires a non-null SYCL queue"); + } + sycl::free(ptr, *q); +#else + std::free(ptr); +#endif + } + + std::array dev_mem_size_map; + std::array dev_mem_ptr_map; +}; + +#if ARK_DNNL + +template +struct always_false : std::false_type {}; + template static inline constexpr dnnl::memory::data_type to_dt() { - if constexpr (std::is_same_v) + if constexpr (std::is_same_v) { return dnnl::memory::data_type::f32; - else if constexpr (std::is_same_v) + } else if constexpr (std::is_same_v) { return dnnl::memory::data_type::f16; - else if constexpr (std::is_same_v) + } else if constexpr (std::is_same_v) { return dnnl::memory::data_type::s8; - else if constexpr (std::is_same_v) + } else if constexpr (std::is_same_v) { return dnnl::memory::data_type::u8; - else if constexpr (std::is_same_v) + } else if constexpr (std::is_same_v) { return dnnl::memory::data_type::bf16; - else - static_assert(0); + } else { + static_assert(always_false::value, "unsupported dnnl dtype"); + } } static inline constexpr dnnl::memory::data_type to_dt(BTLA_DTYPE bt) { @@ -123,6 +221,7 @@ class DnnlContext { size_t check_dnnl_device(sycl::queue* q) { size_t key = 0; + if (q == nullptr) { if (dev_engine_map.find(key) == dev_engine_map.end()) { dev_engine_map[key] = dnnl::engine(dnnl::engine::kind::cpu, 0); @@ -130,75 +229,46 @@ class DnnlContext { } return key; } + #if ARK_XPU - auto uuid = q->get_device().get_info(); - key = UUIDHasher{}(uuid); + key = DeviceMemoryPool::Instance()->get_device_key(q); if (dev_engine_map.find(key) == dev_engine_map.end()) { sycl::device dev = q->get_device(); - // Get the context associated with the queue sycl::context ctx = q->get_context(); dev_engine_map[key] = dnnl::sycl_interop::make_engine(dev, ctx); dev_stream_map[key] = dnnl::sycl_interop::make_stream(dev_engine_map[key], *q); } +#else + if (dev_engine_map.find(key) == dev_engine_map.end()) { + dev_engine_map[key] = dnnl::engine(dnnl::engine::kind::cpu, 0); + dev_stream_map[key] = dnnl::stream(dev_engine_map[key]); + } #endif + return key; } - dnnl::memory get_scratch_mem(dnnl::memory::desc _md, sycl::queue* q) { + dnnl::memory get_scratch_mem(dnnl::memory::desc md, sycl::queue* q) { auto key = check_dnnl_device(q); - auto ptr = get_scratch_ptr(_md.get_size(), 0, q, key); - return dnnl::memory(_md, dev_engine_map[key], ptr); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_ptr(md.get_size(), 0, q, key); + return dnnl::memory(md, dev_engine_map[key], ptr); } - void* get_scratch_mem(size_t _size, size_t buf_loc, sycl::queue* q) { - auto key = check_dnnl_device(q); - auto ptr = get_scratch_ptr(_size, buf_loc, q, key); - return ptr; + void* get_scratch_mem(size_t size, size_t buf_loc, sycl::queue* q) { + return DeviceMemoryPool::Instance()->get_scratch_mem(size, buf_loc, q); } - void* get_scratch_ptr(size_t _size, size_t buf_loc, sycl::queue* q, size_t key) { - if (_size <= 0 || buf_loc >= MaxLocNum) return nullptr; - - auto it = dev_mem_ptr_map[buf_loc].find(key); - if (it == dev_mem_ptr_map[buf_loc].end()) { - dev_mem_size_map[buf_loc][key] = _size; -#if ARK_XPU - auto newptr = sycl::aligned_alloc_device(128, _size, *q); -#else - auto newptr = (int8_t*)malloc(_size); -#endif - dev_mem_ptr_map[buf_loc][key] = newptr; - return newptr; - } else { - auto pre_size = dev_mem_size_map[buf_loc][key]; - if (pre_size < _size) { - auto pre_ptr = it->second; -#if ARK_XPU - sycl::free(pre_ptr, *q); -#else - free(pre_ptr); -#endif - dev_mem_size_map[buf_loc][key] = _size; -#if ARK_XPU - auto newptr = sycl::aligned_alloc_device(128, _size, *q); -#else - auto newptr = (int8_t*)malloc(_size); -#endif - dev_mem_ptr_map[buf_loc][key] = newptr; - return newptr; - } - } - return it->second; + void* get_scratch_ptr(size_t size, size_t buf_loc, sycl::queue* q, size_t key) { + return DeviceMemoryPool::Instance()->get_scratch_ptr(size, buf_loc, q, key); } - typedef std::unordered_map msize_t; - typedef std::unordered_map mptr_t; - static constexpr int MaxLocNum = 8; - std::array dev_mem_size_map; - std::array dev_mem_ptr_map; + + private: std::unordered_map dev_engine_map; std::unordered_map dev_stream_map; }; +#endif // ARK_DNNL + struct QuantParam { int n; int k; diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp index 09fd1ccac..f67774ab5 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp @@ -17,7 +17,18 @@ #include #include #include "utils.hpp" +#if ARK_DNNL #include "dnnl_wrapper.hpp" +#endif + +#if ARK_SYCL_TLA +#include "sycl_tla_dense_gemm.hpp" +#endif + +#if ARK_XPU +#include "sycl_s8_wrapper.hpp" +#endif + #include "sycl_tla_common.hpp" namespace ark { @@ -696,14 +707,19 @@ class XpuWrapper { } auto ret = woq_gemv(q, m, p, a, b, c, bias, acdt); if (ret) { - auto dnnl_dt = to_dt(acdt); + check_compute_type(p); if (p->compute_type != BTLA_DTYPE::S8) { size_t elesize = bestla::utils::bestla_dtype_bytes(acdt); size_t total_size = elesize * p->k * p->n; - auto ptr = DnnlContext::Instance()->get_scratch_mem(total_size, 1, q); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_mem(total_size, 1, q); unpackq(acdt, (int8_t*)b, ptr, p, q); +#if ARK_DNNL + auto dnnl_dt = to_dt(acdt); DnnlWrapper::gemm(q, m, p->n, p->k, a, dnnl_dt, ptr, dnnl_dt, true, c, dnnl_dt, bias); +#elif ARK_SYCL_TLA + ark::sycl_tla_dense_gemm(q, m, p->n, p->k, a, acdt, ptr, acdt, c, acdt, bias, true); +#endif } else { auto bptr = (int8_t*)b; bool _rescale = rescale(p); @@ -713,12 +729,12 @@ class XpuWrapper { } else { size_t elesize = 1; size_t total_size = elesize * p->k * p->n; - bptr = (int8_t*)DnnlContext::Instance()->get_scratch_mem(total_size, 2, q); + bptr = (int8_t*)DeviceMemoryPool::Instance()->get_scratch_mem(total_size, 2, q); unpackq(BTLA_DTYPE::S8, (int8_t*)b, bptr, p, q); scaleb_ptr = _rescale ? (int8_t*)b + get_scalext_offset(p) : (int8_t*)b + get_scale_offset(p); } auto blocksize = rescale_blocksize(p); - DnnlWrapper::woq_s8(q, m, p->n, p->k, a, bptr, true, c, dnnl_dt, scaleb_ptr, (void*)bias, blocksize); + ark::SyclS8Wrapper::woq_s8(q, m, p->n, p->k, a, bptr, true, c, acdt, scaleb_ptr, (void*)bias, blocksize); } } } @@ -819,7 +835,7 @@ class XpuWrapper { template static void print_value_distribution(sycl::queue* q, const T* dev_ptr, size_t count, const char* name) { if (dev_ptr == nullptr || count == 0) { - std::printf("[%s] empty\n", name); + std::fprintf(stdout, "[%s] empty\n", name); return; } @@ -838,7 +854,7 @@ class XpuWrapper { } double mean = sum / static_cast(count); - std::printf("[%s] min=%f max=%f mean=%f\n", name, min_value, max_value, static_cast(mean)); + std::fprintf(stdout, "[%s] min=%f max=%f mean=%f\n", name, min_value, max_value, static_cast(mean)); } // input: num_rows x head_dim matrix, output: int8 quantized matrix + scale (per block) @@ -1495,7 +1511,7 @@ class XpuWrapper { size_t k_bias_size = use_mean_bias ? num_heads_kv * head_dim * batch * sizeof(T) : 0; size_t total_size = k_size + q_size + k_scale_size * sizeof(float) + q_scale_size * sizeof(float) + v_tmp_size + v_scale_size + k_bias_size; - auto ptr = DnnlContext::Instance()->get_scratch_mem(total_size, 1, q); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_mem(total_size, 1, q); auto q_out_ptr = (int8_t*)ptr; auto k_out_ptr = (int8_t*)ptr + q_size; auto qscale = (float*)((int8_t*)ptr + q_size + k_size); @@ -1598,7 +1614,7 @@ class XpuWrapper { size_t k_bias_size = use_mean_bias ? num_heads_kv * head_dim * sizeof(T) : 0; size_t total_size = k_size + q_size + k_scale_size * sizeof(float) + q_scale_size * sizeof(float) + v_tmp_size + v_scale_size + k_bias_size; - auto ptr = DnnlContext::Instance()->get_scratch_mem(total_size, 1, q); + auto ptr = DeviceMemoryPool::Instance()->get_scratch_mem(total_size, 1, q); auto q_out_ptr = (int8_t*)ptr; auto k_out_ptr = (int8_t*)ptr + q_size; auto qscale = (float*)((int8_t*)ptr + q_size + k_size); diff --git a/auto_round_extension/ark/test/test_matmul.py b/auto_round_extension/ark/test/test_matmul.py index 4d4aee98c..bdb52c18f 100644 --- a/auto_round_extension/ark/test/test_matmul.py +++ b/auto_round_extension/ark/test/test_matmul.py @@ -201,7 +201,7 @@ def test_cpu(m, k, n, dt, batch_size, runs, record_property): @pytest.mark.parametrize("m", [1, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("batch_size", [1]) @pytest.mark.parametrize("runs", [1]) def test_xpu_sycl_tla(m, k, n, dt, batch_size, runs, record_property): @@ -231,7 +231,7 @@ def test_xpu_sycl_tla(m, k, n, dt, batch_size, runs, record_property): @pytest.mark.parametrize("m", [1, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) @pytest.mark.parametrize("batch_size", [1]) @pytest.mark.parametrize("runs", [1]) def test_xpu_sycl_tla_no_bias(m, k, n, dt, batch_size, runs, record_property): @@ -262,7 +262,7 @@ def test_xpu_sycl_tla_no_bias(m, k, n, dt, batch_size, runs, record_property): # pytest -vs auto_round_extension/ark/test/test_matmul.py -k compare_dnnl_vs_sycl_tla @pytest.mark.parametrize("m", [1, 8, 16, 32, 128, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) def test_xpu_compare_dnnl_vs_sycl_tla(m, k, n, dt): warmup = 100 runs = 1000 @@ -273,7 +273,7 @@ def test_xpu_compare_dnnl_vs_sycl_tla(m, k, n, dt): @pytest.mark.parametrize("m", [1, 8, 16, 32, 128, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) def test_xpu_compare_dnnl_vs_sycl_tla_no_bias(m, k, n, dt): warmup = 100 runs = 1000 @@ -288,7 +288,7 @@ def compare_matmul_backends(m, k, n, dt, warmup, runs, device="xpu", has_bias=Tr def _matmul_tolerance(dt): if dt == torch.float32: - return 0.001, 0.03 + return 0.1, 0.06 if dt == torch.float16: return 0.1, 0.06 if dt == torch.bfloat16: From c6076c32e672b2a7e79136cd3f60ebb2e17296fd Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Tue, 14 Jul 2026 15:02:57 +0800 Subject: [PATCH 4/8] modify the enable_dnnl variable Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- auto_round_extension/ark/setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/auto_round_extension/ark/setup.py b/auto_round_extension/ark/setup.py index 7bc781f6d..4bda0c841 100644 --- a/auto_round_extension/ark/setup.py +++ b/auto_round_extension/ark/setup.py @@ -99,7 +99,7 @@ def env_flag(name, default=False): requirements = fetch_requirements("requirements.txt") enable_sycl_tla = parse_major_minor(oneapi_version) >= (2025, 3) -enable_dnnl = env_flag("ARK_DNNL", False) +enable_dnnl = env_flag("ARK_DNNL", default=not enable_sycl_tla) def get_system_memory_gb(): From 0f86148b305192c4c077e97c1238b662554b3b3d Mon Sep 17 00:00:00 2001 From: Zhenzhong Xu Date: Tue, 14 Jul 2026 15:05:51 +0800 Subject: [PATCH 5/8] add license Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../wrapper/include/sycl_s8_wrapper.hpp | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp index 25e37655a..2c30a85c7 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp @@ -1,3 +1,15 @@ +// +// MIT license +// Copyright (C) 2026 Intel Corporation +// SPDX-License-Identifier: MIT +// + +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// + #pragma once #include "utils.hpp" From d13c29119c48cbb50948d5b8a7b00f1405463370 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Thu, 16 Jul 2026 15:43:20 +0800 Subject: [PATCH 6/8] sycltla igemm_s8s8 inference pass Signed-off-by: Zhenzhong1 --- .../wrapper/include/sycl_s8_wrapper.hpp | 9 +- .../wrapper/include/sycl_tla_s8_gemm.hpp | 231 ++++++++++++++++++ 2 files changed, 239 insertions(+), 1 deletion(-) create mode 100644 auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp index 2c30a85c7..aeb78e989 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp @@ -9,13 +9,16 @@ // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // - #pragma once #include "utils.hpp" #if ARK_XPU +#if ARK_SYCL_TLA +#include "sycl_tla_s8_gemm.hpp" +#endif + namespace ark { class SyclS8Wrapper { @@ -45,6 +48,9 @@ class SyclS8Wrapper { throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: only B as n x k is supported"); } +#if ARK_SYCL_TLA + ark::sycl_tla_igemm_s8s8_dequant(q, m, n, k, a, b, c, ct, scale_a, scale_b, bias, blocksize); +#else using namespace bestla::sycl_gemm; if (blocksize == k || blocksize == -1) { @@ -77,6 +83,7 @@ class SyclS8Wrapper { } else { throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: k-block path supports only F32/F16 output"); } +#endif } static void woq_s8(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, 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 new file mode 100644 index 000000000..4cfcb175f --- /dev/null +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_s8_gemm.hpp @@ -0,0 +1,231 @@ +#pragma once + +#include +#include + +#ifdef ARK_XPU +#include +#endif + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) +#include +#include "cute/tensor.hpp" +#include "cute/util/compat.hpp" +#endif + +#include "utils.hpp" + +namespace ark { + +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +namespace sycl_tla_s8_detail { + +using namespace cute; + +template +class S8DequantKernelName; + +template +class S8FinalizeKernelName; + +template +void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma, ElementOut* c, float* accum, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int m, int n, + int block_idx, int scale_b_stride) { + auto item = sycl::ext::oneapi::this_work_item::get_nd_item<2>(); + int wg_m = int(item.get_group(1)); + int wg_n = int(item.get_group(0)); + int local_id = int(item.get_local_id(0)); + + Tensor cA = make_identity_tensor(A.shape()); + Tensor cB = make_identity_tensor(B.shape()); + Tensor cC = make_identity_tensor(make_shape(m, n)); + + auto wg_tile = mma.tile_mnk(); + auto wg_coord = make_coord(wg_m, wg_n, 0); + + Tensor gA = local_tile(cA, select<0, 2>(wg_tile), make_coord(wg_m, _)); + Tensor gB = local_tile(cB, select<1, 2>(wg_tile), make_coord(wg_n, _)); + Tensor gC = local_tile(cC, wg_tile, wg_coord, Step<_1, _1, X>{}); + + auto copy_a = make_block_2d_copy_A(mma, A); + auto copy_b = make_block_2d_copy_B(mma, B); + + auto thr_mma = mma.get_slice(local_id); + auto thr_copy_a = copy_a.get_slice(local_id); + auto thr_copy_b = copy_b.get_slice(local_id); + + auto tCrA = thr_mma.partition_sg_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_sg_fragment_B(gB(_, _, 0)); + + auto tArA = thr_copy_a.partition_sg_fragment_D(gA(_, _, 0)); + auto tBrB = thr_copy_b.partition_sg_fragment_D(gB(_, _, 0)); + + Tensor tAgA = thr_copy_a.partition_S(gA); + Tensor tBgB = thr_copy_b.partition_S(gB); + + Tensor tCrC = partition_fragment_C(mma, select<0, 1>(wg_tile)); + Tensor tCgC = thr_mma.partition_C(gC); + + constexpr SPIRVScope barrier_scope = ScopeWorkgroup; + int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); + + clear(tCrC); + + for (int k_tile = 0; k_tile < k_tile_count; ++k_tile) { + barrier_arrive(barrier_scope); + + copy(copy_a, tAgA(_, _, _, k_tile), tArA); + copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + + reorder(tArA, tCrA); + reorder(tBrB, tCrB); + gemm(mma, tCrA, tCrB, tCrC); + + barrier_wait(barrier_scope); + } + + CUTE_UNROLL + for (int i = 0; i < size(tCrC); ++i) { + auto coord = tCgC(i); + int row = int(get<0>(coord)); + int col = int(get<1>(coord)); + if (row >= m || col >= n) continue; + + float acc = static_cast(tCrC(i)); + if constexpr (AccumBlock) { + float sb = static_cast(scale_b[col * scale_b_stride + block_idx]); + accum[row * n + col] += acc * sb; + } else { + float sa = static_cast(scale_a[row]); + float sb = static_cast(scale_b[col]); + float value = acc * sa * sb; + if (bias) value += static_cast(bias[col]); + c[row * n + col] = static_cast(value); + } + } +} + +template +void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, const int8_t* a, const int8_t* b, + ElementOut* c, float* accum, const ElementOut* scale_a, const ElementOut* scale_b, + const ElementOut* bias, int block_idx, int scale_b_stride) { + compat::set_default_queue(*q); + + auto A = make_tensor(make_gmem_ptr(const_cast(a)), make_shape(m, gemm_k), make_stride(lda, _1{})); + auto B = make_tensor(make_gmem_ptr(const_cast(b)), make_shape(n, gemm_k), make_stride(ldb, _1{})); + + using Op = XE_DPAS_TT<8, int32_t, int8_t, int8_t>; + using WGTile = Shape, Int, _64>; + using MMA = typename TiledMMAHelper, Layout, SGLayout>::TiledMMA; + MMA mma{}; + + sycl::range<2> local = {size(mma), 1}; + sycl::range<2> global = {local[0] * ceil_div(n, get<1>(mma.tile_mnk())), + local[1] * ceil_div(m, get<0>(mma.tile_mnk()))}; + + namespace syclex = sycl::ext::oneapi::experimental; + namespace intelex = sycl::ext::intel::experimental; + syclex::properties props{syclex::sub_group_size<16>, intelex::grf_size<256>}; + + q->parallel_for>( + sycl::nd_range<2>(global, local), props, [=](auto) { + igemm_device_impl(A, B, mma, c, accum, scale_a, scale_b, bias, m, n, block_idx, scale_b_stride); + }); +} + +template +void launch_igemm(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, const int8_t* a, const int8_t* b, + ElementOut* c, float* accum, const ElementOut* scale_a, const ElementOut* scale_b, + const ElementOut* bias, int block_idx, int scale_b_stride) { + using SGLayout = Layout, Stride<_4, _1, _0>>; + launch_igemm_tile( + q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); +} + +template +void finalize_block_output(sycl::queue* q, int m, int n, const float* accum, ElementOut* c, + const ElementOut* scale_a, const ElementOut* bias) { + q->parallel_for>(sycl::range<1>(size_t(m) * size_t(n)), [=](sycl::id<1> id) { + size_t idx = id[0]; + int row = int(idx / n); + int col = int(idx - size_t(row) * size_t(n)); + float value = accum[idx] * static_cast(scale_a[row]); + if (bias) value += static_cast(bias[col]); + c[idx] = static_cast(value); + }); +} + +template +void run_typed(sycl::queue* q, int m, int n, int k, const int8_t* a, const int8_t* b, ElementOut* c, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int blocksize) { + bool k_block = !(blocksize == k || blocksize == -1); + + if (!k_block) { + launch_igemm(q, m, n, k, k, k, a, b, c, nullptr, scale_a, scale_b, bias, 0, 1); + return; + } + + if (blocksize <= 0 || k % blocksize != 0) { + throw std::invalid_argument("sycl_tla_igemm_s8s8_dequant: blocksize must divide k"); + } + + int blks = k / blocksize; + size_t bytes = size_t(m) * size_t(n) * sizeof(float); + auto* accum = static_cast(DeviceMemoryPool::Instance()->get_scratch_mem(bytes, 3, q)); + q->memset(accum, 0, bytes); + + for (int ib = 0; ib < blks; ++ib) { + const int8_t* a_blk = a + size_t(ib) * size_t(blocksize); + const int8_t* b_blk = b + size_t(ib) * size_t(blocksize); + launch_igemm(q, m, n, blocksize, k, k, a_blk, b_blk, nullptr, accum, nullptr, scale_b, nullptr, + ib, blks); + } + + finalize_block_output(q, m, n, accum, c, scale_a, bias); +} + +} // namespace sycl_tla_s8_detail + +inline void sycl_tla_igemm_s8s8_dequant(sycl::queue* q, int m, int n, int k, const void* a, const void* b, void* c, + BTLA_DTYPE ct, const void* scale_a, const void* scale_b, const void* bias, + int blocksize) { + if (!q) throw std::invalid_argument("sycl_tla_igemm_s8s8_dequant: queue must not be null"); + if (!a || !b || !c || !scale_a || !scale_b) { + throw std::invalid_argument("sycl_tla_igemm_s8s8_dequant: input pointers must not be null"); + } + if (m <= 0 || n <= 0 || k <= 0) return; + + bool k_block = !(blocksize == k || blocksize == -1); + + switch (ct) { + case BTLA_DTYPE::F32: + sycl_tla_s8_detail::run_typed(q, m, n, k, static_cast(a), static_cast(b), + static_cast(c), static_cast(scale_a), + static_cast(scale_b), static_cast(bias), blocksize); + return; + case BTLA_DTYPE::F16: + sycl_tla_s8_detail::run_typed(q, m, n, k, static_cast(a), static_cast(b), + static_cast(c), static_cast(scale_a), + static_cast(scale_b), static_cast(bias), + blocksize); + return; + case BTLA_DTYPE::BF16: + if (k_block) { + throw std::invalid_argument("sycl_tla_igemm_s8s8_dequant: k-block path supports only F32/F16 output"); + } + sycl_tla_s8_detail::run_typed(q, m, n, k, static_cast(a), static_cast(b), + static_cast(c), + static_cast(scale_a), + static_cast(scale_b), + static_cast(bias), blocksize); + return; + default: + throw std::invalid_argument("sycl_tla_igemm_s8s8_dequant: unsupported output dtype"); + } +} + +#endif // ARK_XPU && ARK_SYCL_TLA + +} // namespace ark \ No newline at end of file From bc7476ea3a9bd01a897c1e92cb0ad7820e8bbd85 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Thu, 16 Jul 2026 16:30:35 +0800 Subject: [PATCH 7/8] add igemm_sycltla&joint_matrix Signed-off-by: Zhenzhong1 --- .../ark/auto_round_kernel/__init__.py | 50 +++++++++++ .../ark/auto_round_kernel/ark.cpp | 18 ++++ .../wrapper/include/sycl_s8_wrapper.hpp | 85 ++++++++++++++++++- auto_round_extension/ark/test/test_matmul.py | 56 +++++++++++- 4 files changed, 206 insertions(+), 3 deletions(-) diff --git a/auto_round_extension/ark/auto_round_kernel/__init__.py b/auto_round_extension/ark/auto_round_kernel/__init__.py index ee37187ed..c3ee9f676 100644 --- a/auto_round_extension/ark/auto_round_kernel/__init__.py +++ b/auto_round_extension/ark/auto_round_kernel/__init__.py @@ -369,6 +369,56 @@ def woqgemm_s8(A: torch.Tensor, B: torch.Tensor, scaleB: torch.Tensor, bias: tor return C +def woqgemm_s8_sycl_tla(A: torch.Tensor, B: torch.Tensor, scaleB: torch.Tensor, bias: torch.Tensor): + m = A.shape[0] + n = B.shape[0] + k = B.shape[1] + lib = get_lib(A) + + C = torch.zeros(m, n, dtype=A.dtype, device=A.device) + stream = get_stream(A) + lib.woqgemm_s8_sycl_tla( + stream, + m, + n, + k, + A.contiguous().data_ptr(), + cvt_dtype(A.dtype), + B.contiguous().data_ptr(), + C.contiguous().data_ptr(), + bias.contiguous().data_ptr(), + True, + scaleB.contiguous().data_ptr(), + ) + return C + + +# A: mxk:DT, B: nxk:s8, scaleB: n:DT +# return: mxn:DT +def woqgemm_s8_joint_matrix(A: torch.Tensor, B: torch.Tensor, scaleB: torch.Tensor, bias: torch.Tensor): + m = A.shape[0] + n = B.shape[0] + k = B.shape[1] + lib = get_lib(A) + + C = torch.zeros(m, n, dtype=A.dtype, device=A.device) + stream = get_stream(A) + lib.woqgemm_s8_joint_matrix( + stream, + m, + n, + k, + A.contiguous().data_ptr(), + cvt_dtype(A.dtype), + B.contiguous().data_ptr(), + C.contiguous().data_ptr(), + bias.contiguous().data_ptr(), + True, + scaleB.contiguous().data_ptr(), + ) + return C + + # A: mxk:DT, B: BS:s8, bias: n:DT # return: C: mxn:DT def woqgemm( diff --git a/auto_round_extension/ark/auto_round_kernel/ark.cpp b/auto_round_extension/ark/auto_round_kernel/ark.cpp index 3ff619f33..262bc1e57 100755 --- a/auto_round_extension/ark/auto_round_kernel/ark.cpp +++ b/auto_round_extension/ark/auto_round_kernel/ark.cpp @@ -58,6 +58,22 @@ static void matmul(torch_ptr stream, int m, int n, int k, torch_ptr A, int Adt, #endif } +#if defined(ARK_XPU) && defined(ARK_SYCL_TLA) + +static void woqgemm_s8_joint_matrix(torch_ptr stream, int m, int n, int k, torch_ptr A, int ACdt, + torch_ptr B, torch_ptr C, torch_ptr bias, bool BT, torch_ptr scaleb) { + ark::SyclS8Wrapper::woq_s8_joint_matrix((sycl::queue*)stream, m, n, k, (void*)A, (void*)B, BT, (void*)C, + (BTLA_DTYPE)ACdt, (void*)scaleb, (void*)bias, k); +} + +static void woqgemm_s8_sycl_tla(torch_ptr stream, int m, int n, int k, torch_ptr A, int ACdt, + torch_ptr B, torch_ptr C, torch_ptr bias, bool BT, torch_ptr scaleb) { + ark::SyclS8Wrapper::woq_s8_sycl_tla((sycl::queue*)stream, m, n, k, (void*)A, (void*)B, BT, (void*)C, + (BTLA_DTYPE)ACdt, (void*)scaleb, (void*)bias, k); +} + +#endif + static void woqgemm_s8(torch_ptr stream, int m, int n, int k, torch_ptr A, int ACdt, torch_ptr B, torch_ptr C, torch_ptr bias, bool BT, torch_ptr scaleb) { #if ARK_XPU @@ -694,5 +710,7 @@ PYBIND11_MODULE(PY_NAME, m) { m.def("sage_dynamic_quant_v_layout", &ark::sage_dynamic_quant_v_layout); m.def("moe_gemm", &ark::moe_gemm_wrapper); m.def("matmul_sycl_tla", &ark::matmul_sycl_tla); + m.def("woqgemm_s8_joint_matrix", &ark::woqgemm_s8_joint_matrix); + m.def("woqgemm_s8_sycl_tla", &ark::woqgemm_s8_sycl_tla); #endif } \ No newline at end of file diff --git a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp index aeb78e989..b26a28a1a 100644 --- a/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp +++ b/auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_s8_wrapper.hpp @@ -86,6 +86,56 @@ class SyclS8Wrapper { #endif } + static void igemm_s8s8_sycl_tla(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE ct, void* scale_a, void* scale_b, void* bias, int blocksize) { + if (!BT) { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: only B as n x k is supported"); + } + + ark::sycl_tla_igemm_s8s8_dequant(q, m, n, k, a, b, c, ct, scale_a, scale_b, bias, blocksize); + } + + + static void igemm_s8s8_joint_matrix(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE ct, void* scale_a, void* scale_b, void* bias, int blocksize) { + if (!BT) { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: only B as n x k is supported"); + } + + using namespace bestla::sycl_gemm; + + if (blocksize == k || blocksize == -1) { + if (ct == BTLA_DTYPE::F32) { + using T = float; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else if (ct == BTLA_DTYPE::F16) { + using T = sycl::half; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else if (ct == BTLA_DTYPE::BF16) { + using T = sycl::ext::oneapi::bfloat16; + Launcher, xmx::IGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b}); + } else { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: unsupported output dtype"); + } + return; + } + + if (ct == BTLA_DTYPE::F32) { + using T = float; + Launcher, xmx::IKblockGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b, blocksize}); + } else if (ct == BTLA_DTYPE::F16) { + using T = sycl::half; + Launcher, xmx::IKblockGemmDQCore>::run( + q, {(void*)a, (void*)b, c, m, n, k, k, k, n, bias, scale_a, scale_b, blocksize}); + } else { + throw std::invalid_argument("SyclS8Wrapper::igemm_s8s8: k-block path supports only F32/F16 output"); + } + } + static void woq_s8(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, BTLA_DTYPE act, void* scale_b, void* bias, int blocksize) { size_t qa_size = size_t(m) * size_t(k); @@ -97,8 +147,41 @@ class SyclS8Wrapper { auto scalea_ptr = tmp_ptr + scalea_offset; dyn_quant_s8(q, m, k, a, act, qa_ptr, scalea_ptr, 0); - igemm_s8s8(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); +#if ARK_SYCL_TLA + igemm_s8s8_sycl_tla(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); +#else + igemm_s8s8_joint_matrix(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); +#endif + } + + static void woq_s8_joint_matrix(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE act, void* scale_b, void* bias, int blocksize) { + size_t qa_size = size_t(m) * size_t(k); + size_t scalea_offset = (qa_size + alignof(float) - 1) & ~(size_t(alignof(float)) - 1); + size_t tmp_size = scalea_offset + size_t(m) * sizeof(float); + + auto tmp_ptr = static_cast(DeviceMemoryPool::Instance()->get_scratch_mem(tmp_size, 1, q)); + auto qa_ptr = tmp_ptr; + auto scalea_ptr = tmp_ptr + scalea_offset; + + dyn_quant_s8(q, m, k, a, act, qa_ptr, scalea_ptr, 0); + igemm_s8s8_joint_matrix(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); } + + static void woq_s8_sycl_tla(sycl::queue* q, int m, int n, int k, const void* a, const void* b, bool BT, void* c, + BTLA_DTYPE act, void* scale_b, void* bias, int blocksize) { + size_t qa_size = size_t(m) * size_t(k); + size_t scalea_offset = (qa_size + alignof(float) - 1) & ~(size_t(alignof(float)) - 1); + size_t tmp_size = scalea_offset + size_t(m) * sizeof(float); + + auto tmp_ptr = static_cast(DeviceMemoryPool::Instance()->get_scratch_mem(tmp_size, 1, q)); + auto qa_ptr = tmp_ptr; + auto scalea_ptr = tmp_ptr + scalea_offset; + + dyn_quant_s8(q, m, k, a, act, qa_ptr, scalea_ptr, 0); + igemm_s8s8_sycl_tla(q, m, n, k, qa_ptr, b, BT, c, act, scalea_ptr, scale_b, bias, blocksize); + } + }; } // namespace ark diff --git a/auto_round_extension/ark/test/test_matmul.py b/auto_round_extension/ark/test/test_matmul.py index bdb52c18f..52c66773d 100644 --- a/auto_round_extension/ark/test/test_matmul.py +++ b/auto_round_extension/ark/test/test_matmul.py @@ -164,6 +164,58 @@ def woqgemm(m, k, n, dt, batch_size, runs, record_property, device): print(f"[Performance] Time: {dur*1000:.4f} ms, GFLOPS: {gflops:.2f}, Bandwidth: {bandwidth:.2f} GB/s") +@pytest.mark.parametrize("m", [1, 8, 32, 128, 1024, 4096]) +@pytest.mark.parametrize("k, n", [(4096, 4096)]) +@pytest.mark.parametrize("dt", [torch.float16, torch.float32]) +@pytest.mark.parametrize("runs", [1000]) +def test_woqgemm_s8_joint_matrix_vs_sycl_tla_perf(m, k, n, dt, runs, record_property): + if not torch.xpu.is_available(): + pytest.skip("No XPU Device") + if not hasattr(ark, "woqgemm_s8_joint_matrix") or not hasattr(ark, "woqgemm_s8_sycl_tla"): + pytest.skip("Both woqgemm_s8 backends are not available in this build") + + torch.manual_seed(0) + A = torch.rand(m, k, dtype=dt, device="xpu") - 0.5 + B = torch.randint(-128, 127, (n, k), dtype=torch.int8, device="xpu") + scaleB = torch.rand(n, 1, dtype=dt, device="xpu") / 100 + bias = torch.rand(1, n, dtype=dt, device="xpu") + 2 + + out_joint = ark.woqgemm_s8_joint_matrix(A, B, scaleB, bias) + out_tla = ark.woqgemm_s8_sycl_tla(A, B, scaleB, bias) + torch.xpu.synchronize() + + diff = (out_joint - out_tla).abs() + print(f"\n Max Diff joint_matrix vs sycl_tla: {diff.max().item():.5f}, Mean Diff: {diff.mean().item():.5f}") + assert torch.allclose(out_tla, out_joint, atol=1, rtol=0.1) + + def bench(op): + warmup = 200 + for _ in range(warmup): + _ = op(A, B, scaleB, bias) + torch.xpu.synchronize() + + st = time.perf_counter() + for _ in range(runs): + _ = op(A, B, scaleB, bias) + torch.xpu.synchronize() + return (time.perf_counter() - st) / runs + + joint_dur = bench(ark.woqgemm_s8_joint_matrix) + tla_dur = bench(ark.woqgemm_s8_sycl_tla) + + ops = m * n * k * 2 + joint_tflops = ops / joint_dur / 1e12 + tla_tflops = ops / tla_dur / 1e12 + speedup = joint_dur / tla_dur + + print(f" [joint_matrix] : {joint_dur * 1000:8.4f} ms {joint_tflops:7.3f} TFLOPS") + print(f" [sycl_tla] : {tla_dur * 1000:8.4f} ms {tla_tflops:7.3f} TFLOPS speedup={speedup:5.2f}x") + + record_property("joint_matrix_time_ms", round(joint_dur * 1000, 4)) + record_property("sycl_tla_time_ms", round(tla_dur * 1000, 4)) + record_property("speedup", round(speedup, 4)) + + @pytest.mark.parametrize("m", [4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) @pytest.mark.parametrize("dt", [torch.float16, torch.float32]) @@ -201,7 +253,7 @@ def test_cpu(m, k, n, dt, batch_size, runs, record_property): @pytest.mark.parametrize("m", [1, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("batch_size", [1]) @pytest.mark.parametrize("runs", [1]) def test_xpu_sycl_tla(m, k, n, dt, batch_size, runs, record_property): @@ -231,7 +283,7 @@ def test_xpu_sycl_tla(m, k, n, dt, batch_size, runs, record_property): @pytest.mark.parametrize("m", [1, 8, 16, 32, 64, 128, 256, 1024, 2048, 4096]) @pytest.mark.parametrize("k, n", [(4096, 4096)]) -@pytest.mark.parametrize("dt", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dt", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("batch_size", [1]) @pytest.mark.parametrize("runs", [1]) def test_xpu_sycl_tla_no_bias(m, k, n, dt, batch_size, runs, record_property): From f18a732063e25ab496ff674afe1ee6bb014f8049 Mon Sep 17 00:00:00 2001 From: Zhenzhong1 Date: Thu, 16 Jul 2026 16:56:55 +0800 Subject: [PATCH 8/8] perf ipmprove Signed-off-by: Zhenzhong1 --- .../wrapper/include/sycl_tla_s8_gemm.hpp | 51 ++++++++++++++++--- 1 file changed, 44 insertions(+), 7 deletions(-) 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 4cfcb175f..ab909b305 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 @@ -72,19 +72,39 @@ void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma, int k_tile_count = ceil_div(shape<1>(A), get<2>(wg_tile)); clear(tCrC); + auto prefetch_a = make_block_2d_prefetch(copy_a); + auto prefetch_b = make_block_2d_prefetch(copy_b); - for (int k_tile = 0; k_tile < k_tile_count; ++k_tile) { + auto thr_prefetch_A = prefetch_a.get_slice(local_id); + auto thr_prefetch_B = prefetch_b.get_slice(local_id); + + auto pAgA = thr_prefetch_A.partition_S(gA); + auto pBgB = thr_prefetch_B.partition_S(gB); + + constexpr int prefetch_dist = 3; + int k_tile_prefetch = 0; + + CUTE_UNROLL + for (; k_tile_prefetch < prefetch_dist; ++k_tile_prefetch) { + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + } + + for (int k_tile = 0; k_tile < k_tile_count; ++k_tile, ++k_tile_prefetch) { barrier_arrive(barrier_scope); copy(copy_a, tAgA(_, _, _, k_tile), tArA); copy(copy_b, tBgB(_, _, _, k_tile), tBrB); + prefetch(prefetch_a, pAgA(_, _, _, k_tile_prefetch)); + prefetch(prefetch_b, pBgB(_, _, _, k_tile_prefetch)); + reorder(tArA, tCrA); reorder(tBrB, tCrB); gemm(mma, tCrA, tCrB, tCrC); barrier_wait(barrier_scope); - } + } CUTE_UNROLL for (int i = 0; i < size(tCrC); ++i) { @@ -107,6 +127,7 @@ void igemm_device_impl(ATensor const& A, BTensor const& B, TiledMMA const& mma, } } + template void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, const int8_t* a, const int8_t* b, ElementOut* c, float* accum, const ElementOut* scale_a, const ElementOut* scale_b, @@ -136,12 +157,28 @@ void launch_igemm_tile(sycl::queue* q, int m, int n, int gemm_k, int lda, int ld } template -void launch_igemm(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, const int8_t* a, const int8_t* b, - ElementOut* c, float* accum, const ElementOut* scale_a, const ElementOut* scale_b, +void launch_igemm(sycl::queue* q, int m, int n, int gemm_k, int lda, int ldb, + const int8_t* a, const int8_t* b, ElementOut* c, float* accum, + const ElementOut* scale_a, const ElementOut* scale_b, const ElementOut* bias, int block_idx, int scale_b_stride) { - using SGLayout = Layout, Stride<_4, _1, _0>>; - launch_igemm_tile( - q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); + using SmallTileSG = Layout, Stride<_0, _1, _0>>; + using SmallMidTileSG = Layout, Stride<_4, _1, _0>>; + using MediumTileSG = Layout, Stride<_4, _1, _0>>; + using LargeTileSG = Layout, Stride<_4, _1, _0>>; + + if (m < 16) { + launch_igemm_tile( + q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); + } else if (m < 128) { + launch_igemm_tile( + q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); + } else if (m <= 1024) { + launch_igemm_tile( + q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); + } else { + launch_igemm_tile( + q, m, n, gemm_k, lda, ldb, a, b, c, accum, scale_a, scale_b, bias, block_idx, scale_b_stride); + } } template