diff --git a/README.md b/README.md index 7ecee701d6a..73c37fba485 100644 --- a/README.md +++ b/README.md @@ -104,10 +104,30 @@ If you happen to forget the `--recursive` during `clone`, you can use the follow git submodule sync && git submodule update --init --recursive ``` +### Native Windows ROCm (experimental) + +Windows keeps the Triton-only installation mode by default. To build the HIP +and Composable Kernel extensions, opt in with `AITER_ENABLE_HIP=1`. Native GPU +detection supports gfx1100 through gfx1103, gfx1151, and gfx1201; `GPU_ARCHS` +can be set explicitly for an offline or cross build. + +```powershell +$env:AITER_ENABLE_HIP = "1" +$env:GPU_ARCHS = "gfx1100" # optional when building on the target GPU +python -m pip install --no-build-isolation . +``` + +GPU models sharing an architecture can expose different CU counts. When +cross-building for a binned model, set `CU_NUM` to the target device's actual +compute-unit count. + ### FlyDSL AITER uses [FlyDSL](https://github.com/ROCm/FlyDSL)-based kernels across a range of operators (e.g., GEMM and MoE). FlyDSL is a required dependency and is installed automatically when you run `python3 setup.py develop`. +FlyDSL is not installed or AOT-compiled on Windows; the experimental native +Windows build uses Composable Kernel instead. + To install it manually: ```bash diff --git a/aiter/__init__.py b/aiter/__init__.py index 569b5295a9f..59dadb16fb2 100644 --- a/aiter/__init__.py +++ b/aiter/__init__.py @@ -58,11 +58,18 @@ def getLogger(): logger = getLogger() AITER_AOT_IMPORT = os.getenv("AITER_AOT_IMPORT", "0") == "1" +from .utility import dtypes as dtypes # noqa: E402 + # Triton-only: expose only the Triton ops, skipping the C++/CK/HIP ops and their -# JIT build. Always on for Windows (no CK/HIP there); elsewhere opt in via the -# env var, e.g. Triton-backend users with no C++ toolchain or CK. +# JIT build. Native Windows ROCm support is opt-in so existing Windows installs +# keep the upstream Triton-only default unless AITER_ENABLE_HIP=1 is set. +AITER_ENABLE_HIP = os.getenv("AITER_ENABLE_HIP", "0") == "1" or ( + sys.platform == "win32" + and os.path.isfile(os.path.join(os.path.dirname(__file__), "jit", "module_aiter_core.pyd")) +) AITER_TRITON_ONLY = ( - os.getenv("AITER_TRITON_ONLY", "0") == "1" or sys.platform == "win32" + os.getenv("AITER_TRITON_ONLY", "0") == "1" + or (sys.platform == "win32" and not AITER_ENABLE_HIP) ) # Use bundled pre-compiled FlyDSL cache unless the user overrides via env var. @@ -85,7 +92,6 @@ def getLogger(): # non-gfx950) inside aiter/ops/opus/__init__.py, so its import line # is safe to put at top-level without try/except. from .jit import core as core # noqa: E402 - from .utility import dtypes as dtypes # noqa: E402 from .ops.enum import * # noqa: F403,E402 from .ops.norm import * # noqa: F403,E402 from .ops.quant import * # noqa: F403,E402 @@ -111,7 +117,8 @@ def getLogger(): from .ops.pos_encoding import * # noqa: F403,E402 from .ops.cache import * # noqa: F403,E402 from .ops.rmsnorm import * # noqa: F403,E402 - from .ops.communication import * # noqa: F403,E402 + if sys.platform != "win32" or torch.distributed.is_available(): + from .ops.communication import * # noqa: F403,E402 from .ops.rope import * # noqa: F403,E402 from .ops.topk import * # noqa: F403,E402 from .ops.topk_plain import topk_plain # noqa: F403,F401,E402 diff --git a/aiter/jit/core.py b/aiter/jit/core.py index 5652b97a438..333d9d76ea4 100644 --- a/aiter/jit/core.py +++ b/aiter/jit/core.py @@ -8,9 +8,10 @@ import multiprocessing import os import re -import shlex import shutil +import subprocess import sys +import tempfile import time import traceback import types @@ -22,8 +23,14 @@ this_dir = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, f"{this_dir}/utils/") from chip_info import get_gfx, get_gfx_list, get_gfx_runtime # noqa: E402 -from cpp_extension import _jit_compile, executable_path, get_hip_version # noqa: E402 +from cpp_extension import ( # noqa: E402 + LIB_EXT, + _jit_compile, + executable_path, + get_hip_version, +) from file_baton import FileBaton # noqa: E402 +from blob_gen import windows_blob_gen_argv # noqa: E402 from torch_guard import torch_compile_guard # noqa: E402 AITER_REBUILD = int(os.environ.get("AITER_REBUILD", "0")) @@ -331,7 +338,7 @@ def update_config_files(self, file_path: str, merge_name: str): ) from pathlib import Path - config_path = Path("/tmp/aiter_configs/") + config_path = Path(tempfile.gettempdir()) / "aiter_configs" if not config_path.exists(): config_path.mkdir(parents=True, exist_ok=True) new_file_path = f"{config_path}/{merge_name}.csv" @@ -492,14 +499,22 @@ def validate_and_update_archs(): def hip_flag_checker(flag_hip: str) -> bool: import subprocess + null_device = "NUL" if sys.platform == "win32" else "/dev/null" cmd = ( [executable_path("hipcc")] + flag_hip.split() - + ["-x", "hip", "-E", "-P", "/dev/null", "-o", "/dev/null"] + + ["-x", "hip", "-E", "-P", "-", "-o", null_device] ) try: - subprocess.check_output(cmd, stderr=subprocess.DEVNULL) - except subprocess.CalledProcessError: + subprocess.run( + cmd, + input="", + text=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=True, + ) + except (OSError, subprocess.CalledProcessError): logger.warning(f"Current hipcc not support: {flag_hip}, skip it.") return False return True @@ -516,16 +531,41 @@ def check_LLVM_MAIN_REVISION(): import subprocess try: - hipcc = shlex.quote(executable_path("hipcc")) - cmd = f"""echo "#include -__host__ __device__ void func(){{std::tuple t = std::tuple(1, 1);}}" | {hipcc} -x hip -P -c -Wno-unused-command-line-argument -o /dev/null -""" - subprocess.check_output(cmd, shell=True, text=True, stderr=subprocess.STDOUT) - except (subprocess.CalledProcessError, AssertionError): + source = """#include +__host__ __device__ void func(){std::tuple t = std::tuple(1, 1);}""" + null_device = "NUL" if sys.platform == "win32" else "/dev/null" + subprocess.run( + [ + executable_path("hipcc"), + "-x", + "hip", + "-P", + "-c", + "-Wno-unused-command-line-argument", + "-o", + null_device, + "-", + ], + input=source, + text=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.STDOUT, + check=True, + ) + except (OSError, subprocess.CalledProcessError, AssertionError): return 554785 return 554785 - 1 def check_and_set_ninja_worker(): + max_jobs_env = os.environ.get("MAX_JOBS") + if max_jobs_env is not None: + try: + if int(max_jobs_env) > 0: + return + except ValueError: + pass + max_num_jobs_cores = max(1, os.cpu_count() * 0.8) import psutil @@ -535,19 +575,7 @@ def check_and_set_ninja_worker(): # pick lower value of jobs based on cores vs memory metric to minimize oom and swap usage during compilation max_jobs = int(max(1, min(max_num_jobs_cores, max_num_jobs_memory))) - max_jobs_env = os.environ.get("MAX_JOBS") - if max_jobs_env is not None: - try: - max_processes = int(max_jobs_env) - # too large value - if max_processes > max_jobs: - os.environ["MAX_JOBS"] = str(max_jobs) - # error value - except ValueError: - os.environ["MAX_JOBS"] = str(max_jobs) - # none value - else: - os.environ["MAX_JOBS"] = str(max_jobs) + os.environ["MAX_JOBS"] = str(max_jobs) def rename_cpp_to_cu(els, dst, hipify, recursive=False): @@ -583,6 +611,8 @@ def do_rename_and_mv(name, src, dst, ret): @torch_compile_guard() def check_numa_custom_op() -> None: + if sys.platform == "win32": + return numa_balance_set = os.popen("cat /proc/sys/kernel/numa_balancing").read().strip() if numa_balance_set == "1": logger.warning( @@ -640,7 +670,7 @@ def _needs_arch_rebuild(md_name): except Exception: # running arch undetectable (e.g. no GPU) -> keep normal behaviour return False - so_path = os.path.join(get_user_jit_dir(), f"{md_name}.so") + so_path = os.path.join(get_user_jit_dir(), f"{md_name}{LIB_EXT}") built = _so_offload_archs(so_path) if not built or cur in built: return False @@ -789,11 +819,13 @@ def check_git_version(required_major, required_minor): def rm_module(md_name): - os.system(f"rm -rf {get_user_jit_dir()}/{md_name}.so") + module_path = os.path.join(get_user_jit_dir(), f"{md_name}{LIB_EXT}") + if os.path.isfile(module_path): + os.remove(module_path) def clear_build(md_name): - os.system(f"rm -rf {bd_dir}/{md_name}") + shutil.rmtree(os.path.join(bd_dir, md_name), ignore_errors=True) def build_module( @@ -815,7 +847,7 @@ def build_module( os.makedirs(bd_dir, exist_ok=True) lock_path = f"{bd_dir}/lock_{md_name}" startTS = time.perf_counter() - target_name = f"{md_name}.so" if not is_standalone else md_name + target_name = f"{md_name}{LIB_EXT}" if not is_standalone else md_name for tp in third_party: clone_3rdparty(tp) @@ -923,9 +955,20 @@ def exec_blob(blob_gen_cmd, op_dir, src_dir, sources): if blob_gen_cmd: blob_dir = f"{op_dir}/blob/" os.makedirs(blob_dir, exist_ok=True) - if AITER_LOG_MORE: - logger.info(f"exec_blob ---> {PY} {blob_gen_cmd.format(blob_dir)}") - os.system(f"{PY} {blob_gen_cmd.format(blob_dir)}") + if sys.platform == "win32": + blob_gen_argv = windows_blob_gen_argv(PY, blob_gen_cmd, blob_dir) + if AITER_LOG_MORE: + logger.info( + f"exec_blob ---> {subprocess.list2cmdline(blob_gen_argv)}" + ) + subprocess.run(blob_gen_argv, check=True) + else: + blob_gen_shell_cmd = ( + f"{PY} {blob_gen_cmd.format(blob_dir)}" + ) + if AITER_LOG_MORE: + logger.info(f"exec_blob ---> {blob_gen_shell_cmd}") + os.system(blob_gen_shell_cmd) sources += rename_cpp_to_cu([blob_dir], src_dir, hipify, recursive=True) return sources @@ -1255,7 +1298,7 @@ def _ctypes_call(func, fc_name, md_name): def _ensure_loaded(): if _cache: return - so_path = os.path.join(get_user_jit_dir(), f"{md_name}.so") + so_path = os.path.join(get_user_jit_dir(), f"{md_name}{LIB_EXT}") if not os.path.exists(so_path) or _needs_arch_rebuild(md_name): d_args = get_args_of_build(md_name) d_args["torch_exclude"] = True diff --git a/aiter/jit/utils/blob_gen.py b/aiter/jit/utils/blob_gen.py new file mode 100644 index 00000000000..4b3b78bbdc8 --- /dev/null +++ b/aiter/jit/utils/blob_gen.py @@ -0,0 +1,67 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import re +import shlex +from typing import List + + +def windows_blob_gen_argv( + python_executable: str, blob_gen_cmd: str, blob_dir: str +) -> List[str]: + """Convert a blob generator command to argv without invoking cmd.exe. + + Blob generator commands are stored as shell-like strings whose first + argument is always a Python script and whose remaining arguments are + option/value pairs. On Windows, paths in those strings may contain spaces + (including the checkout, Python installation, output, and tune-file paths). + Group each option's value before passing the result to CreateProcess so no + shell quoting is required. + """ + command = blob_gen_cmd.format(blob_dir).strip() + script_match = re.match(r"^(.*?\.py)(?:\s+(.*))?$", command, re.IGNORECASE) + if script_match is None: + raise ValueError( + f"Blob generator command must start with a .py script: {command}" + ) + + script, raw_args = script_match.groups() + argv = [python_executable, script] + if not raw_args: + return argv + + lexer = shlex.shlex(raw_args, posix=False) + lexer.whitespace_split = True + lexer.commenters = "" + + option = None + values: List[str] = [] + + def append_option(): + nonlocal option, values + if option is None: + return + argv.append(option) + if values: + value = " ".join(values) + if len(value) >= 2 and value[0] == value[-1] and value[0] in "\"'": + value = value[1:-1] + argv.append(value) + option = None + values = [] + + for token in lexer: + is_option = re.match(r"^-{1,2}[A-Za-z]", token) is not None + if is_option: + append_option() + if "=" in token: + option, first_value = token.split("=", 1) + values.append(first_value) + else: + option = token + elif option is None: + raise ValueError(f"Unexpected positional blob generator argument: {token}") + else: + values.append(token) + append_option() + return argv diff --git a/aiter/jit/utils/build_targets.py b/aiter/jit/utils/build_targets.py index 5a1a93b421c..0b93ba95f0a 100644 --- a/aiter/jit/utils/build_targets.py +++ b/aiter/jit/utils/build_targets.py @@ -39,6 +39,12 @@ GFX_CU_NUM_MAP = { "gfx942": 304, # MI300X (SPX, full GPU); MI308X shares gfx942 — use CU_NUM override "gfx950": 256, # MI350 + "gfx1100": 96, # Radeon RX 7900 XTX + "gfx1101": 60, # Radeon RX 7800 XT + "gfx1102": 32, # Radeon RX 7600 + "gfx1103": 12, # Radeon 780M + "gfx1151": 20, # Strix Halo / Radeon 8060S + "gfx1201": 64, # Radeon RX 9070 XT / Radeon AI PRO R9700 "gfx1250": 256, # Gfx1250 } diff --git a/aiter/jit/utils/chip_info.py b/aiter/jit/utils/chip_info.py index 50ff1068ab2..315dec0c711 100644 --- a/aiter/jit/utils/chip_info.py +++ b/aiter/jit/utils/chip_info.py @@ -5,22 +5,44 @@ import os import re import subprocess - -from cpp_extension import executable_path -from torch_guard import torch_compile_guard - -from build_targets import ( # noqa: F401 -- re-exported for callers - GFX_MAP, - _parse_gpu_archs_env, - filter_tune_df, - get_build_targets_env, -) +import sys + +try: + from .cpp_extension import executable_path + from .torch_guard import torch_compile_guard + from .build_targets import ( # noqa: F401 -- re-exported for callers + GFX_MAP, + _parse_gpu_archs_env, + filter_tune_df, + get_build_targets_env, + ) +except ImportError: + # core.py also imports this module directly after adding utils/ to sys.path. + from cpp_extension import executable_path + from torch_guard import torch_compile_guard + from build_targets import ( # noqa: F401 -- re-exported for callers + GFX_MAP, + _parse_gpu_archs_env, + filter_tune_df, + get_build_targets_env, + ) logger = logging.getLogger("aiter") @functools.lru_cache(maxsize=1) def _detect_native() -> list[str]: + if sys.platform == "win32": + try: + import torch + + props = torch.cuda.get_device_properties(0) + arch = props.gcnArchName.split(":", 1)[0].lower() + if arch.startswith("gfx"): + return [arch] + except Exception as e: + raise RuntimeError(f"Get GPU arch from PyTorch failed: {e}") from e + try: rocminfo = executable_path("rocminfo") result = subprocess.run( @@ -141,6 +163,10 @@ def get_gfx_list() -> list[str]: def get_cu_num_custom_op() -> int: cu_num = int(os.getenv("CU_NUM", 0)) if cu_num == 0: + if sys.platform == "win32": + import torch + + return int(torch.cuda.get_device_properties(0).multi_processor_count) try: rocminfo = executable_path("rocminfo") result = subprocess.run( diff --git a/aiter/jit/utils/cpp_extension.py b/aiter/jit/utils/cpp_extension.py index 444a5669f51..28de5a23502 100644 --- a/aiter/jit/utils/cpp_extension.py +++ b/aiter/jit/utils/cpp_extension.py @@ -20,27 +20,42 @@ from typing import Dict, List, Optional, Tuple, Union import setuptools -from _cpp_extension_versioner import ExtensionVersioner -from file_baton import FileBaton -from hipify import hipify_python -from hipify.hipify_python import GeneratedFileCleaner +try: + from ._cpp_extension_versioner import ExtensionVersioner + from .file_baton import FileBaton + from .hipify import hipify_python + from .hipify.hipify_python import GeneratedFileCleaner +except ImportError: + from _cpp_extension_versioner import ExtensionVersioner + from file_baton import FileBaton + from hipify import hipify_python + from hipify.hipify_python import GeneratedFileCleaner from packaging.version import Version from setuptools.command.build_ext import build_ext IS_WINDOWS = sys.platform == "win32" IS_LINUX = sys.platform.startswith("linux") -LIB_EXT = ".so" -EXEC_EXT = "" -CLIB_PREFIX = "lib" -CLIB_EXT = ".so" +LIB_EXT = ".pyd" if IS_WINDOWS else ".so" +EXEC_EXT = ".exe" if IS_WINDOWS else "" +CLIB_PREFIX = "" if IS_WINDOWS else "lib" +CLIB_EXT = ".dll" if IS_WINDOWS else ".so" SHARED_FLAG = "-shared" -SUBPROCESS_DECODE_ARGS = () +SUBPROCESS_DECODE_ARGS = ("oem",) if IS_WINDOWS else () MINIMUM_GCC_VERSION = (5, 0, 0) MINIMUM_MSVC_VERSION = (19, 0, 24215) VersionRange = Tuple[Tuple[int, ...], Tuple[int, ...]] VersionMap = Dict[str, VersionRange] + + +def _nt_quote_args(args: Optional[List[str]]) -> List[str]: + """Quote command-line arguments using Windows command-line rules.""" + if not args: + return [] + return [f'"{arg}"' if " " in arg else arg for arg in args] + + # The following values were taken from the following GitHub gist that # summarizes the minimum valid major versions of g++/clang++ for each supported # CUDA version: https://gist.github.com/ax3l/9489132 @@ -155,7 +170,15 @@ def _find_rocm_home() -> Optional[str]: spec = None if spec is not None and spec.submodule_search_locations: candidate = spec.submodule_search_locations[0] - if os.path.exists(os.path.join(candidate, "bin", "hipconfig")): + hipconfig_names = ( + ("hipconfig.exe", "hipconfig") + if IS_WINDOWS + else ("hipconfig",) + ) + if any( + os.path.exists(os.path.join(candidate, "bin", name)) + for name in hipconfig_names + ) and os.path.exists(os.path.join(candidate, "lib")): rocm_home = candidate if rocm_home is None: # Guess #3 @@ -293,23 +316,32 @@ def _join_rocm_home(*paths) -> str: ] COMMON_HIP_FLAGS = [ - "-fPIC", "-D__HIP_PLATFORM_AMD__=1", "-DUSE_ROCM=1", "-DHIPBLAS_V2", ] +if not IS_WINDOWS: + COMMON_HIP_FLAGS.append("-fPIC") COMMON_HIPCC_FLAGS = [ "-DCUDA_HAS_FP16=1", "-D__HIP_NO_HALF_OPERATORS__=1", "-D__HIP_NO_HALF_CONVERSIONS__=1", - "-mcmodel=large", - "-fno-unique-section-names", - "-ffunction-sections", - "-fdata-sections", ] -if not int(os.environ.get("AITER_SYMBOL_VISIBLE", "0")): +if IS_WINDOWS: + COMMON_HIPCC_FLAGS.extend(["-fms-extensions", "-Wno-ignored-attributes"]) +else: + COMMON_HIPCC_FLAGS.extend( + [ + "-mcmodel=large", + "-fno-unique-section-names", + "-ffunction-sections", + "-fdata-sections", + ] + ) + +if not IS_WINDOWS and not int(os.environ.get("AITER_SYMBOL_VISIBLE", "0")): COMMON_HIPCC_FLAGS.extend(["-fvisibility=hidden", "-fvisibility-inlines-hidden"]) JIT_EXTENSION_VERSIONER = ExtensionVersioner() @@ -320,8 +352,19 @@ def _join_rocm_home(*paths) -> str: } +def _get_vc_env(vc_arch: str) -> Dict[str, str]: + try: + from setuptools._distutils import _msvccompiler + + return _msvccompiler._get_vc_env(vc_arch) + except AttributeError: + from setuptools._distutils.compilers.C import msvc + + return msvc._get_vc_env(vc_arch) + + def get_cxx_compiler(): - return os.environ.get("CXX", "c++") + return os.environ.get("CXX", "cl" if IS_WINDOWS else "c++") def _is_binary_build() -> bool: @@ -1475,29 +1518,54 @@ def verify_ninja_availability(): def _prepare_ldflags(extra_ldflags, with_cuda, verbose, is_standalone, torch_exclude): - extra_ldflags.append("-mcmodel=large") - extra_ldflags.append("-ffunction-sections") - extra_ldflags.append("-fdata-sections ") - extra_ldflags.append("-Wl,--gc-sections") - extra_ldflags.append("-Wl,--cref") + if not IS_WINDOWS: + extra_ldflags.append("-mcmodel=large") + extra_ldflags.append("-ffunction-sections") + extra_ldflags.append("-fdata-sections") + extra_ldflags.append("-Wl,--gc-sections") + extra_ldflags.append("-Wl,--cref") + if not torch_exclude: import torch _TORCH_PATH = os.path.join(os.path.dirname(torch.__file__)) TORCH_LIB_PATH = os.path.join(_TORCH_PATH, "lib") - extra_ldflags.append(f"-L{TORCH_LIB_PATH}") - extra_ldflags.append("-lc10") - if with_cuda: - extra_ldflags.append("-lc10_hip" if IS_HIP_EXTENSION else "-lc10_cuda") - extra_ldflags.append("-ltorch_cpu") - if with_cuda: - extra_ldflags.append("-ltorch_hip" if IS_HIP_EXTENSION else "-ltorch_cuda") - extra_ldflags.append("-ltorch") - if not is_standalone: - extra_ldflags.append("-ltorch_python") - - if is_standalone: - extra_ldflags.append(f"-Wl,-rpath,{TORCH_LIB_PATH}") + if IS_WINDOWS: + extra_ldflags.extend( + [ + f"-L{TORCH_LIB_PATH}", + "-lc10", + *(["-lc10_hip"] if with_cuda and IS_HIP_EXTENSION else []), + "-ltorch_cpu", + *(["-ltorch_hip"] if with_cuda and IS_HIP_EXTENSION else []), + "-ltorch", + ] + ) + if not is_standalone: + extra_ldflags.append("-ltorch_python") + else: + extra_ldflags.append(f"-L{TORCH_LIB_PATH}") + extra_ldflags.append("-lc10") + if with_cuda: + extra_ldflags.append( + "-lc10_hip" if IS_HIP_EXTENSION else "-lc10_cuda" + ) + extra_ldflags.append("-ltorch_cpu") + if with_cuda: + extra_ldflags.append( + "-ltorch_hip" if IS_HIP_EXTENSION else "-ltorch_cuda" + ) + extra_ldflags.append("-ltorch") + if not is_standalone: + extra_ldflags.append("-ltorch_python") + + if is_standalone: + extra_ldflags.append(f"-Wl,-rpath,{TORCH_LIB_PATH}") + + if IS_WINDOWS and not is_standalone: + python_lib_path = os.path.join(sys.base_exec_prefix, "libs") + python_lib = f"-lpython{sys.version_info.major}{sys.version_info.minor}" + extra_ldflags.extend([f"-L{python_lib_path}", python_lib]) if with_cuda and IS_HIP_EXTENSION: if verbose: @@ -1536,8 +1604,7 @@ def _get_rocm_arch_flags(cflags: Optional[List[str]] = None) -> List[str]: def _get_num_workers(verbose: bool) -> Optional[int]: max_jobs = os.environ.get("MAX_JOBS") if max_jobs is not None and max_jobs.isdigit(): - if int(max_jobs) > int(max(1, os.cpu_count() * 0.8)): - max_jobs = int(max(1, os.cpu_count() * 0.8)) + max_jobs = int(max_jobs) if verbose: print( f"Using envvar MAX_JOBS ({max_jobs}) as the number of workers...", @@ -1561,6 +1628,14 @@ def _run_ninja_build(build_directory: str, verbose: bool, error_prefix: str) -> if num_workers is not None: command.extend(["-j", str(num_workers)]) env = os.environ.copy() + if IS_WINDOWS and "VSCMD_ARG_TGT_ARCH" not in env: + from setuptools import distutils + + plat_spec = PLAT_TO_VCVARS[distutils.util.get_platform()] + vc_env = {k.upper(): v for k, v in _get_vc_env(plat_spec).items()} + for key, value in env.items(): + vc_env.setdefault(key.upper(), value) + env = vc_env try: sys.stdout.flush() @@ -1580,6 +1655,7 @@ def _run_ninja_build(build_directory: str, verbose: bool, error_prefix: str) -> stdout_fileno = 1 subprocess.run( command, + shell=IS_WINDOWS and IS_HIP_EXTENSION, stdout=stdout_fileno if verbose else subprocess.PIPE, stderr=subprocess.STDOUT, cwd=build_directory, @@ -1669,7 +1745,9 @@ def _write_ninja_file_to_build_library( # Explicitly specify 'posix_prefix' scheme on non-Windows platforms to workaround error on some MacOS # installations where default `get_path` points to non-existing `/Library/Python/M.m/include` folder if is_python_module: - python_include_path = sysconfig.get_path("include", scheme="posix_prefix") + python_include_path = sysconfig.get_path( + "include", scheme="nt" if IS_WINDOWS else "posix_prefix" + ) if python_include_path is not None: system_includes.append(python_include_path) @@ -1683,16 +1761,31 @@ def _write_ninja_file_to_build_library( # common_cflags += [f"{x}" for x in _get_pybind11_abi_build_flags()] # common_cflags += [f"{x}" for x in _get_glibcxx_abi_build_flags()] - # Windows does not understand `-isystem` and quotes flags later. - common_cflags += [f"-I{shlex.quote(include)}" for include in user_includes] - common_cflags += [f"-isystem {shlex.quote(include)}" for include in system_includes] - - cflags = common_cflags + ["-fPIC", "-std=c++20"] + extra_cflags + if IS_WINDOWS: + common_cflags += [f"-I{include}" for include in user_includes + system_includes] + cflags = common_cflags + ["/std:c++20"] + extra_cflags + cflags += COMMON_MSVC_FLAGS + cflags = _nt_quote_args(cflags) + else: + common_cflags += [f"-I{shlex.quote(include)}" for include in user_includes] + common_cflags += [ + f"-isystem {shlex.quote(include)}" for include in system_includes + ] + cflags = common_cflags + ["-fPIC", "-std=c++20"] + extra_cflags if with_cuda and IS_HIP_EXTENSION: - cuda_flags = ["-DWITH_HIP"] + cflags + COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS - cuda_flags += extra_cuda_cflags + cuda_flags = ( + ["-DWITH_HIP"] + + common_cflags + + extra_cflags + + COMMON_HIP_FLAGS + + COMMON_HIPCC_FLAGS + + ["-std=c++20"] + ) cuda_flags += _get_rocm_arch_flags(cuda_flags) + cuda_flags += extra_cuda_cflags + if IS_WINDOWS: + cuda_flags = _nt_quote_args(cuda_flags) def object_file_path(source_file: str) -> str: # '/path/to/file.cpp' -> 'file' @@ -1707,6 +1800,8 @@ def object_file_path(source_file: str) -> str: objects = [object_file_path(src) for src in sources] ldflags = ([] if is_standalone else [SHARED_FLAG]) + extra_ldflags + if IS_WINDOWS: + ldflags = _nt_quote_args(ldflags) ext = EXEC_EXT if is_standalone else LIB_EXT library_target = f"{name}{ext}" @@ -1785,8 +1880,14 @@ def sanitize_flags(flags): # Version 1.3 is required for the `deps` directive. config = ["ninja_required_version = 1.3"] config.append(f"cxx = {compiler}") + linker = executable_path("hipcc") if IS_WINDOWS and IS_HIP_EXTENSION else compiler + config.append(f"linker = {linker}") if with_cuda or cuda_dlink_post_cflags: - nvcc = _join_rocm_home("bin", "hipcc") + nvcc = ( + executable_path("hipcc") + if IS_WINDOWS + else _join_rocm_home("bin", "hipcc") + ) config.append(f"nvcc = {nvcc}") if IS_HIP_EXTENSION: @@ -1804,11 +1905,16 @@ def sanitize_flags(flags): # See https://ninja-build.org/build.ninja.html for reference. compile_rule = ["rule compile"] - compile_rule.append( - " command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags" - ) - compile_rule.append(" depfile = $out.d") - compile_rule.append(" deps = gcc") + if IS_WINDOWS: + compile_rule.append( + " command = $cxx /showIncludes $cflags -c $in /Fo$out $post_cflags" + ) + else: + compile_rule.append( + " command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags" + ) + compile_rule.append(" depfile = $out.d") + compile_rule.append(" deps = gcc") if with_cuda: cuda_compile_rule = ["rule cuda_compile"] @@ -1851,6 +1957,9 @@ def _resolve_per_source_flags(src_abs): _resolve_per_source_flags(source_file) if is_cuda_source else None ) + if IS_WINDOWS: + source_file = source_file.replace(":", "$:") + object_file = object_file.replace(":", "$:") source_file_q = source_file.replace(" ", "$ ") object_file_q = object_file.replace(" ", "$ ") build.append(f"build {object_file_q}: {rule} {source_file_q}") @@ -1875,9 +1984,10 @@ def _resolve_per_source_flags(src_abs): if library_target is not None: link_rule = ["rule link"] - link_rule.append( - " command = $cxx @$out.rsp $ldflags -o $out\n rspfile = $out.rsp\n rspfile_content = $in" + " command = $linker @$out.rsp $ldflags -o $out\n" + " rspfile = $out.rsp\n" + " rspfile_content = $in" ) link = [f'build {library_target}: link {" ".join(objects)}'] diff --git a/aiter/jit/utils/file_baton.py b/aiter/jit/utils/file_baton.py index 5c47a3d63e9..4922d930187 100644 --- a/aiter/jit/utils/file_baton.py +++ b/aiter/jit/utils/file_baton.py @@ -114,6 +114,12 @@ def _pid_alive(pid): return False except PermissionError: return True # exists but owned by another user + except OSError: + # Windows reports a dead or otherwise invalid PID as WinError 87 + # rather than ProcessLookupError. + if os.name == "nt": + return False + raise return True def _is_stale(self): diff --git a/aiter/jit/utils/mha_recipes.py b/aiter/jit/utils/mha_recipes.py index 4430ed432bc..2d025fcc549 100755 --- a/aiter/jit/utils/mha_recipes.py +++ b/aiter/jit/utils/mha_recipes.py @@ -1,23 +1,30 @@ -from typing import List, Dict, Tuple +from typing import Dict, List, Tuple + + +def _ck_targets_flag_for_arch(gfx: str) -> str: + if gfx.startswith("gfx9"): + return "" + return f" --targets {gfx}" def _ck_targets_flag() -> str: """Return ``--targets `` when the runtime GPU is not gfx9. ck-tile's ``generate.py`` defaults to ``--targets gfx9,gfx950``, so any - non-gfx9 host (gfx10/11/12) ends up with an empty kernel set and ``mha_fwd`` - fails at dispatch with "invalid argument for fmha_fwd". For gfx9 hosts we - keep the default (covers both gfx942 and gfx950 like before). + RDNA host (gfx11/12) ends up with an empty kernel set and ``mha_fwd`` fails + at dispatch with "invalid argument for fmha_fwd". For gfx9 hosts we keep + the default (covers both gfx942 and gfx950 like before). """ try: + from .chip_info import get_gfx + except ImportError: from chip_info import get_gfx + try: gfx = get_gfx() except Exception: return "" - if gfx.startswith("gfx9"): - return "" - return f" --targets {gfx}" + return _ck_targets_flag_for_arch(gfx) def compose_mha_fwd_variant_suffix_and_filter( diff --git a/aiter/jit/utils/torch_guard.py b/aiter/jit/utils/torch_guard.py index dd70debe134..099f47c21cd 100644 --- a/aiter/jit/utils/torch_guard.py +++ b/aiter/jit/utils/torch_guard.py @@ -342,10 +342,10 @@ def outer_wrapper_dummy(dummy, *args, **kwargs): tags = () else: tags = (torch.Tag.needs_fixed_stride_order,) - op_schema = f"aiter::{loadName}" + schema + op_schema = f"{loadName}" + schema aiter_lib.define(op_schema, tags=tags) - aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CUDA") - aiter_lib.impl(f"aiter::{loadName}", custom_func, dispatch_key="CPU") + aiter_lib.impl(loadName, custom_func, dispatch_key="CUDA") + aiter_lib.impl(loadName, custom_func, dispatch_key="CPU") aiter_lib._register_fake(f"{loadName}", fake_func) return wrapper_custom diff --git a/aiter/ops/mha.py b/aiter/ops/mha.py index 5c3464a310b..06642a999ee 100644 --- a/aiter/ops/mha.py +++ b/aiter/ops/mha.py @@ -11,6 +11,7 @@ from ..jit.utils.chip_info import get_cu_num, get_gfx from ..jit.utils.torch_guard import torch_compile_guard from ..jit.utils.mha_recipes import ( + _ck_targets_flag, compose_mha_fwd_variant_suffix_and_filter, get_mha_varlen_prebuild_variants_by_names, ) @@ -97,7 +98,9 @@ def cmdGenFunc_mha_fwd( blob_gen_cmd = [ f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd " - "--receipt 100 --filter {} --output_dir {{}}".format(filter), + "--receipt 100 --filter {} --output_dir {{}}{}".format( + filter, _ck_targets_flag() + ), ] return { "md_name": md_name, @@ -831,11 +834,15 @@ def cmdGenFunc_mha_varlen_fwd( filter_fwd_splitkv = f"{filter_fwd_splitkv1}@{filter_fwd_splitkv2}" blob_gen_cmd = [ f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd " - "--receipt 200 --filter {} --output_dir {{}}".format('" "') + "--receipt 200 --filter {} --output_dir {{}}{}".format( + '" "', _ck_targets_flag() + ) ] blob_gen_cmd.append( f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d fwd_splitkv " - "--receipt 200 --filter {} --output_dir {{}}".format(filter_fwd_splitkv) + "--receipt 200 --filter {} --output_dir {{}}{}".format( + filter_fwd_splitkv, _ck_targets_flag() + ) ) return { "md_name": md_name, @@ -1502,7 +1509,9 @@ def cmdGenFunc_mha_batch_prefill( filter_fwd += "_nsink*" blob_gen_cmd = [ f"{CK_DIR}/example/ck_tile/01_fmha/generate.py -d batch_prefill " - "--receipt 200 --filter {} --output_dir {{}}".format(filter_fwd) + "--receipt 200 --filter {} --output_dir {{}}{}".format( + filter_fwd, _ck_targets_flag() + ) ] return { "md_name": md_name, diff --git a/aiter/utility/dtypes.py b/aiter/utility/dtypes.py index 2dc70310b1c..601fb4d9968 100644 --- a/aiter/utility/dtypes.py +++ b/aiter/utility/dtypes.py @@ -1,10 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. import torch -from ..jit.utils.chip_info import get_gfx_runtime -from ..ops.enum import QuantType, ActivationType from .aiter_types import aiter_dtypes, aiter_tensor_t import argparse +import sys defaultDtypes = { "gfx942": {"fp8": torch.float8_e4m3fnuz}, @@ -18,7 +17,16 @@ def get_dtype_fp8(): - return defaultDtypes.get(get_gfx_runtime(), {"fp8": _8bit_fallback})["fp8"] + try: + if sys.platform == "win32": + gfx = torch.cuda.get_device_properties(0).gcnArchName.split(":", 1)[0] + else: + from ..jit.utils.chip_info import get_gfx_runtime + + gfx = get_gfx_runtime() + except Exception: + gfx = "unknown" + return defaultDtypes.get(gfx, {"fp8": _8bit_fallback})["fp8"] i4x2 = getattr(torch, "int4", _8bit_fallback) @@ -117,6 +125,8 @@ def str2tuple(v): def str2Dtype(v): + from ..ops.enum import QuantType + def _convert(s): if s.lower() == "none": return None @@ -142,4 +152,6 @@ def _convert(s): def str2ActivationType(s): """Convert string to ActivationType.""" + from ..ops.enum import ActivationType + return getattr(ActivationType, s.capitalize()) diff --git a/csrc/kernels/activation_kernels.cu b/csrc/kernels/activation_kernels.cu index 59df4e697be..a03cc9b7a13 100644 --- a/csrc/kernels/activation_kernels.cu +++ b/csrc/kernels/activation_kernels.cu @@ -1,6 +1,13 @@ // SPDX-License-Identifier: MIT // Copyright (C) 2025-2026, Advanced Micro Devices, Inc. All rights reserved. +#if defined(_WIN32) +// hipcc pre-includes , so include UCRT after enabling its M_* constants. +#ifndef _USE_MATH_DEFINES +#define _USE_MATH_DEFINES +#endif +#include +#endif #include #include "aiter_hip_common.h" diff --git a/op_tests/test_windows_rdna_ck_targets.py b/op_tests/test_windows_rdna_ck_targets.py new file mode 100644 index 00000000000..8dac3fd9228 --- /dev/null +++ b/op_tests/test_windows_rdna_ck_targets.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import os +import sys +import unittest +from unittest import mock + + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +JIT_UTILS = os.path.join(REPO_ROOT, "aiter", "jit", "utils") +sys.path.insert(0, JIT_UTILS) + +from build_targets import get_build_targets_env # noqa: E402 +from blob_gen import windows_blob_gen_argv # noqa: E402 +from mha_recipes import _ck_targets_flag_for_arch # noqa: E402 + + +WINDOWS_RDNA_TARGETS = { + "gfx1100": 96, + "gfx1101": 60, + "gfx1102": 32, + "gfx1103": 12, + "gfx1151": 20, + "gfx1201": 64, +} + + +class TestWindowsRDNACKTargets(unittest.TestCase): + def test_offline_build_target_defaults(self): + for gfx, cu_num in WINDOWS_RDNA_TARGETS.items(): + with self.subTest(gfx=gfx), mock.patch.dict( + os.environ, {"GPU_ARCHS": gfx}, clear=True + ): + self.assertEqual(get_build_targets_env(), [(gfx, cu_num)]) + + def test_ck_codegen_receives_each_rdna_target(self): + for gfx in WINDOWS_RDNA_TARGETS: + with self.subTest(gfx=gfx): + self.assertEqual( + _ck_targets_flag_for_arch(gfx), f" --targets {gfx}" + ) + + def test_cdna_keeps_ck_generator_defaults(self): + self.assertEqual(_ck_targets_flag_for_arch("gfx942"), "") + self.assertEqual(_ck_targets_flag_for_arch("gfx950"), "") + + def test_blob_generator_paths_with_spaces_are_single_arguments(self): + argv = windows_blob_gen_argv( + r"C:\Program Files\Python\python.exe", + ( + r"C:\Work Tree\aiter\3rdparty\composable_kernel" + r"\example\ck_tile\01_fmha\generate.py " + '-d fwd --receipt 100 --filter " @ " --output_dir {}' + ), + r"C:\Work Tree\build\blob", + ) + + self.assertEqual( + argv, + [ + r"C:\Program Files\Python\python.exe", + ( + r"C:\Work Tree\aiter\3rdparty\composable_kernel" + r"\example\ck_tile\01_fmha\generate.py" + ), + "-d", + "fwd", + "--receipt", + "100", + "--filter", + " @ ", + "--output_dir", + r"C:\Work Tree\build\blob", + ], + ) + + def test_blob_generator_equals_path_with_spaces(self): + argv = windows_blob_gen_argv( + r"C:\Python\python.exe", + ( + r"C:\Work Tree\gen.py --working_path {} " + r"--compiled_kids_sidecar=C:\Work Tree\compiled.json" + ), + r"C:\Work Tree\blob", + ) + + self.assertEqual( + argv, + [ + r"C:\Python\python.exe", + r"C:\Work Tree\gen.py", + "--working_path", + r"C:\Work Tree\blob", + "--compiled_kids_sidecar", + r"C:\Work Tree\compiled.json", + ], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 48659ad98b1..694cbc84402 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ requires = [ "psutil", "ninja", "pandas", - "flydsl==0.2.4" + "flydsl==0.2.4; sys_platform != 'win32'" ] [tool.setuptools_scm] diff --git a/setup.py b/setup.py index ec2d84330e2..27ac0499b59 100644 --- a/setup.py +++ b/setup.py @@ -20,8 +20,12 @@ PRETUNE_MODULES = os.environ.get("PRETUNE_MODULES", "") ENABLE_CK = int(os.environ.get("ENABLE_CK", "1")) IS_WINDOWS = sys.platform == "win32" -# Single skip-C++/HIP-build gate; Windows enables it automatically. -AITER_TRITON_ONLY = os.environ.get("AITER_TRITON_ONLY", "0") == "1" or IS_WINDOWS +# Native Windows ROCm builds are opt-in while the default wheel remains the +# lightweight Triton-only package published by upstream. +AITER_ENABLE_HIP = os.environ.get("AITER_ENABLE_HIP", "0") == "1" +AITER_TRITON_ONLY = os.environ.get("AITER_TRITON_ONLY", "0") == "1" or ( + IS_WINDOWS and not AITER_ENABLE_HIP +) if AITER_TRITON_ONLY: ENABLE_CK = False PREBUILD_KERNELS = False @@ -56,7 +60,7 @@ def is_develop_mode(): return False -if not AITER_TRITON_ONLY and is_develop_mode(): +if not AITER_TRITON_ONLY and not IS_WINDOWS and is_develop_mode(): try: from importlib.metadata import version as pkg_version from packaging.version import Version @@ -161,7 +165,10 @@ def prepare_packaging(): shutil.copytree("3rdparty", "aiter_meta/3rdparty") else: os.makedirs("aiter_meta/3rdparty", exist_ok=True) - if not AITER_TRITON_ONLY: + # The precompiled HSA code objects are Linux/CDNA-specific. Bundling them + # in a Windows wheel adds over 100 MiB but cannot provide gfx115x kernels; + # those architectures use the CK/JIT sources packaged above instead. + if not AITER_TRITON_ONLY and not IS_WINDOWS: shutil.copytree("hsa", "aiter_meta/hsa") else: os.makedirs("aiter_meta/hsa", exist_ok=True) @@ -389,18 +396,22 @@ def build_one_module(one_opt_args): os.environ["PREBUILD_THREAD_NUM"] = str(prebuid_thread_num) # --- FlyDSL AOT pre-compilation (MOE + GEMM, before CK) --- - _prev_aot_import = os.environ.get("AITER_AOT_IMPORT") - os.environ["AITER_AOT_IMPORT"] = "1" - try: - from aiter.aot.flydsl.common import run_aot + # FlyDSL does not publish Windows wheels; native Windows builds use CK. + if not IS_WINDOWS: + _prev_aot_import = os.environ.get("AITER_AOT_IMPORT") + os.environ["AITER_AOT_IMPORT"] = "1" + try: + from aiter.aot.flydsl.common import run_aot - flydsl_cache_dir = os.path.join(this_dir, "aiter", "jit", "flydsl_cache") - run_aot(flydsl_cache_dir) - finally: - if _prev_aot_import is None: - os.environ.pop("AITER_AOT_IMPORT", None) - else: - os.environ["AITER_AOT_IMPORT"] = _prev_aot_import + flydsl_cache_dir = os.path.join( + this_dir, "aiter", "jit", "flydsl_cache" + ) + run_aot(flydsl_cache_dir) + finally: + if _prev_aot_import is None: + os.environ.pop("AITER_AOT_IMPORT", None) + else: + os.environ["AITER_AOT_IMPORT"] = _prev_aot_import # --- CK kernel builds --- with ThreadPoolExecutor(max_workers=prebuid_thread_num) as executor: @@ -470,8 +481,9 @@ def has_ext_modules(self): "einops", "psutil", "packaging", - FLYDSL_VERSION, ] + if not IS_WINDOWS: + install_requires.append(FLYDSL_VERSION) setup( name=PACKAGE_NAME, @@ -480,10 +492,12 @@ def has_ext_modules(self): include_package_data=True, package_data={ "": ["*"], + "aiter": ["jit/*.pyd"], }, classifiers=[ "Programming Language :: Python :: 3", "License :: OSI Approved :: BSD License", + "Operating System :: Microsoft :: Windows", "Operating System :: Unix", ], cmdclass={"build_ext": NinjaBuildExtension},