Decoupling oneDNN

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

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")
 
@@ -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(

auto_round_extension/ark/auto_round_kernel/ark.cpp

@@ -21,6 +21,7 @@ typedef uintptr_t torch_ptr;
 #if ARK_XPU
 #include <sycl/sycl.hpp>
 #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,27 +33,59 @@ 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
 }
 
+#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
+  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,
@@ -677,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
 }

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<int*>(
-      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<int*>(
-      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;

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