Fix xetla import for wheel

Author: ZzEeKkAa

benchmarks/triton_kernels_benchmark/flash_attention_fwd_benchmark.py

@@ -4,7 +4,7 @@
 import triton.language as tl
 
 import triton_kernels_benchmark as benchmark_suit
-import xetla_kernel
+import triton_kernels_benchmark.xetla_kernel
 
 if benchmark_suit.USE_IPEX_OPTION:
     import intel_extension_for_pytorch  # type: ignore # noqa: F401
@@ -262,7 +262,7 @@ def benchmark(Z, H, N_CTX, D_HEAD, CAUSAL, provider):
 
     elif provider == 'xetla':
         module_name = f'flash_attn_causal_{CAUSAL}'.lower()
-        func = getattr(xetla_kernel, module_name)
+        func = getattr(triton_kernels_benchmark.xetla_kernel, module_name)
         out = torch.empty_like(q, device='xpu', dtype=dtype)
         size_score = Z * H * N_CTX * N_CTX
         size_attn_mask = Z * N_CTX * N_CTX

benchmarks/triton_kernels_benchmark/fused_softmax.py

@@ -13,7 +13,7 @@
 from triton.runtime import driver
 
 import triton_kernels_benchmark as benchmark_suit
-import xetla_kernel
+import triton_kernels_benchmark.xetla_kernel
 
 
 @torch.jit.script
@@ -140,7 +140,7 @@ def benchmark(M, N, provider):
 
     elif provider == "xetla":
         name = f"softmax_shape_{M}_{N}"
-        func = getattr(xetla_kernel, name)
+        func = getattr(triton_kernels_benchmark.xetla_kernel, name)
         out = torch.empty_like(x, device="xpu")
         xetla_fn = lambda: func(x, out, 0)
         torch_fn = lambda: torch.softmax(x, axis=-1)

benchmarks/triton_kernels_benchmark/gemm_benchmark.py

@@ -15,7 +15,7 @@
 import triton_kernels_benchmark as benchmark_suit
 from triton_kernels_benchmark.benchmark_testing import do_bench_elapsed_time, BENCHMARKING_METHOD
 
-import xetla_kernel
+import triton_kernels_benchmark.xetla_kernel
 
 if benchmark_suit.USE_IPEX_OPTION:
     import intel_extension_for_pytorch  # type: ignore # noqa: F401
@@ -313,7 +313,7 @@ def benchmark(B, M, N, K, provider):
         # better performance.
         if (B, M, N, K) == (1, 3072, 4096, 3072):
             name = 'gemm_streamk_shape_3072_4096_3072'
-        func = getattr(xetla_kernel, name)
+        func = getattr(triton_kernels_benchmark.xetla_kernel, name)
         xetla_fn = lambda: func(a, b, c, acc, cnt)
         torch_fn = lambda: torch.matmul(a, b).to(torch.float32)
 

benchmarks/triton_kernels_benchmark/gemm_splitk_benchmark.py

@@ -3,7 +3,7 @@
 import triton.language as tl
 
 import triton_kernels_benchmark as benchmark_suit
-import xetla_kernel
+import triton_kernels_benchmark.xetla_kernel
 
 if benchmark_suit.USE_IPEX_OPTION:
     import intel_extension_for_pytorch  # type: ignore # noqa: F401
@@ -165,7 +165,7 @@ def benchmark(M, N, K, provider):
         cnt = torch.empty((M, N), device='xpu', dtype=torch.int32)
 
         name = f'gemm_splitk_shape_{M}_{K}_{N}'
-        func = getattr(xetla_kernel, name)
+        func = getattr(triton_kernels_benchmark.xetla_kernel, name)
         xetla_fn = lambda: func(a, b, c, acc, cnt)
         torch_fn = lambda: torch.matmul(a, b).to(torch.float32)
 

benchmarks/triton_kernels_benchmark/gemm_streamk_benchmark.py

@@ -10,7 +10,7 @@
 import triton.language as tl
 
 import triton_kernels_benchmark as benchmark_suit
-import xetla_kernel
+import triton_kernels_benchmark.xetla_kernel
 
 if benchmark_suit.USE_IPEX_OPTION:
     import intel_extension_for_pytorch  # type: ignore # noqa: F401
@@ -288,7 +288,7 @@ def benchmark(M, N, K, provider):
         cnt = torch.empty((M, N), device='xpu', dtype=torch.int32)
 
         name = f'gemm_streamk_shape_{M}_{K}_{N}'
-        func = getattr(xetla_kernel, name)
+        func = getattr(triton_kernels_benchmark.xetla_kernel, name)
         xetla_fn = lambda: func(a, b, c, acc, cnt)
         torch_fn = lambda: torch.matmul(a, b).to(torch.float32)