[FlyDSL] [Feature] Precompile only the GEMM kernels the wheel is built for

.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

@@ -296,6 +296,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/aot/flydsl/README.md

@@ -98,14 +98,18 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 | `AITER_FLYDSL_AOT_TIMEOUT` | Per-kernel wall-clock cap (seconds). A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` |
 | `AITER_FLYDSL_AOT_MAX_RETRIES` | Retries for a worker that **died abnormally** (OOM-kill / segfault / timeout-kill). A clean compile error is never retried. `0` disables. | `2` |
 | `AITER_CONFIGS` | Resolves the default CSV lookup path (same as the runtime JIT) | repo built-in |
-| `ARCH` / `GPU_ARCHS` | Selects which jobs to build, not what arch a job compiles *for* (that comes from the CSV's `cu_num`). `conv.py` applies it inside `parse_csv`, so both `python -m` and the `setup.py` path (`run_aot`) honour it. `gemm.py` still filters in `main()` only, so `run_aot` builds all of its archs. | auto-detect |
-
-> **About the compile target arch.** The arch each kernel is actually compiled
-> for is derived per-job from the CSV's `cu_num` column (`cu_num_to_arch(...)`)
-> and applied internally via `FLYDSL_GPU_ARCH`. That internal var is overwritten
-> for every job, so setting `ARCH` / `GPU_ARCHS` / `FLYDSL_GPU_ARCH` in your shell
-> does **not** change what gets built. To cross-compile, edit the `cu_num`
-> column in the CSV.
+| `AITER_GPU_TARGETS` / `ARCH` / `GPU_ARCHS` | Filter **GEMM** AOT jobs by target. The first one set wins, in that order. Unset builds every job. | unset |
+
+> **Compile target.** Each job is compiled for the arch in its `gfx` field, or
+> from the CSV `cu_num` via `cu_num_to_arch(...)` when `gfx` is empty. AOT
+> applies that through `FLYDSL_GPU_ARCH` and overwrites it per job, so setting
+> `FLYDSL_GPU_ARCH` in the shell does not change what gets built.
+
+`AITER_GPU_TARGETS=gfx950:128` is the exact `(gfx, cu_num)` form on both
+paths. The other two differ by path. Running `python -m aiter.aot.flydsl.gemm`
+directly, `ARCH` selects every GEMM job of that arch and `GPU_ARCHS` does the
+same (`CU_NUM` is ignored); packaging (`setup.py`) reads no `ARCH` and pairs
+`GPU_ARCHS` with `CU_NUM` into exact SKU pairs instead.
 
 Example:
 

aiter/aot/flydsl/common.py

@@ -133,6 +133,18 @@ def override_env(var_name: str, value: str | None) -> Iterator[None]:
             os.environ[var_name] = prev
 
 
+def _filter_collected_aot_jobs(
+    kind: OpKind, jobs: list[dict[str, Any]]
+) -> list[dict[str, Any]]:
+    # GEMM only: parse_csv for the others keep just one of gfx/cu_num,
+    # so a (gfx, cu_num) filter would drop all of them.
+    if kind is OpKind.GEMM:
+        from .gemm import filter_jobs_for_build_targets
+
+        return filter_jobs_for_build_targets(jobs)
+    return jobs
+
+
 def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]:
     """Load DEFAULT_CSVS + parse_csv for the named kind and return its
     job list. Note: importing .gemm / .moe / .chunk_gdn_h here also
@@ -161,7 +173,8 @@ def _collect_aot_jobs_for(kind: OpKind) -> list[dict[str, Any]]:
         from .chunk_gdn_h import DEFAULT_CSVS, parse_csv
     else:
         raise ValueError(f"unknown FlyDSL AOT kind: {kind!r}")
-    return collect_aot_jobs(DEFAULT_CSVS, parse_csv)
+    jobs = collect_aot_jobs(DEFAULT_CSVS, parse_csv)
+    return _filter_collected_aot_jobs(kind, jobs)
 
 
 def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:

aiter/aot/flydsl/gemm.py

@@ -29,7 +29,12 @@
 
 Environment variables:
     FLYDSL_RUNTIME_CACHE_DIR  Cache directory (default: ~/.flydsl/cache)
-    GPU_ARCHS / ARCH          Target GPU architecture information for logging.
+    AITER_GPU_TARGETS         ';'/','-separated 'gfx' or 'gfx:cu_num' build
+                              targets; missing pairs retain runtime fallback.
+    GPU_ARCHS                 Arch list. Packaging resolves it with CU_NUM into
+                              exact targets; the CLI keeps every job of that arch.
+    CU_NUM                    Compute-unit count paired with GPU_ARCHS.
+    ARCH                      Bare-arch filter on the CLI; packaging ignores it.
 """
 
 from __future__ import annotations
@@ -55,6 +60,11 @@
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
+from aiter.jit.utils.build_targets import (
+    KNOWN_GFX,
+    _parse_gpu_targets_env,
+)
+from aiter.jit.utils.chip_info import get_build_targets, get_gfx_runtime
 from aiter.ops.flydsl.batched_gemm_a8w8_gfx950 import (
     BMM_MFMA_NAME_PREFIX,
     compile_bmm_a8w8_mxfp8_gfx950,
@@ -243,7 +253,7 @@ def parse_csv(csv_path: str):
             m = int(row["M"])
             n = int(row["N"])
             k = int(row["K"])
-            cu_num = int(row.get("cu_num", "0"))
+            cu_num = int(row.get("cu_num") or 0)
             gfx = row.get("gfx", "").strip()
 
             if row.get("kernelId", "").strip() == "bmm":  # the batched GEMM, B = 1
@@ -892,6 +902,106 @@ def job_arch(cu_num: int = 0, gfx: str = "") -> str:
     return gfx or cu_num_to_arch(cu_num, default=GEMM_AOT_ARCH_DEFAULT)
 
 
+def active_target_env(*, cli: bool = True) -> tuple[str, str] | None:
+    """First configured target variable; packaging ignores the CLI-only ARCH."""
+    if cli:
+        names = ("AITER_GPU_TARGETS", "ARCH", "GPU_ARCHS")
+    else:
+        names = ("AITER_GPU_TARGETS", "GPU_ARCHS")
+    for var_name in names:
+        value = (os.environ.get(var_name) or "").strip()
+        if value:
+            return var_name, value
+    return None
+
+
+def _archs_from_env(value: str) -> set[str]:
+    """Arch names in an ARCH/GPU_ARCHS value. Raises ValueError on a bad one."""
+    archs = {a.strip().lower() for a in re.split(r"[;,]", value) if a.strip()}
+    if not archs:
+        raise ValueError("contains no valid architecture names")
+    if "native" in archs and len(archs) != 1:
+        raise ValueError("'native' cannot be combined with other targets")
+    unknown = archs - KNOWN_GFX - {"native"}
+    if unknown:
+        raise ValueError(f"contains unknown target(s): {sorted(unknown)}")
+    return archs
+
+
+def _requested_targets(
+    var_name: str, value: str, *, arch_wide: bool
+) -> set[tuple[str, int | None]]:
+    """Resolve (gfx, CU count) targets; None means every CU count of an arch."""
+    if var_name == "AITER_GPU_TARGETS":
+        return set(_parse_gpu_targets_env())
+
+    archs = _archs_from_env(value)
+    if arch_wide:
+        if archs == {"native"}:
+            archs = {get_gfx_runtime()}
+        return {(gfx, None) for gfx in archs}
+
+    # Normalize separators/case before resolving packaging targets. The shared
+    # resolver handles CU_NUM, native detection, and live-device CU counts.
+    with override_env("GPU_ARCHS", ";".join(sorted(archs))):
+        return set(get_build_targets())
+
+
+def filter_jobs_for_build_targets(
+    jobs: list[dict], *, arch_wide: bool = False
+) -> list[dict]:
+    """Select jobs for AITER_GPU_TARGETS, or for GPU_ARCHS with CU_NUM.
+
+    The CLI passes arch_wide=True: ARCH takes priority over GPU_ARCHS and both
+    select every CU count of the requested architectures. AITER_GPU_TARGETS
+    always selects exact targets and takes priority over either variable.
+    """
+    active = active_target_env(cli=arch_wide)
+    if active is None:
+        return jobs
+    var_name, value = active
+    try:
+        targets = _requested_targets(var_name, value, arch_wide=arch_wide)
+    except ValueError as e:
+        raise RuntimeError(f"{var_name} {e}.") from None
+
+    selected = []
+    covered = set()
+    legacy_cu_nums = set()
+    for job in jobs:
+        gfx = str(job.get("gfx") or "").strip().lower()
+        cu_num = int(job.get("cu_num") or 0)
+        if not gfx:
+            gfx = job_arch(cu_num)
+            legacy_cu_nums.add(cu_num)
+
+        # A job with no recorded CU count (0) matches every target of its arch.
+        matches = {
+            (target_gfx, target_cu)
+            for target_gfx, target_cu in targets
+            if gfx == target_gfx and (target_cu is None or cu_num in (0, target_cu))
+        }
+        if matches:
+            selected.append(job)
+            covered.update(matches)
+
+    for cu_num in sorted(legacy_cu_nums):
+        print(
+            f"  [WARN] FlyDSL GEMM row has no gfx; inferring {job_arch(cu_num)} "
+            f"from legacy cu_num={cu_num}. Re-run the tuner to record gfx."
+        )
+    missing = [
+        gfx if cu is None else f"{gfx}:{cu}" for gfx, cu in sorted(targets - covered)
+    ]
+    if missing:
+        print(
+            f"  [WARN] The configured GEMM CSVs have no FlyDSL GEMM jobs for "
+            f"{', '.join(missing)}; those targets will use their existing "
+            "runtime fallback."
+        )
+    return selected
+
+
 def compile_one_config(
     kernel_name: str,
     kind: str,
@@ -982,17 +1092,17 @@ def main():
     cache_dir = os.path.expanduser(
         os.environ.get("FLYDSL_RUNTIME_CACHE_DIR", "~/.flydsl/cache")
     )
-    arch = os.environ.get("ARCH") or os.environ.get("GPU_ARCHS")
-
     all_jobs = collect_aot_jobs(csv_paths, parse_csv)
-    if arch:
-        # GPU_ARCHS may be a ';'- or ','-separated list (e.g. "gfx942;gfx950").
-        arch_set = {a.strip() for a in re.split(r"[;,]", arch) if a.strip()}
-        n_before = len(all_jobs)
-        all_jobs = [
-            j for j in all_jobs if job_arch(j["cu_num"], j.get("gfx", "")) in arch_set
-        ]
-        print(f"[aiter] ARCH={arch}: {len(all_jobs)}/{n_before} jobs match")
+    n_before = len(all_jobs)
+    try:
+        all_jobs = filter_jobs_for_build_targets(all_jobs, arch_wide=True)
+    except RuntimeError as e:
+        print(f"Error: {e}")
+        sys.exit(1)
+    active = active_target_env()
+    target_spec = active[1] if active else ""
+    if target_spec:
+        print(f"[aiter] targets={target_spec}: {len(all_jobs)}/{n_before} jobs match")
 
     print("=" * 72)
     print("FlyDSL GEMM AOT Pre-compilation")
@@ -1003,7 +1113,7 @@ def main():
         print(f"  {kind} jobs:{' ' * max(1, 15 - len(kind))}{count}")
     print(f"  Total jobs:       {len(all_jobs)}")
     print(f"  Cache dir:        {cache_dir}")
-    print(f"  Target arch:      {arch or '(all archs found in CSVs)'}")
+    print(f"  Target arch:      {target_spec or '(all archs found in CSVs)'}")
     print("=" * 72)
 
     total_t0 = time.time()

aiter/jit/core.py

@@ -22,6 +22,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 get_gfx, get_gfx_list, get_gfx_runtime
 from cpp_extension import _jit_compile, executable_path, get_hip_version
 from file_baton import FileBaton
@@ -671,29 +677,21 @@ 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]
-    # List of allowed architectures
-    allowed_archs = [
-        "native",
-        "gfx90a",
-        "gfx940",
-        "gfx941",
-        "gfx942",
-        "gfx1100",
-        "gfx1101",
-        "gfx1102",
-        "gfx1103",
-        "gfx1150",
-        "gfx1151",
-        "gfx1152",
-        "gfx1153",
-        "gfx1200",
-        "gfx1201",
-        "gfx1250",
-        "gfx950",
-        "gfx1250",
-    ]
+    # 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"))
+
+    allowed_archs = {"native", *KNOWN_GFX}
 
     # Validate if each element in archs is in allowed_archs
     assert all(
@@ -1202,13 +1200,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

@@ -350,8 +350,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": {