diff --git a/omni/omni_xpu_kernel/README.md b/omni/omni_xpu_kernel/README.md index c2b5ad78..a4f7d6ff 100644 --- a/omni/omni_xpu_kernel/README.md +++ b/omni/omni_xpu_kernel/README.md @@ -134,7 +134,7 @@ output = rotary.rotary_emb(x, cos_cache, sin_cache, seq_len, heads) - Intel oneAPI DPC++/C++ Compiler (icpx) - PyTorch >= 2.0 with XPU support - Intel GPU: Arc B-series (BMG), Data Center GPU Max (PVC), or compatible -- oneDNN (for INT4/FP8 GEMM; auto-detected from oneAPI) +- oneDNN (for INT4/FP8 GEMM; auto-detected from oneAPI when enabled) ## Installation @@ -151,6 +151,30 @@ OMNI_XPU_DEVICE=pvc pip install -e . --no-build-isolation # Data Center GPU Ma On Windows, see [WHL_BUILD_INSTALL.md](WHL_BUILD_INSTALL.md). +### Windows status + +Windows wheel build and runtime validation currently use this known-good stack: + +| Component | Version | +| --- | --- | +| Python | 3.12 | +| PyTorch | `2.10.0+xpu` | +| torchvision | `0.25.0+xpu` | +| torchaudio | `2.10.0+xpu` | +| oneAPI compiler | `2025.3` | +| Intel XPU runtime | `2025.3.x` | +| MSVC toolchain | VS 2022 Build Tools | + +The critical Windows requirement is **compiler/runtime alignment**. Building with a newer oneAPI compiler and loading into an older PyTorch XPU runtime can fail at import time or surface misleading runtime errors. + +Current Windows guidance: + +- set `OMNI_XPU_ONEAPI_VERSION=2025.3` +- use `OMNI_XPU_DEVICE=bmg` unless you are targeting another validated device +- keep `OMNI_XPU_ENABLE_ONEDNN=0` unless you are explicitly validating oneDNN-backed kernels on Windows + +The package now also configures Windows DLL search paths for the active Python runtime, `torch\lib`, and the oneAPI compiler locale directory. See [WHL_BUILD_INSTALL.md](WHL_BUILD_INSTALL.md) for the full build, install, and troubleshooting guide. + ## Debug Logging Controlled by `OMNI_XPU_DEBUG` environment variable. **Disabled by default.** @@ -211,8 +235,8 @@ To switch config at compile time: `-DSDP_CONFIG_PVC` ### Build System The package builds multiple extension modules: -- `_C.so` — Main extension (norm, gguf, svdq, rotary, sdp loader, fp8) -- `lgrf_sdp.so` — SDP ESIMD sidecar (AOT, doubleGRF) +- `_C.so` / `_C.pyd` — Main extension (norm, gguf, svdq, rotary, sdp loader, fp8) +- `lgrf_sdp.so` / `lgrf_sdp*.pyd` — SDP ESIMD sidecar (AOT, doubleGRF) ## License diff --git a/omni/omni_xpu_kernel/WHL_BUILD_INSTALL.md b/omni/omni_xpu_kernel/WHL_BUILD_INSTALL.md index af21b331..b0cb9a62 100644 --- a/omni/omni_xpu_kernel/WHL_BUILD_INSTALL.md +++ b/omni/omni_xpu_kernel/WHL_BUILD_INSTALL.md @@ -1,123 +1,232 @@ -# omni_xpu_kernel (llm-scaler) Build and Install Notes - -This document summarizes how to build the updated omni_xpu_kernel from llm-scaler on Windows using a conda environment, then install it into the ComfyUI embedded Python environment. - -## 1. Prerequisites -- Intel oneAPI Base Toolkit installed -- VS2022 C++ components installed -- conda environment: omni_env -- ComfyUI embedded Python located at: %EMBED_PYTHON_DIR% -- Variables (adjust as needed): - - %WORKSPACE%: workspace root - - %EMBED_PYTHON_DIR%: embedded Python directory (example: %WORKSPACE%\omni\comfyui_windows_setup\python_embeded) - - %LLM_SCALER_DIR%: llm-scaler repo directory (example: %WORKSPACE%\llm-scaler) - - %OUTPUT_DIR%: build output directory (example: %WORKSPACE%\llm_scaler_dist) - -### Variable setup example (cmd) -You can put these variables at the top of your command prompt or script for consistent configuration: - -``` -set "WORKSPACE=C:\workspace" -set "LLM_SCALER_DIR=%WORKSPACE%\llm-scaler" -set "EMBED_PYTHON_DIR=%WORKSPACE%\omni\comfyui_windows_setup\python_embeded" -set "OUTPUT_DIR=%WORKSPACE%\llm_scaler_dist" -``` - -## 2. Build (in conda environment) -Build script location: -- [scripts/build_llm_scaler_conda.cmd](scripts/build_llm_scaler_conda.cmd) - -Script contents (full example, with variables): - -``` -@echo off -set "VS2022INSTALLDIR=C:\Program Files\Microsoft Visual Studio\18\Community" -call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" -call C:\ProgramData\miniforge3\Scripts\activate.bat omni_env -set "WORKSPACE=C:\workspace" -set "LLM_SCALER_DIR=%WORKSPACE%\llm-scaler" -set "EMBED_PYTHON_DIR=%WORKSPACE%\omni\comfyui_windows_setup\python_embeded" -set "OUTPUT_DIR=%WORKSPACE%\llm_scaler_dist" -set "PATH=%CONDA_PREFIX%\Library\bin;%CONDA_PREFIX%\Lib\site-packages\torch\lib;%PATH%" -cd /d %LLM_SCALER_DIR%\omni\omni_xpu_kernel -python -m pip wheel . -w %OUTPUT_DIR% --no-build-isolation --no-deps > %WORKSPACE%\build_log.txt 2>&1 -``` - -How to run the .cmd file (common options): -- Double-click it in File Explorer. -- From Command Prompt: - - `cmd /c scripts\build_llm_scaler_conda.cmd` -- From PowerShell: - - `cmd /c "scripts\build_llm_scaler_conda.cmd"` - -Key points: -- Initialize oneAPI: setvars.bat -- Activate conda: omni_env -- Set torch DLL search path -- Produce wheel to %OUTPUT_DIR% - -Output: -- %OUTPUT_DIR%\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl - -## 3. Install the newly built wheel (embedded Python) -Use --no-deps to avoid pulling dependencies again: - - - Wheel path: - - %OUTPUT_DIR%\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl - - Install command: - - pip install --force-reinstall --no-deps - -Command example (with variables): - -``` -"%EMBED_PYTHON_DIR%\python.exe" -m pip install --force-reinstall --no-deps "%OUTPUT_DIR%\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl" -``` - -## 4. Verify (optional) -- Import omni_xpu_kernel -- Check __version__ - -Command example (with variables): - -``` -"%EMBED_PYTHON_DIR%\python.exe" -c "import omni_xpu_kernel as ok, importlib.metadata as im; print('omni_xpu_kernel:', ok); print('version:', im.version('omni-xpu-kernel'))" -``` - ---- - -To update or rebuild, repeat steps 2 → 3. - ---- - -## FAQ - -### 1) XPU symbol link errors (c10::xpu::XPUStream) -Add missing link libraries in llm-scaler code: - - - Source file: - - [setup.py](setup.py) - - - Add on Windows: - - torch_xpu.lib - - c10_xpu.lib - - - Add on Linux: - - -ltorch_xpu - - -lc10_xpu - -### 2) Restore embedded PyTorch to XPU build -Reinstall the XPU build of PyTorch in embedded Python (to avoid fallback to standard builds): - - - Reference script: - - [setup_portable_env.bat](../../../omni/comfyui_windows_setup/setup_portable_env.bat) - - Key versions: - - torch==2.9.0+xpu - - torchvision==0.24.0+xpu - - torchaudio==2.9.0+xpu - - --index-url https://download.pytorch.org/whl/xpu - -Command example (with variables): - -``` -"%EMBED_PYTHON_DIR%\python.exe" -m pip install torch==2.9.0+xpu torchvision==0.24.0+xpu torchaudio==2.9.0+xpu --index-url https://download.pytorch.org/whl/xpu -``` +# omni_xpu_kernel Windows wheel build guide + +This document describes the Windows workflow that is currently known to build and run `omni_xpu_kernel` successfully. + +It is based on the configuration validated during local bring-up of: + +- wheel build on Windows +- import of `omni_xpu_kernel._C` +- `norm.rms_norm(...)` +- `svdq.unpack_int4(...)` +- `sdp.sdp(...)` + +## Validated version matrix + +Use this combination unless you are intentionally re-validating a different stack: + +| Component | Version | +| --- | --- | +| OS | Windows x64 | +| Python | 3.12 | +| PyTorch | `2.10.0+xpu` | +| torchvision | `0.25.0+xpu` | +| torchaudio | `2.10.0+xpu` | +| oneAPI compiler | `2025.3` | +| Intel XPU runtime packages | `2025.3.x` | +| MSVC toolchain | Visual Studio 2022 Build Tools | + +The most important rule on Windows is: **the compiler/runtime used to build the wheel must match the SYCL runtime shipped with the active PyTorch XPU environment**. + +The working combination above uses `sycl8.dll` from the `2025.3` stack. Building with a newer oneAPI toolchain and loading into an older PyTorch XPU environment can produce import failures or misleading runtime errors. + +## Prerequisites + +Install the following before building: + +1. Intel oneAPI C++ Compiler `2025.3` +2. Visual Studio 2022 Build Tools with MSVC C++ components +3. Miniconda or Miniforge +4. Intel GPU driver capable of running PyTorch XPU + +Recommended oneAPI component on Windows: + +- Intel oneAPI C++ Essentials / DPC++ Compiler `2025.3` + +## Create the build environment + +The examples below use a conda environment named `omni_env`. + +```cmd +conda create -n omni_env python=3.12 -y +conda activate omni_env +python -m pip install --upgrade pip setuptools wheel +python -m pip install --index-url https://download.pytorch.org/whl/xpu ^ + torch==2.10.0+xpu torchvision==0.25.0+xpu torchaudio==2.10.0+xpu +python -m pip install ^ + dpcpp-cpp-rt==2025.3.1 ^ + intel-cmplr-lib-rt==2025.3.1 ^ + intel-cmplr-lib-ur==2025.3.1 ^ + intel-cmplr-lic-rt==2025.3.1 ^ + intel-sycl-rt==2025.3.1 ^ + intel-opencl-rt==2025.3.1 ^ + intel-openmp==2025.3.1 ^ + intel-pti==0.15.0 ^ + onemkl-sycl-blas==2025.3.0 ^ + onemkl-sycl-dft==2025.3.0 ^ + onemkl-sycl-lapack==2025.3.0 ^ + onemkl-sycl-rng==2025.3.0 ^ + onemkl-sycl-sparse==2025.3.0 ^ + mkl==2025.3.0 ^ + tbb==2022.3.0 ^ + tcmlib==1.4.1 ^ + umf==1.0.2 ^ + pytorch-triton-xpu==3.5.0 +``` + +Quick sanity check: + +```cmd +python -c "import torch; print(torch.__version__); print(torch.xpu.is_available())" +``` + +## Build the wheel + +The commands below assume: + +- repo root: `C:\workspace\llm-scaler` +- package dir: `C:\workspace\llm-scaler\omni\omni_xpu_kernel` +- wheel output dir: `C:\workspace\llm_scaler_dist` + +Open a Developer Command Prompt for VS 2022, or otherwise ensure the MSVC build tools are on `PATH`, then run: + +```cmd +call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" +call C:\ProgramData\miniforge3\Scripts\activate.bat omni_env + +set "WORKSPACE=C:\workspace" +set "LLM_SCALER_DIR=%WORKSPACE%\llm-scaler" +set "OUTPUT_DIR=%WORKSPACE%\llm_scaler_dist" + +set "OMNI_XPU_ONEAPI_VERSION=2025.3" +set "OMNI_XPU_DEVICE=bmg" +set "OMNI_XPU_ENABLE_ONEDNN=0" +set "PATH=%CONDA_PREFIX%\Library\bin;%CONDA_PREFIX%\DLLs;%CONDA_PREFIX%\Lib\site-packages\torch\lib;%PATH%" + +cd /d "%LLM_SCALER_DIR%\omni\omni_xpu_kernel" +python -m pip wheel . -w "%OUTPUT_DIR%" --no-build-isolation --no-deps +``` + +Expected wheel: + +```text +C:\workspace\llm_scaler_dist\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl +``` + +### Important build variables + +| Variable | Meaning | +| --- | --- | +| `OMNI_XPU_ONEAPI_VERSION` | Prefer a specific installed oneAPI compiler version on Windows | +| `OMNI_XPU_DEVICE` | AOT target for the ESIMD SDP sidecar, default `bmg` | +| `OMNI_XPU_ENABLE_ONEDNN` | Enable oneDNN-backed kernels explicitly; current Windows guidance is `0` | + +Notes: + +- `OMNI_XPU_ENABLE_ONEDNN=0` is the recommended Windows default today. +- `OMNI_XPU_DEVICE=bmg` is the validated target for Arc B-series testing. +- If you want to validate oneDNN-backed kernels on Windows, treat that as a separate compatibility pass. + +## Install the wheel into the build environment + +```cmd +conda activate omni_env +python -m pip install --force-reinstall --no-deps ^ + C:\workspace\llm_scaler_dist\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl +``` + +## Verify in the build environment + +Use the same environment that built the wheel: + +```cmd +python -c "import omni_xpu_kernel as ok; print(ok.__version__); print(ok.is_available())" +python -c "import torch, omni_xpu_kernel as ok; x=torch.randn(2,2560,device='xpu',dtype=torch.float16); w=torch.randn(2560,device='xpu',dtype=torch.float16); y=ok.norm.rms_norm(w,x,1e-6); torch.xpu.synchronize(); print(y.shape, y.device, y.dtype)" +``` + +Optional SDP smoke test: + +```cmd +python -c "import torch; from omni_xpu_kernel import sdp; q=torch.randn(1,64,30,128,device='xpu',dtype=torch.bfloat16); k=torch.randn(1,64,30,128,device='xpu',dtype=torch.bfloat16); v=torch.randn(1,64,30,128,device='xpu',dtype=torch.bfloat16); y=sdp.sdp(q,k,v); torch.xpu.synchronize(); print(y.shape, y.device, y.dtype)" +``` + +## Install into another Python environment + +For example, to install into an embedded ComfyUI Python: + +```cmd +"%EMBED_PYTHON_DIR%\python.exe" -m pip install --force-reinstall --no-deps ^ + C:\workspace\llm_scaler_dist\omni_xpu_kernel-0.1.0-cp312-cp312-win_amd64.whl +``` + +The target environment should also use the same validated XPU stack: + +```cmd +"%EMBED_PYTHON_DIR%\python.exe" -m pip install --force-reinstall --index-url https://download.pytorch.org/whl/xpu ^ + torch==2.10.0+xpu torchvision==0.25.0+xpu torchaudio==2.10.0+xpu +``` + +## Windows runtime notes + +The wheel now relies on Windows DLL search setup performed in `omni_xpu_kernel.__init__`. + +At import time it adds these locations when present: + +1. active Python `Library\bin` +2. active Python `DLLs` +3. `torch\lib` +4. oneAPI compiler locale directory such as `compiler\2025.3\bin\1033` +5. any extra directories from `OMNI_XPU_DLL_DIRS` + +This is important on Windows because the extension can depend on Intel compiler runtime resources such as `irc_msg.dll`. + +## Troubleshooting + +### `_C` import fails with a missing DLL or "找不到指定的模块" + +Usually this means one of: + +- oneAPI compiler/runtime version does not match the PyTorch XPU runtime +- the target environment is missing required Intel runtime DLLs +- the embedded environment is not using the XPU PyTorch build + +Check: + +```cmd +python -c "import torch; print(torch.__version__); print(torch.__file__); print(torch.xpu.is_available())" +python -c "import omni_xpu_kernel as ok; print(ok.is_available())" +``` + +### Wheel built, but kernels fail at runtime with "device or resource busy" + +Do not assume this is only a device contention issue. On Windows this can also be a symptom of a mismatched SYCL runtime. + +Re-check: + +- oneAPI compiler version +- PyTorch XPU version +- XPU runtime package versions + +The validated combination in this document is the known-good baseline. + +### SDP fails to load the sidecar + +The Windows loader now scans `lgrf_uni\` for `lgrf_sdp*.pyd`. + +If this fails, confirm the wheel contains a packaged sidecar under: + +```text +omni_xpu_kernel\lgrf_uni\ +``` + +### oneDNN kernels on Windows + +Windows builds currently default to `OMNI_XPU_ENABLE_ONEDNN=0`. + +That keeps the validated Windows path focused on: + +- `_C` import +- normalization kernels +- SVDQ ESIMD helpers +- SDP sidecar loading + +If you enable oneDNN-backed kernels on Windows, re-validate that configuration explicitly. diff --git a/omni/omni_xpu_kernel/omni_xpu_kernel/__init__.py b/omni/omni_xpu_kernel/omni_xpu_kernel/__init__.py index 0620a5b8..53205787 100644 --- a/omni/omni_xpu_kernel/omni_xpu_kernel/__init__.py +++ b/omni/omni_xpu_kernel/omni_xpu_kernel/__init__.py @@ -17,12 +17,63 @@ import os import sys +from pathlib import Path __version__ = "0.1.0" __author__ = "Intel" # Lazy loading of native extension _native_module = None +_dll_dir_handles = [] +_dll_dir_paths = set() + + +def _add_windows_dll_directory(path: Path) -> None: + """Register a DLL search path and keep the handle alive for process lifetime.""" + if not path.is_dir(): + return + + resolved = path.resolve() + if resolved in _dll_dir_paths: + return + + handle = os.add_dll_directory(str(resolved)) + _dll_dir_handles.append(handle) + _dll_dir_paths.add(resolved) + + +def _configure_windows_dll_search_paths() -> None: + """Ensure the active Python environment's runtime DLLs are discoverable.""" + if sys.platform != "win32": + return + + python_root = Path(sys.executable).resolve().parent + _add_windows_dll_directory(python_root / "Library" / "bin") + _add_windows_dll_directory(python_root / "DLLs") + + program_files_x86 = Path(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)")) + oneapi_root = program_files_x86 / "Intel" / "oneAPI" / "compiler" + preferred_version = os.environ.get("OMNI_XPU_ONEAPI_VERSION") + locale_dirs = [] + if preferred_version: + locale_dirs.append(oneapi_root / preferred_version / "bin" / "1033") + if oneapi_root.is_dir(): + locale_dirs.extend(sorted(oneapi_root.glob("*\\bin\\1033"), reverse=True)) + for locale_dir in locale_dirs: + _add_windows_dll_directory(locale_dir) + + try: + import torch + + torch_root = Path(torch.__file__).resolve().parent + _add_windows_dll_directory(torch_root / "lib") + except Exception: + pass + + extra_dirs = os.environ.get("OMNI_XPU_DLL_DIRS", "") + for raw_path in extra_dirs.split(os.pathsep): + if raw_path: + _add_windows_dll_directory(Path(raw_path)) def _load_extension(): """Load the native C++ extension module.""" @@ -31,6 +82,7 @@ def _load_extension(): return _native_module try: + _configure_windows_dll_search_paths() from omni_xpu_kernel import _C _native_module = _C return _native_module diff --git a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bindings.cpp b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bindings.cpp index 979ed2a7..0587a570 100644 --- a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bindings.cpp +++ b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/bindings.cpp @@ -33,9 +33,11 @@ namespace svdq { torch::Tensor dequantize_svdq_w4(const torch::Tensor& packed, const torch::Tensor& scales, torch::ScalarType out_dtype); torch::Tensor unpack_svdq_int4(const torch::Tensor& packed, bool is_signed); std::tuple quantize_svdq_act_int4(const torch::Tensor& input, int64_t group_size); +#if OMNI_XPU_HAS_ONEDNN torch::Tensor onednn_int4_gemm(const torch::Tensor& act, const torch::Tensor& packed, const torch::Tensor& wscales); torch::Tensor onednn_int4_gemm_preconverted(const torch::Tensor& act, const torch::Tensor& packed_u4, const torch::Tensor& scales_f16); void onednn_int4_gemm_add_to_output(const torch::Tensor& act, const torch::Tensor& packed_u4, const torch::Tensor& scales_f16, torch::Tensor& dst); +#endif void fused_convert_add(torch::Tensor& out, const torch::Tensor& result, const torch::Tensor& residual); torch::Tensor fused_smooth_convert(const torch::Tensor& x, const torch::Tensor& smooth_factor); torch::Tensor fused_smooth_mul_convert(const torch::Tensor& x, const torch::Tensor& rcp_smooth); @@ -46,11 +48,13 @@ namespace rotary { namespace sdp { torch::Tensor sdp(torch::Tensor q, torch::Tensor k, torch::Tensor v); } +#if OMNI_XPU_HAS_ONEDNN namespace linear { torch::Tensor onednn_w8a16_fp8(torch::Tensor input, torch::Tensor weight, torch::Tensor scale_w, std::optional bias); void fp8_cache_clear(); std::tuple fp8_cache_stats(); } +#endif } PYBIND11_MODULE(_C, m) { @@ -127,6 +131,7 @@ PYBIND11_MODULE(_C, m) { "Output: (packed [M, K/2] uint8, scales [num_groups, M])", py::arg("input"), py::arg("group_size") = 64); +#if OMNI_XPU_HAS_ONEDNN svdq.def("onednn_int4_gemm", &omni_xpu::svdq::onednn_int4_gemm, "Fused INT4 dequant + GEMM using oneDNN u4 matmul primitive\n" "Converts signed INT4 to u4 and bf16 scales to f16 per call\n" @@ -147,6 +152,7 @@ PYBIND11_MODULE(_C, m) { "Input: act [M, K] f16, packed_u4 [N, K/2] uint8, scales_f16 [G, N] f16, dst [M, N] bf16\n" "Output: dst modified in-place (dst += GEMM result)", py::arg("act"), py::arg("packed_u4"), py::arg("scales_f16"), py::arg("dst")); +#endif svdq.def("fused_convert_add", &omni_xpu::svdq::fused_convert_add, "Fused f16->bf16 conversion + bf16 addition in single ESIMD kernel\n" @@ -181,6 +187,7 @@ PYBIND11_MODULE(_C, m) { // FP8 Linear (oneDNN W8A16) auto linear = m.def_submodule("linear", "FP8 linear kernels"); +#if OMNI_XPU_HAS_ONEDNN linear.def("onednn_w8a16_fp8", &omni_xpu::linear::onednn_w8a16_fp8, "FP8 GEMM: W8A16 matmul with E4M3/E5M2 weights via oneDNN.\n" "Input: x [M, K] fp16/bf16, weight [N, K] float8, scales [N] f32\n" @@ -190,6 +197,7 @@ PYBIND11_MODULE(_C, m) { "Clear FP8 primitive cache"); linear.def("fp8_cache_stats", &omni_xpu::linear::fp8_cache_stats, "Return FP8 cache stats as (hits, misses, size)"); +#endif // Scaled Dot-Product Attention (ESIMD Flash Attention) auto sdp = m.def_submodule("sdp", "Scaled dot-product attention kernels"); diff --git a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/norm.cpp b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/norm.cpp index 4b2b8769..2f8408bb 100644 --- a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/norm.cpp +++ b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/norm.cpp @@ -314,12 +314,17 @@ using rms_fn_t = void(*)(const void*, const void*, void*, float, const int, cons template rms_fn_t select_rms_kernel(int nb) { +#ifdef _WIN32 + (void)nb; + return rms_norm_kernel; +#else if (nb <= 1) return rms_norm_kernel; if (nb <= 2) return rms_norm_kernel; if (nb <= 4) return rms_norm_kernel; if (nb <= 8) return rms_norm_kernel; if (nb <= 16) return rms_norm_kernel; return rms_norm_kernel; +#endif } // LayerNorm dispatch @@ -328,12 +333,17 @@ using ln_fn_t = void(*)(const void*, const uint64_t, const uint64_t, void*, floa template ln_fn_t select_ln_kernel(int nb) { +#ifdef _WIN32 + (void)nb; + return layer_norm_kernel; +#else if (nb <= 1) return layer_norm_kernel; if (nb <= 2) return layer_norm_kernel; if (nb <= 4) return layer_norm_kernel; if (nb <= 8) return layer_norm_kernel; if (nb <= 16) return layer_norm_kernel; return layer_norm_kernel; +#endif } // Fused add rms norm dispatch @@ -342,12 +352,17 @@ using fused_fn_t = void(*)(const void*, void*, void*, float, const int, const in template fused_fn_t select_fused_kernel(int nb) { +#ifdef _WIN32 + (void)nb; + return fused_add_rms_norm_kernel; +#else if (nb <= 1) return fused_add_rms_norm_kernel; if (nb <= 2) return fused_add_rms_norm_kernel; if (nb <= 4) return fused_add_rms_norm_kernel; if (nb <= 8) return fused_add_rms_norm_kernel; if (nb <= 16) return fused_add_rms_norm_kernel; return fused_add_rms_norm_kernel; +#endif } // ============================================================================ diff --git a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/sdp.cpp b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/sdp.cpp index 202b41f7..f0d498a4 100644 --- a/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/sdp.cpp +++ b/omni/omni_xpu_kernel/omni_xpu_kernel/csrc/sdp.cpp @@ -22,6 +22,12 @@ #include #ifdef _WIN32 +#ifndef NOMINMAX +#define NOMINMAX +#endif +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif #include #else #include @@ -91,10 +97,31 @@ KernelLibrary& get_kernel_library() { path_buffer.resize(path_length); package_dir = fs::path(path_buffer).parent_path(); - fs::path library_path = package_dir / "lgrf_uni" / "lgrf_sdp.pyd"; + fs::path library_dir = package_dir / "lgrf_uni"; + TORCH_CHECK(fs::exists(library_dir), "missing lgrf sidecar directory: ", library_dir.string()); + + fs::path library_path; + for (const auto& entry : fs::directory_iterator(library_dir)) { + if (!entry.is_regular_file()) { + continue; + } + + const auto name = entry.path().filename().string(); + if (name.rfind("lgrf_sdp", 0) == 0 && entry.path().extension() == ".pyd") { + library_path = entry.path(); + break; + } + } + + if (!fs::exists(library_path)) { + load_error = "missing lgrf sidecar artifact under " + library_dir.string(); + return; + } + library.handle = LoadLibraryW(library_path.wstring().c_str()); if (library.handle == nullptr) { - load_error = "failed to load lgrf sidecar at " + library_path.string(); + load_error = "failed to load lgrf sidecar at " + library_path.string() + + " (WinError " + std::to_string(GetLastError()) + ")"; return; } diff --git a/omni/omni_xpu_kernel/setup.py b/omni/omni_xpu_kernel/setup.py index 140fbead..7f485b81 100644 --- a/omni/omni_xpu_kernel/setup.py +++ b/omni/omni_xpu_kernel/setup.py @@ -41,7 +41,16 @@ def get_icpx_path(): # Try common oneAPI installation paths on Windows program_files = os.environ.get("ProgramFiles", "C:\\Program Files") program_files_x86 = os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)") + preferred_version = os.environ.get("OMNI_XPU_ONEAPI_VERSION") candidates = [ + *( + [ + os.path.join(program_files_x86, "Intel", "oneAPI", "compiler", preferred_version, "bin", "icx.exe"), + os.path.join(program_files, "Intel", "oneAPI", "compiler", preferred_version, "bin", "icx.exe"), + ] + if preferred_version + else [] + ), os.path.join(program_files, "Intel", "oneAPI", "compiler", "latest", "bin", "icx.exe"), os.path.join(program_files, "Intel", "oneAPI", "compiler", "2025.1", "bin", "icx.exe"), os.path.join(program_files, "Intel", "oneAPI", "compiler", "2024.2", "bin", "icx.exe"), @@ -63,6 +72,47 @@ def get_icpx_path(): return None +def get_windows_oneapi_tool_dirs(): + """Find auxiliary oneAPI tool directories needed for Windows AOT builds.""" + if not IS_WINDOWS: + return [] + + program_files_x86 = Path(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)")) + oneapi_root = program_files_x86 / "Intel" / "oneAPI" + if not oneapi_root.exists(): + return [] + + preferred_version = os.environ.get("OMNI_XPU_ONEAPI_VERSION") + candidates = [] + patterns = [] + if preferred_version: + patterns.extend([ + f"compiler\\{preferred_version}\\bin\\compiler", + f"ocloc\\{preferred_version}\\bin", + ]) + patterns.extend([ + "compiler\\*\\bin\\compiler", + "ocloc\\*\\bin", + ]) + + for pattern in patterns: + matches = sorted(oneapi_root.glob(pattern), reverse=True) + for match in matches: + if match.is_dir(): + candidates.append(str(match)) + break + + return candidates + + +def should_enable_onednn(): + """Decide whether oneDNN-backed kernels should be built.""" + env_value = os.environ.get("OMNI_XPU_ENABLE_ONEDNN") + if env_value is not None: + return env_value.strip().lower() in {"1", "true", "yes", "on"} + return not IS_WINDOWS + + class ICPXBuildExt(build_ext): """Build extension using Intel icpx compiler directly.""" @@ -117,29 +167,58 @@ def build_extension(self, ext): onednn_include = os.environ.get("ONEDNN_INCLUDE", "") onednn_lib = os.environ.get("ONEDNN_LIB", "") - if not onednn_include or not onednn_lib: + want_onednn = should_enable_onednn() + + if want_onednn and (not onednn_include or not onednn_lib): # Auto-detect from common oneAPI paths - onednn_candidates = [ - "/opt/intel/oneapi/dnnl/2025.1", - "/opt/intel/oneapi/dnnl/latest", - "/opt/intel/oneapi/2025.1", - ] + if IS_WINDOWS: + program_files_x86 = os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)") + onednn_candidates = [ + os.path.join(program_files_x86, "Intel", "oneAPI", "dnnl", "latest"), + os.path.join(program_files_x86, "Intel", "oneAPI", "dnnl", "2026.0"), + os.path.join(program_files_x86, "Intel", "oneAPI", "dnnl", "2025.1"), + os.path.join(program_files_x86, "Intel", "oneAPI", "2026.0"), + ] + else: + onednn_candidates = [ + "/opt/intel/oneapi/dnnl/2025.1", + "/opt/intel/oneapi/dnnl/latest", + "/opt/intel/oneapi/2025.1", + ] for candidate in onednn_candidates: inc = os.path.join(candidate, "include") lib = os.path.join(candidate, "lib") - if os.path.exists(os.path.join(inc, "oneapi", "dnnl", "dnnl.hpp")): + if ( + os.path.exists(os.path.join(inc, "oneapi", "dnnl", "dnnl.hpp")) + or os.path.exists(os.path.join(inc, "dnnl.hpp")) + ): if not onednn_include: onednn_include = inc if not onednn_lib: onednn_lib = lib break - has_onednn = bool(onednn_include and os.path.isdir(onednn_include)) + has_onednn = bool( + want_onednn + and + onednn_include + and onednn_lib + and os.path.isdir(onednn_include) + and os.path.isdir(onednn_lib) + ) if has_onednn: print(f"oneDNN include: {onednn_include}") print(f"oneDNN lib: {onednn_lib}") - else: + elif want_onednn: print("WARNING: oneDNN not found. onednn_int4_gemm will not be available.") + else: + print("oneDNN disabled for this build.") + + if not has_onednn: + sources = [ + source for source in sources + if source.name not in {"onednn_int4_gemm.cpp", "onednn_fp8.cpp"} + ] if IS_WINDOWS: # Windows compile command using icx @@ -165,11 +244,17 @@ def build_extension(self, ext): cmd.append(f"/I{onednn_include}") cmd += [str(s) for s in sources] else: + device_target = os.environ.get("OMNI_XPU_DEVICE", "bmg") cmd += [ + "-fsycl-targets=spir64_gen", + "-Xs", f"-device {device_target}", "-fsycl-esimd-force-stateless-mem", "/O2", "/DNDEBUG", "/EHsc", # Enable C++ exception handling "/std:c++17", + "/DNOMINMAX", + "/DWIN32_LEAN_AND_MEAN", + f"/DOMNI_XPU_HAS_ONEDNN={1 if has_onednn else 0}", f"/I{python_include}", f"/I{torch_include}", f"/I{torch_include}\\torch\\csrc\\api\\include", @@ -215,6 +300,7 @@ def build_extension(self, ext): "-O3", "-DNDEBUG", "-fPIC", "-shared", "-std=c++17", + f"-DOMNI_XPU_HAS_ONEDNN={1 if has_onednn else 0}", f"-I{python_include}", f"-I{torch_include}", f"-I{torch_include}/torch/csrc/api/include", @@ -237,7 +323,14 @@ def build_extension(self, ext): print(f"Compile command: {' '.join(cmd)}") # Run compiler - result = subprocess.run(cmd, capture_output=True, text=True) + env = os.environ.copy() + if IS_WINDOWS: + extra_tool_dirs = get_windows_oneapi_tool_dirs() + if extra_tool_dirs: + print(f"Adding oneAPI tool dirs to PATH: {extra_tool_dirs}") + env["PATH"] = os.pathsep.join(extra_tool_dirs + [env.get("PATH", "")]) + + result = subprocess.run(cmd, capture_output=True, text=True, env=env) if result.returncode != 0: print("STDOUT:", result.stdout) @@ -258,7 +351,7 @@ def __init__(self, name, sourcedir=""): # Read version def get_version(): version_file = Path(__file__).parent / "omni_xpu_kernel" / "__init__.py" - with open(version_file) as f: + with open(version_file, encoding="utf-8") as f: for line in f: if line.startswith("__version__"): return line.split("=")[1].strip().strip('"\'')