diff --git a/build2cmake/src/templates/xpu/torch-extension.cmake b/build2cmake/src/templates/xpu/torch-extension.cmake
index a359f8c7..d2a95e47 100644
--- a/build2cmake/src/templates/xpu/torch-extension.cmake
+++ b/build2cmake/src/templates/xpu/torch-extension.cmake
@@ -1,6 +1,7 @@
 define_gpu_extension_target(
   {{ ops_name }}
   DESTINATION {{ ops_name }}
+  INCLUDE_DIRECTORIES "${ONEDNN_XPU_INCLUDE_DIR}"
   LANGUAGE ${GPU_LANG}
   SOURCES ${SRC}
   COMPILE_FLAGS ${sycl_flags}
diff --git a/lib/torch-extension/default.nix b/lib/torch-extension/default.nix
index 5e35282c..41e90860 100644
--- a/lib/torch-extension/default.nix
+++ b/lib/torch-extension/default.nix
@@ -52,6 +52,7 @@ let
   );
 
   oneapi-torch-dev = xpuPackages.oneapi-torch-dev.override { inherit stdenv; };
+  onednn-xpu = xpuPackages.onednn-xpu.override { inherit stdenv oneapi-torch-dev; };
 
   # On Darwin, we need the host's xcrun for `xcrun metal` to compile Metal shaders.
   # t's not supported by the nixpkgs shim.
@@ -132,7 +133,7 @@ stdenv.mkDerivation (prevAttrs: {
   ++ lib.optionals rocmSupport (with rocmPackages; [ hipsparselt ])
   ++ lib.optionals xpuSupport ([
     oneapi-torch-dev
-    xpuPackages.onednn-xpu
+    onednn-xpu
   ])
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
     apple-sdk_15
@@ -172,6 +173,9 @@ stdenv.mkDerivation (prevAttrs: {
     (lib.cmakeFeature "CMAKE_HIP_COMPILER_ROCM_ROOT" "${clr}")
     (lib.cmakeFeature "HIP_ROOT_DIR" "${clr}")
   ]
+  ++ lib.optionals xpuSupport [
+    (lib.cmakeFeature "ONEDNN_XPU_INCLUDE_DIR" "${onednn-xpu}/include")
+  ]
   ++ lib.optionals stdenv.hostPlatform.isDarwin [
     # Use host compiler for Metal. Not included in the redistributable SDK.
     (lib.cmakeFeature "METAL_COMPILER" "${xcrunHost}/bin/xcrunHost")