diff --git a/aiter/aot/flydsl/grouped_moe.py b/aiter/aot/flydsl/grouped_moe.py index 91b6467c43a..692b2fc38af 100644 --- a/aiter/aot/flydsl/grouped_moe.py +++ b/aiter/aot/flydsl/grouped_moe.py @@ -9,6 +9,7 @@ import argparse import csv +import os import sys import time import traceback @@ -21,6 +22,7 @@ run_jobs_parallel, ) from aiter.jit.core import AITER_CONFIGS +from aiter.jit.utils.build_targets import _parse_gpu_archs_env from aiter.ops.flydsl.kernels.tensor_shim import ptr_arg DEFAULT_CSVS = [AITER_CONFIGS.AITER_CONFIG_GROUPED_FMOE_FILE] @@ -28,6 +30,14 @@ _TILE_K = 256 +def _target_gfx_archs() -> set[str] | None: + """Requested build archs from GPU_ARCHS, or None when unset/native.""" + gfx_env = os.getenv("GPU_ARCHS") + if not gfx_env or gfx_env.strip().lower() == "native": + return None + return set(_parse_gpu_archs_env(gfx_env)) + + def _align_up(value: int, alignment: int) -> int: return ((int(value) + int(alignment) - 1) // int(alignment)) * int(alignment) @@ -124,6 +134,7 @@ def _scheduler_variants(row, base_job): def parse_csv(csv_path: str): jobs = [] seen = set() + target_archs = _target_gfx_archs() with open(csv_path, newline="") as f: for row in csv.DictReader(f): if not row: @@ -134,6 +145,10 @@ def parse_csv(csv_path: str): for col in ("model_dim", "inter_dim", "expert", "token") ): continue + # Skip rows whose target arch is not in the requested GPU_ARCHS + row_gfx = (row.get("gfx") or "").strip() + if target_archs is not None and row_gfx and row_gfx not in target_archs: + continue n_warp = int(row.get("n_warp") or 4) token_num = int(row["token"]) tile_m = int(row.get("tile_m") or 64)