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

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

@@ -31,7 +31,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
@@ -57,6 +62,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,
@@ -213,7 +223,7 @@ def _bmm_mfma_job(
         "m": 1,
         "n": n,
         "k": k,
-        "cu_num": 0,  # the gfx column names the arch
+        "cu_num": int(row.get("cu_num") or 0),
         "gfx": row.get("gfx", "").strip(),
         "b": b,
         # A table from before the column holds only 128x128 rows (as at runtime).
@@ -253,7 +263,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
@@ -994,6 +1004,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,
@@ -1086,17 +1196,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")
@@ -1107,7 +1217,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()