[CK] [FlyDSL] [CI] [Bugfix] Explicit gfx in shipped fused-MoE tuning CSVs

aiter/aot/flydsl/README.md

@@ -71,11 +71,12 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 | `ARCH` / `GPU_ARCHS` | **Banner/logging only** — printed as the "Target arch" line. Does **not** control the compiled target. | 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.
+> for comes from `resolve_job_arch(...)`: an explicit CSV `gfx` if present,
+> otherwise only the known historical `cu_num` mapping (80/304 → gfx942,
+> 256 → gfx950). Missing or unknown values fail rather than silently compiling
+> for gfx950. The chosen arch is applied internally via `FLYDSL_GPU_ARCH`.
+> Setting `ARCH` / `GPU_ARCHS` / `FLYDSL_GPU_ARCH` in your shell does **not**
+> change what gets built. To cross-compile, set the `gfx` column in the CSV.
 
 Example:
 

aiter/aot/flydsl/chunk_gdn_h.py

@@ -36,16 +36,15 @@
 from aiter.aot.flydsl.common import (
     collect_aot_jobs,
     compile_only_env,
-    cu_num_to_arch,
     job_identity,
     override_env,
+    resolve_job_arch,
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
 from aiter.ops.flydsl.kernels.gdr_prefill import compile_chunk_gated_delta_h
 from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled
 
-CHUNK_GDN_H_AOT_ARCH_DEFAULT = "gfx950"
 _KERNEL_NAME = "chunk_gdn_fwd_h_flydsl_opt"
 
 _TORCH_DTYPE = {
@@ -123,6 +122,7 @@ def parse_csv(csv_path: str) -> list[dict[str, Any]]:
                 H = int(row["H"])
                 Hg = int(row["Hg"])
                 cu_num = int(row.get("cu_num") or 0)
+                gfx = (row.get("gfx") or "").strip()
                 is_varlen = _parse_bool(row.get("is_varlen") or "True")
                 use_h0 = _parse_bool(row.get("use_h0") or "True")
                 store_fs = _parse_bool(row.get("store_fs") or "True")
@@ -148,6 +148,7 @@ def parse_csv(csv_path: str) -> list[dict[str, Any]]:
                     "kernel_name": _KERNEL_NAME,
                     "dtype": dtype,
                     "cu_num": cu_num,
+                    "gfx": gfx,
                     "K": K,
                     "V": V,
                     "BT": BT,
@@ -302,9 +303,9 @@ def _format_shape_str(job: dict) -> str:
     )
 
 
-def compile_one_config(*, cu_num: int = 0, **kwargs) -> dict:
+def compile_one_config(*, cu_num: int = 0, gfx: str = "", **kwargs) -> dict:
     """Compile one opt configuration and save it to cache."""
-    aot_arch = cu_num_to_arch(cu_num, default=CHUNK_GDN_H_AOT_ARCH_DEFAULT)
+    aot_arch = resolve_job_arch(cu_num, gfx)
     kwargs.pop("kernel_name", None)
     shape_str = _format_shape_str(kwargs)
     result = {
@@ -364,7 +365,7 @@ def main():
     for csv_path in csv_paths:
         print(f"  CSV:          {csv_path}")
     print(f"  Total jobs:   {len(jobs)}")
-    print("  Compile arch: (from cu_num)")
+    print("  Compile arch: (from gfx, else known cu_num mapping)")
     print(f"  Cache dir:    {cache_dir}")
     print(f"  Target arch:  {arch}")
     print("=" * 72)

aiter/aot/flydsl/common.py

@@ -18,6 +18,8 @@
 from multiprocessing.connection import wait as wait_for_sentinels
 from typing import Any
 
+from aiter.jit.utils.gfx_placeholders import GFX_PLACEHOLDERS, LEGACY_CU_NUM_TO_GFX
+
 _DEFAULT_KERNEL_TIMEOUT = 1200.0
 _DEFAULT_MAX_WORKERS = 64
 _DEFAULT_MAX_RETRIES = 2
@@ -47,16 +49,33 @@ def __str__(self) -> str:
         return f"{self.kind.name} {self.kernel_name}"
 
 
-_CU_NUM_TO_ARCH = {
-    80: "gfx942",
-    304: "gfx942",
-    256: "gfx950",
-}
+def cu_num_to_arch(cu_num: int) -> str:
+    """Map a known legacy compute-unit count to its historical architecture."""
+    try:
+        mapped = LEGACY_CU_NUM_TO_GFX.get(int(cu_num))
+    except (TypeError, ValueError):
+        mapped = None
+    if mapped is None:
+        raise ValueError(
+            f"cannot map cu_num={cu_num!r} to an architecture; known values are "
+            f"{sorted(LEGACY_CU_NUM_TO_GFX)}"
+        )
+    return mapped
+
 
+def resolve_job_arch(cu_num: int = 0, gfx: str = "") -> str:
+    """Architecture this AOT job should compile for.
 
-def cu_num_to_arch(cu_num: int, default: str = "gfx950") -> str:
-    """Map compute-unit count to GPU architecture string."""
-    return _CU_NUM_TO_ARCH.get(cu_num, default)
+    Prefer an explicit ``gfx`` from the row. Placeholder values (``0``, empty,
+    ``nan``, ``None``) are treated as missing. If gfx is missing, fall back
+    only to the known historical CU-count mapping (80/304 -> gfx942,
+    256 -> gfx950). Unknown or missing values raise rather than inventing
+    an architecture.
+    """
+    gfx = "" if gfx is None else str(gfx).strip()
+    if gfx and gfx not in GFX_PLACEHOLDERS:
+        return gfx
+    return cu_num_to_arch(cu_num)
 
 
 def job_identity(job: dict[str, Any]) -> tuple:

aiter/aot/flydsl/fhmoe.py

@@ -9,7 +9,7 @@
 
 from dataclasses import dataclass
 
-from aiter.aot.flydsl.common import cu_num_to_arch
+from aiter.aot.flydsl.common import resolve_job_arch
 
 
 def _shared_weight(device, n_in: int, k_in: int):
@@ -149,6 +149,7 @@ def precompile_fhmoe_to_cache(
     b_dtype: str = "fp4",
     act: str = "silu",
     cu_num: int = 0,
+    gfx: str = "",
     enable_bias: bool = False,
     **kwargs,
 ):
@@ -167,8 +168,11 @@ def precompile_fhmoe_to_cache(
         raise ValueError("FHMoE AOT does not support expert bias")
     if act != "silu":
         raise ValueError(f"FHMoE AOT supports only SiLU, got {act=}")
-    if cu_num_to_arch(cu_num) != "gfx950":
-        raise ValueError(f"FHMoE AOT supports only gfx950, got {cu_num=}")
+    aot_arch = resolve_job_arch(cu_num, gfx)
+    if aot_arch != "gfx950":
+        raise ValueError(
+            f"FHMoE AOT supports only gfx950, got gfx={gfx!r} cu_num={cu_num}"
+        )
 
     from aiter.aot.flydsl.moe import _precompile_to_cache
 

aiter/aot/flydsl/gemm.py

@@ -46,9 +46,9 @@
 from aiter.aot.flydsl.common import (
     collect_aot_jobs,
     compile_only_env,
-    cu_num_to_arch,
     job_identity,
     override_env,
+    resolve_job_arch,
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
@@ -100,7 +100,6 @@
     AITER_CONFIGS.AITER_CONFIG_BF16_BATCHED_GEMM_FILE,
     AITER_CONFIGS.AITER_CONFIG_GEMM_BF16_FILE,
 ]
-GEMM_AOT_ARCH_DEFAULT = "gfx950"
 
 _PRESHUFFLE_RE = re.compile(
     r"^flydsl_bpreshuflle_"
@@ -649,11 +648,6 @@ def _compile_ptpc_wmma_to_cache(
             )
 
 
-def job_arch(cu_num: int = 0, gfx: str = "") -> str:
-    """Target arch a job would compile for -- shared by dispatch and ARCH filtering."""
-    return gfx or cu_num_to_arch(cu_num, default=GEMM_AOT_ARCH_DEFAULT)
-
-
 def compile_one_config(
     kernel_name: str,
     kind: str,
@@ -667,7 +661,7 @@ def compile_one_config(
     """Compile one GEMM kernel configuration and save it to cache."""
     from torch._subclasses.fake_tensor import FakeTensorMode
 
-    aot_arch = job_arch(cu_num, gfx)
+    aot_arch = resolve_job_arch(cu_num, gfx)
     shape_str = f"{kernel_name}  M={m} N={n} K={k}"
     result = {
         "kernel_name": kernel_name,
@@ -744,7 +738,9 @@ def main():
         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
+            j
+            for j in all_jobs
+            if resolve_job_arch(j["cu_num"], j.get("gfx", "")) in arch_set
         ]
         print(f"[aiter] ARCH={arch}: {len(all_jobs)}/{n_before} jobs match")
 

aiter/aot/flydsl/grouped_moe.py

@@ -18,6 +18,7 @@
     compile_only_env,
     job_identity,
     override_env,
+    resolve_job_arch,
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
@@ -136,6 +137,7 @@ def parse_csv(csv_path: str):
                 "data_format": (
                     "fp4" if "float4" in row.get("q_dtype_a", "") else "a8w4"
                 ),
+                "cu_num": int(row.get("cu_num") or 0),
                 "gfx": row.get("gfx", ""),
             }
             for job in _scheduler_variants(row, base_job):
@@ -147,9 +149,6 @@ def parse_csv(csv_path: str):
     return jobs
 
 
-GROUPED_MOE_AOT_ARCH_DEFAULT = "gfx1250"
-
-
 def _compile_grouped_moe_aux_kernels(job, *, dtype, quant_mode, wmma_rep, contiguous):
     """Precompile the non-GEMM FlyDSL kernels the run-only grouped MoE fast path
     launches around gemm1/gemm2."""
@@ -377,7 +376,8 @@ def compile_one_config(**job):
     import torch
     from torch._subclasses.fake_tensor import FakeTensorMode
 
-    aot_arch = job.pop("gfx", "") or GROUPED_MOE_AOT_ARCH_DEFAULT
+    cu_num = job.pop("cu_num", 0)
+    gfx = job.pop("gfx", "")
     shape_str = (
         # Use .get() so a missing key can't raise here, outside the try below:
         # an escaping exception would crash the worker (exitcode != 0), which the
@@ -388,8 +388,10 @@ def compile_one_config(**job):
         f"contiguous={bool(job.get('grouped_contiguous_m', False))}"
     )
 
+    aot_arch = None
     t0 = time.time()
     try:
+        aot_arch = resolve_job_arch(cu_num, gfx)
         # TODO(aot): only the auxiliary (non-GEMM) kernels are precompiled here.
         # The grouped GEMM itself moved to the TDM batched kernel
         # (aiter.ops.flydsl.batched_gemm_mxfp4), which has no AOT wiring yet, so

aiter/aot/flydsl/moe.py

@@ -31,12 +31,13 @@
 from aiter.aot.flydsl.common import (
     collect_aot_jobs,
     compile_only_env,
-    cu_num_to_arch,
     job_identity,
     override_env,
+    resolve_job_arch,
     run_jobs_parallel,
 )
 from aiter.jit.core import AITER_CONFIGS
+from aiter.jit.utils.chip_info import warn_legacy_gfx_inference
 from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg as _ptr_view_safe
 from aiter.ops.flydsl.moe_kernels import (
     _get_compiled_silu_fused,
@@ -59,7 +60,6 @@
     AITER_CONFIGS.AITER_CONFIG_FMOE_FILE,
     AITER_CONFIGS.AITER_CONFIG_FHMOE_FILE,
 ]
-MOE_AOT_ARCH_DEFAULT = "gfx950"
 
 
 def parse_csv(csv_path: str):
@@ -77,6 +77,9 @@ def parse_csv(csv_path: str):
 
     with open(csv_path, newline="") as f:
         reader = csv.DictReader(f)
+        has_gfx = "gfx" in (reader.fieldnames or [])
+        if not has_gfx:
+            warn_legacy_gfx_inference(csv_path)
         for row in reader:
             token = int(row["token"])
             model_dim = int(row["model_dim"])
@@ -85,6 +88,7 @@ def parse_csv(csv_path: str):
             topk = int(row["topk"])
             doweight_stage1 = bool(int(row.get("doweight_stage1", "0")))
             cu_num = int(row.get("cu_num", "0"))
+            gfx = row.get("gfx", "").strip()
             block_m = int(row.get("block_m", "0") or "0")
             shared_expert_id = int(row.get("shared_expert_id", "-1") or "-1")
             act_type = row.get("act_type", "")
@@ -137,6 +141,7 @@ def parse_csv(csv_path: str):
                     "inter_dim": inter_dim,
                     "topk": topk,
                     "cu_num": cu_num,
+                    "gfx": gfx,
                     # Not used by the epilogue compile; zeroed so dedup keys on
                     # (act, inter_dim, topk, cu_num) only.
                     "model_dim": 0,
@@ -174,6 +179,7 @@ def parse_csv(csv_path: str):
                         "topk": topk,
                         "doweight_stage1": doweight_stage1,
                         "cu_num": cu_num,
+                        "gfx": gfx,
                         "act": act,
                         "enable_bias": enable_bias,
                         "token_num": token,
@@ -1003,6 +1009,7 @@ def compile_one_config(
     experts: int,
     topk: int,
     cu_num: int = 0,
+    gfx: str = "",
     **kwargs,
 ) -> dict:
     """Compile one MoE kernel configuration and save to cache.
@@ -1012,7 +1019,7 @@ def compile_one_config(
 
     Returns a dict with timing info.
     """
-    aot_arch = cu_num_to_arch(cu_num, default=MOE_AOT_ARCH_DEFAULT)
+    aot_arch = resolve_job_arch(cu_num, gfx)
     is_epilogue = kwargs.get("stage") == "epilogue"
     shape_str = (
         f"{kernel_name}  inter_dim={inter_dim} topk={topk}"
@@ -1074,6 +1081,7 @@ def compile_one_config(
 
                         precompile = precompile_fhmoe_to_cache
                         kwargs["shared_expert_id"] = shared_expert_id
+                        kwargs["gfx"] = gfx
                     else:
                         precompile = _precompile_to_cache
                     precompile(
@@ -1131,7 +1139,7 @@ def main():
     print(f"  Stage2 jobs:    {len(stage2_jobs)}")
     print(f"  Epilogue jobs:  {len(epilogue_jobs)}")
     print(f"  Total jobs:     {len(all_jobs)}")
-    print("  Compile arch: (from cu_num)")
+    print("  Compile arch: (from gfx, else known cu_num mapping)")
     print(f"  Cache dir:    {cache_dir}")
     print(f"  Target arch:  {arch}")
     print("=" * 72)