From 9a9792e950655cceb06e16ea4acc0ecc2065538b Mon Sep 17 00:00:00 2001 From: simondanielsson Date: Fri, 3 Jul 2026 10:10:28 +0200 Subject: [PATCH] fix: grouped moe build respects GPU_ARCHS Signed-off-by: simondanielsson --- aiter/aot/flydsl/grouped_moe.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/aiter/aot/flydsl/grouped_moe.py b/aiter/aot/flydsl/grouped_moe.py index 68ba380c391..587fd47f853 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,12 +22,21 @@ run_jobs_parallel, ) from aiter.jit.core import AITER_CONFIGS +from aiter.jit.utils.build_targets import _parse_gpu_archs_env DEFAULT_CSVS = [AITER_CONFIGS.AITER_CONFIG_GROUPED_FMOE_FILE] _WARP_TILE_N = 64 _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) @@ -123,6 +133,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: @@ -133,6 +144,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)