[FlyDSL] [Bugfix] Use FLYDSL_GPU_ARCH LDS limit so gfx950 8-wave AOT isn't excluded from builds on gfx942

aiter/ops/flydsl/utils.py

@@ -4,6 +4,7 @@
 """General utilities shared across all FlyDSL kernel families."""
 
 import importlib.util
+import os
 from functools import cache, lru_cache
 
 import torch
@@ -47,7 +48,15 @@ def _get_shared_memory_per_block_cached(device_index: int, fallback_gfx: str) ->
 
 
 def get_shared_memory_per_block(device=None, fallback_gfx: str = "") -> int:
-    """Return per-block shared memory/LDS limit for the active device."""
+    """Return per-block shared memory/LDS limit for the compile/run target.
+
+    AOT sets FLYDSL_GPU_ARCH to the job's target, which may differ from the
+    live GPU. Prefer that over torch.cuda device properties so a gfx942
+    builder can compile gfx950 kernels that need more than 64 KiB LDS.
+    """
+    aot_arch = os.environ.get("FLYDSL_GPU_ARCH", "").strip()
+    if aot_arch:
+        return addressable_lds_bytes_for_gfx(aot_arch)
     if device is None:
         device = _default_cuda_device_index()
     elif isinstance(device, torch.device):

op_tests/flydsl_tests/test_flydsl_utils.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+"""CPU unit tests for FlyDSL helpers (no GPU required)."""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import torch
+
+_UTILS_PATH = (
+    Path(__file__).resolve().parents[2] / "aiter" / "ops" / "flydsl" / "utils.py"
+)
+
+
+def _load_utils():
+    spec = importlib.util.spec_from_file_location(
+        "flydsl_utils_under_test", _UTILS_PATH
+    )
+    mod = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def test_lds_limit_follows_aot_arch_not_live_device(monkeypatch):
+    """LDS budget must come from the arch we compile for, not the GPU we build on.
+
+    A multi-arch wheel (GPU_ARCHS='gfx942;gfx950') compiles gfx950 8-wave tiles
+    even on a gfx942 host. Reading the host's 64 KiB limit rejected tiles that
+    fit gfx950's 160 KiB, aborting FlyDSL AOT before any wheel was written.
+    """
+
+    class Props:
+        shared_memory_per_block = 65536
+        gcnArchName = "gfx942"
+
+    monkeypatch.setenv("FLYDSL_GPU_ARCH", "gfx950")
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(torch.cuda, "get_device_properties", lambda i: Props())
+
+    flydsl_utils = _load_utils()
+    flydsl_utils._default_cuda_device_index.cache_clear()
+    flydsl_utils._get_shared_memory_per_block_cached.cache_clear()
+
+    assert flydsl_utils.get_shared_memory_per_block(fallback_gfx="gfx950") == 163840