Add native Windows RDNA HIP and CK support

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

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

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 <tuple>
-__host__ __device__ void func(){{std::tuple<int, int> 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 <tuple>
+__host__ __device__ void func(){std::tuple<int, int> 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

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

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
 }