[HIP] [CI] [JIT] [Feature] Add gfx:cu_num build targets so one build serves multiple SKUs of an arch

.github/workflows/aiter-release.yaml

@@ -82,6 +82,11 @@ on:
         type: string
         required: false
         default: 'gfx942;gfx950'
+      gpu_targets:
+        description: 'Explicit gfx:cu_num build targets; overrides gpu_archs when set'
+        type: string
+        required: false
+        default: ''
       runner:
         description: 'Build host'
         type: string
@@ -162,6 +167,9 @@ jobs:
       BUILD_DOCKER_IMAGE: ${{ matrix.docker_image }}
       PYTHON_VERSION: ${{ matrix.python_version }}
       GPU_ARCHS: ${{ inputs.gpu_archs || github.event.inputs.gpu_archs }}
+      # workflow_call only, since workflow_dispatch is at GitHub's 10-input cap;
+      # blank resolves to unset, leaving GPU_ARCHS in charge of target selection.
+      AITER_GPU_TARGETS: ${{ inputs.gpu_targets || '' }}
       RELEASE_TYPE: ${{ inputs.release_type || github.event.inputs.release_type || 'none' }}
       ADD_DATE_STAMP: ${{ inputs.add_date_stamp || github.event.inputs.add_date_stamp }}
       USE_MANYLINUX: ${{ inputs.use_manylinux || github.event.inputs.use_manylinux || (startsWith(matrix.docker_image, 'pytorch/manylinux') && 'true') || 'false' }}
@@ -305,13 +313,15 @@ jobs:
             -e SETUPTOOLS_SCM_PRETEND_VERSION="${{ steps.version.outputs.nightly_version }}" \
             -e PREBUILD_KERNELS=1 \
             -e GPU_ARCHS="${{ env.GPU_ARCHS }}" \
+            -e AITER_GPU_TARGETS="${{ env.AITER_GPU_TARGETS }}" \
             -w /workspace \
             aiter_build_${{ matrix.python_version }} \
             bash -c "python3 setup.py bdist_wheel && ls dist/*.whl"
           else
             docker exec \
             -e PREBUILD_KERNELS=1 \
             -e GPU_ARCHS="${{ env.GPU_ARCHS }}" \
+            -e AITER_GPU_TARGETS="${{ env.AITER_GPU_TARGETS }}" \
             -w /workspace \
             aiter_build_${{ matrix.python_version }} \
             bash -c "python3 setup.py bdist_wheel && ls dist/*.whl"
@@ -412,6 +422,7 @@ jobs:
             -e SETUPTOOLS_SCM_PRETEND_VERSION="${PRETEND_VER}" \
             -e PREBUILD_KERNELS=1 \
             -e GPU_ARCHS="${GPU_ARCHS}" \
+            -e AITER_GPU_TARGETS="${AITER_GPU_TARGETS}" \
             -e PYBIN="${PYBIN}" \
             -w /workspace \
             aiter_build_${{ matrix.python_version }} \

.github/workflows/nightly.yaml

@@ -18,6 +18,10 @@ on:
         description: 'GPU architectures (default: gfx942;gfx950)'
         type: string
         required: false
+      gpu_targets:
+        description: 'Explicit gfx:cu_num build targets; overrides gpu_archs when set (default: gfx942:304;gfx950:256, dropped if gpu_archs is overridden)'
+        type: string
+        required: false
       skip_tests:
         description: 'Skip test job for faster iteration'
         type: boolean
@@ -41,6 +45,7 @@ jobs:
       build_py310: ${{ github.event.inputs.python_version == '' || github.event.inputs.python_version == '3.10' }}
       build_py312: ${{ github.event.inputs.python_version == '' || github.event.inputs.python_version == '3.12' }}
       gpu_archs: ${{ github.event.inputs.gpu_archs || 'gfx942;gfx950' }}
+      gpu_targets: ${{ github.event.inputs.gpu_targets || (github.event.inputs.gpu_archs == '' && 'gfx942:304;gfx950:256') || '' }}
       runner: build-only-aiter
       add_date_stamp: true
     permissions:

.github/workflows/release-automation.yaml

@@ -255,6 +255,7 @@ jobs:
       build_py310: true
       build_py312: true
       gpu_archs: 'gfx942;gfx950'
+      gpu_targets: 'gfx942:304;gfx950:256'
       runner: build-only-aiter
       add_date_stamp: false
       use_manylinux: true

aiter/jit/core.py

@@ -21,6 +21,12 @@
 
 this_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, f"{this_dir}/utils/")
+from build_targets import (
+    KNOWN_GFX,
+    _parse_gpu_archs_env,
+    get_build_archs_env,
+    gpu_archs_env_names,
+)
 from chip_info import _detect_native, get_gfx, get_gfx_list, get_gfx_runtime
 from cpp_extension import _jit_compile, executable_path, get_hip_version
 from file_baton import FileBaton
@@ -690,35 +696,26 @@ def get_user_jit_dir() -> str:
 
 
 def validate_and_update_archs():
-    archs = os.getenv("GPU_ARCHS", "native").split(";")
-    archs = [arch.strip() for arch in archs]
+    # AITER_GPU_TARGETS is authoritative for the arch set; a conflicting
+    # GPU_ARCHS warns and loses.
+    named = get_build_archs_env()
+    if named is not None:
+        archs = list(named)
+        explicit = set(gpu_archs_env_names())
+        if explicit and explicit != set(archs):
+            logger.warning(
+                f"GPU_ARCHS={sorted(explicit)} disagrees with "
+                f"AITER_GPU_TARGETS={archs}; compiling for {archs}."
+            )
+    else:
+        archs = _parse_gpu_archs_env(os.getenv("GPU_ARCHS", "native"))
 
     # clang's `--offload-arch=native` shells out to offload-arch, which cannot
     # detect the GPU on Windows, so resolve the arch ourselves instead.
     if IS_WINDOWS and archs == ["native"]:
         archs = _detect_native()
 
-    # List of allowed architectures
-    allowed_archs = [
-        "native",
-        "gfx90a",
-        "gfx940",
-        "gfx941",
-        "gfx942",
-        "gfx1100",
-        "gfx1101",
-        "gfx1102",
-        "gfx1103",
-        "gfx1150",
-        "gfx1151",
-        "gfx1152",
-        "gfx1153",
-        "gfx1200",
-        "gfx1201",
-        "gfx1250",
-        "gfx950",
-        "gfx1250",
-    ]
+    allowed_archs = {"native", *KNOWN_GFX}
 
     # Validate if each element in archs is in allowed_archs
     assert all(
@@ -1273,13 +1270,11 @@ def raise_build_error(error):
         # Cluster launch is a HOST-side API question (hipDrvLaunchKernelEx +
         # HIP_LAUNCH_CONFIG appear in ROCm 7.0), not a question about which GPU
         # this machine has -- so the arch test must accept a cross-compile for
-        # gfx1250 the way the gfx1250 flags in optCompilerConfig.json already do.
-        # get_gfx() alone reads the LAST entry of a multi-arch GPU_ARCHS, which
-        # left "gfx1250;gfx942" building gfx1250 kernels whose cluster launch
+        # gfx1250, and must test membership rather than a single resolved arch:
+        # "gfx1250;gfx942" otherwise builds gfx1250 kernels whose cluster launch
         # path was compiled out (AiterAsmKernelFast then rejects them at launch).
-        if (
-            get_gfx() == "gfx1250" or "gfx1250" in os.environ.get("GPU_ARCHS", "")
-        ) and hip_version >= Version("7.0.0"):
+        # Same test as the gfx1250 flags in optCompilerConfig.json.
+        if "gfx1250" in get_gfx_list() and hip_version >= Version("7.0.0"):
             flags_hip += ["-DAITER_ENABLE_CLUSTER_LAUNCH"]
 
         if not torch_exclude:

aiter/jit/optCompilerConfig.json

@@ -353,8 +353,8 @@
         "flags_extra_cc": [],
         "flags_extra_hip": [
             "'-mllvm --amdgpu-mfma-vgpr-form'",
-            "('-mllvm -amdgpu-expert-scheduling-mode' if (get_gfx() == 'gfx1250' or 'gfx1250' in os.environ.get('GPU_ARCHS', '')) else '')",
-            "('-mllvm -enable-post-misched=1' if (get_gfx() == 'gfx1250' or 'gfx1250' in os.environ.get('GPU_ARCHS', '')) else '')",
+            "('-mllvm -amdgpu-expert-scheduling-mode' if 'gfx1250' in get_gfx_list() else '')",
+            "('-mllvm -enable-post-misched=1' if 'gfx1250' in get_gfx_list() else '')",
             "f'-DOPUS_GEN_CO_DIR=\\\\\"{AITER_CSRC_DIR}/opus_gemm/gen_co\\\\\"'"
         ],
         "flags_extra_hip_per_source": {

aiter/jit/utils/build_targets.py

@@ -5,6 +5,7 @@
 # No torch dependency — safe to import in build scripts, gen_instances, and tests
 # that run without a GPU or a full PyTorch install.
 import os
+import re
 
 GFX_MAP = {
     0: "native",
@@ -42,6 +43,9 @@
     "gfx1250": 256,  # Gfx1250
 }
 
+# Valid arch names; those without a GFX_CU_NUM_MAP default must name their CU count.
+KNOWN_GFX = frozenset(v for v in GFX_MAP.values() if v != "native")
+
 
 def _parse_gpu_archs_env(gfx_env: str) -> list[str]:
     """Split a GPU_ARCHS string into a list of non-empty architecture names.
@@ -50,7 +54,7 @@ def _parse_gpu_archs_env(gfx_env: str) -> list[str]:
     on ';' and stripping whitespace — e.g. GPU_ARCHS=" ; " would otherwise
     silently produce an empty target list and fall back to heuristic kernels.
     """
-    archs = [g.strip() for g in gfx_env.split(";") if g.strip()]
+    archs = [g.strip().lower() for g in gfx_env.split(";") if g.strip()]
     if not archs:
         raise RuntimeError(
             f"GPU_ARCHS={gfx_env!r} contains no valid architecture names after splitting on ';'. "
@@ -59,19 +63,111 @@ def _parse_gpu_archs_env(gfx_env: str) -> list[str]:
     return archs
 
 
+def _cu_num_or_none(value: str) -> int | None:
+    """The positive CU count in value, or None when it is not one."""
+    try:
+        cu_num = int(value)
+    except (TypeError, ValueError):
+        return None
+    return cu_num if cu_num > 0 else None
+
+
+def _parse_gpu_targets_env() -> list[tuple[str, int]] | None:
+    """Parse AITER_GPU_TARGETS into (gfx, cu_num) targets, or None if it is unset.
+
+    gfx950:128;gfx950:256  -> [("gfx950", 128), ("gfx950", 256)]
+    gfx950                 -> [("gfx950", 256)]  # CU from GFX_CU_NUM_MAP
+    gfx942;gfx950:128      -> [("gfx942", 304), ("gfx950", 128)]
+    (unset or blank)       -> None
+
+    Entries split on ';' or ','; arch names are case-folded and validated,
+    CU counts must be positive integers.
+    """
+    targets_env = os.getenv("AITER_GPU_TARGETS")
+    if not targets_env or not targets_env.strip():
+        return None
+
+    targets = []
+    for entry in re.split(r"[;,]", targets_env):
+        entry = entry.strip()
+        if not entry:
+            continue
+        where = f"AITER_GPU_TARGETS entry {entry!r}"
+        gfx, sep, cu = entry.partition(":")
+        gfx = gfx.strip().lower()
+        if gfx not in KNOWN_GFX:
+            raise RuntimeError(
+                f"{where}: unknown gfx {gfx!r}. Known targets: {sorted(KNOWN_GFX)}"
+            )
+        if sep:
+            cu_num = _cu_num_or_none(cu.strip())
+            if cu_num is None:
+                raise RuntimeError(
+                    f"{where}: CU count {cu.strip()!r} must be a positive integer."
+                )
+            targets.append((gfx, cu_num))
+        elif gfx in GFX_CU_NUM_MAP:
+            targets.append((gfx, GFX_CU_NUM_MAP[gfx]))
+        else:
+            raise RuntimeError(
+                f"{where}: {gfx!r} has no default CU count — add it to "
+                f"GFX_CU_NUM_MAP in build_targets.py, or name the count "
+                f"explicitly as '{gfx}:<cu_num>'."
+            )
+
+    if not targets:
+        raise RuntimeError(
+            f"AITER_GPU_TARGETS={targets_env!r} names no targets. "
+            f"Expected entries of the form "
+            f"'gfx' or 'gfx:cu_num'."
+        )
+
+    # Preserve caller order, drop exact duplicates.
+    return list(dict.fromkeys(targets))
+
+
+def get_build_archs_env() -> list[str] | None:
+    """Deduped arch names from AITER_GPU_TARGETS, or None if it is unset."""
+    targets = _parse_gpu_targets_env()
+    if targets is None:
+        return None
+    return list(dict.fromkeys(gfx for gfx, _ in targets))
+
+
+def gpu_archs_env_names() -> list[str]:
+    """Arch names GPU_ARCHS explicitly requests, empty when unset or native."""
+    archs = _parse_gpu_archs_env(os.getenv("GPU_ARCHS") or "native")
+    return [a for a in archs if a != "native"]
+
+
 def get_build_targets_env() -> list[tuple[str, int]]:
-    """Resolve build targets from GPU_ARCHS env var only.  No live GPU detection.
+    """Resolve build targets from env only. No live GPU detection.
 
-    Raises RuntimeError if GPU_ARCHS is not set or contains an unknown arch.
-    Intended for CI nodes, build scripts, and tests that run without a GPU.
-    Use chip_info.get_build_targets() when live GPU fallback is also desired.
+    AITER_GPU_TARGETS, when set, overrides GPU_ARCHS + CU_NUM. Raises
+    RuntimeError if neither is set or an arch is unknown. Use
+    chip_info.get_build_targets() when live-GPU fallback is also desired.
+
+    AITER_GPU_TARGETS=gfx950:128;gfx950:256 -> [("gfx950", 128), ("gfx950", 256)]
+    GPU_ARCHS=gfx942;gfx950                 -> [("gfx942", 304), ("gfx950", 256)]
+    GPU_ARCHS=gfx942 CU_NUM=80              -> [("gfx942", 80)]
     """
+    targets = _parse_gpu_targets_env()
+    if targets is not None:
+        return targets
+
     gfx_env = os.getenv("GPU_ARCHS")
     if not gfx_env:
         raise RuntimeError(
-            "GPU_ARCHS is not set. "
+            "Neither AITER_GPU_TARGETS nor GPU_ARCHS is set. "
             "Set GPU_ARCHS=gfx942 (or similar) to resolve build targets without a GPU."
         )
+    cu_env = os.getenv("CU_NUM")
+    cu_override = None
+    if cu_env:
+        cu_override = _cu_num_or_none(cu_env)
+        if cu_override is None:
+            raise RuntimeError(f"CU_NUM={cu_env!r} must be a positive integer.")
+
     targets = []
     for gfx in _parse_gpu_archs_env(gfx_env):
         if gfx not in GFX_CU_NUM_MAP:
@@ -80,9 +176,12 @@ def get_build_targets_env() -> list[tuple[str, int]]:
                 f"GFX_CU_NUM_MAP in build_targets.py. Known targets: "
                 f"{list(GFX_CU_NUM_MAP.keys())}"
             )
-        cu_num = int(os.getenv("CU_NUM", GFX_CU_NUM_MAP[gfx]))
-        targets.append((gfx, cu_num))
-    return targets
+        targets.append((gfx, cu_override or GFX_CU_NUM_MAP[gfx]))
+    return list(dict.fromkeys(targets))
+
+
+def _target_mask(tune_df, gfx: str, cu_num: int):
+    return (tune_df["gfx"] == gfx) & (tune_df["cu_num"] == cu_num)
 
 
 def filter_tune_df(tune_df, targets: list):
@@ -101,5 +200,14 @@ def filter_tune_df(tune_df, targets: list):
 
     mask = pd.Series([False] * len(tune_df), index=tune_df.index)
     for gfx, cu_num in targets:
-        mask |= (tune_df["gfx"] == gfx) & (tune_df["cu_num"] == cu_num)
+        mask |= _target_mask(tune_df, gfx, cu_num)
     return tune_df[mask]
+
+
+def unmatched_targets(tune_df, targets: list) -> list[str]:
+    """The targets with no row in tune_df, as 'gfx:cu_num' strings."""
+    return [
+        f"{gfx}:{cu_num}"
+        for gfx, cu_num in targets
+        if not _target_mask(tune_df, gfx, cu_num).any()
+    ]