[Issue]: `generate.py: error: argument --target_gpus: expected at least one argument` if only unsupported GPUs are selected

[ScottTodd]

Problem Description

Summary

We're seeing build errors like https://github.com/ROCm/TheRock/actions/runs/25201744067/job/73894021744:

-- [AOTriton] Skipping triton due to AOTRITON_NOIMAGE_MODE
CMAKE_SOURCE_DIR /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime
CMAKE_CURRENT_SOURCE_DIR /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime/v3src
CMAKE_CURRENT_SOURCE_PARENT_DIR /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime
CMAKE_CURRENT_LIST_DIR /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime/v3src
CMAKE_CURRENT_BINARY_DIR /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime-build/v3src
-- AOTRITON_TARGET_ARCH gfx900
-- AOTRITON_OVERRIDE_TARGET_GPUS 
-- EFFECTIVE_TARGET_GPUS 
AOTRITON_COMPILER /__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime/v3python/compile.py
usage: generate.py [-h]
'/opt/_internal/cpython-3.12.10/lib/python3.12/site-packages/cmake/data/bin/cmake' '-E' 'env' 'VIRTUAL_ENV=/__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime-build/venv' 'AOTRITON_ENABLE_FP32=1' '/__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime-build/venv/bin/python' '-X' 'utf8' '-m' 'v3python.generate' '--target_gpus' '--build_dir' '/__w/TheRock/TheRock/external-builds/pytorch/pytorch/build/aotriton/src/aotriton_runtime-build/v3src' '--noimage_mode'
                   [--target_gpus {gfx90a_mod0,gfx942_mod0,gfx950_mod0,gfx1100_mod0,gfx1101_mod0,gfx1102_mod0,gfx1151_mod0,gfx1150_mod0,gfx1201_mod0,gfx1200_mod0,gfx1250_mod0} [{gfx90a_mod0,gfx942_mod0,gfx950_mod0,gfx1100_mod0,gfx1101_mod0,gfx1102_mod0,gfx1151_mod0,gfx1150_mod0,gfx1201_mod0,gfx1200_mod0,gfx1250_mod0} ...]]
                   [--build_dir BUILD_DIR] [--root_dir ROOT_DIR]
                   [--archive_only] [--library_suffix LIBRARY_SUFFIX]
                   [--noimage_mode] [--build_for_tuning]
                   [--build_for_tuning_second_pass]
                   [--build_for_tuning_but_skip_kernel [BUILD_FOR_TUNING_BUT_SKIP_KERNEL ...]]
                   [--verbose] [--lut_sanity_check]
generate.py: error: argument --target_gpus: expected at least one argument
CMake Error at v3src/CMakeLists.txt:71 (execute_process):
  execute_process failed command indexes:

    1: "Child return code: 2"
-- Configuring incomplete, errors occurred!

Context

In https://github.com/ROCm/TheRock we build pytorch using https://github.com/ROCm/TheRock/blob/main/external-builds/pytorch/build_prod_wheels.py in a few configurations:

• Per-family releases: separate builds for PYTORCH_ROCM_ARCH=gfx900, PYTORCH_ROCM_ARCH=gfx1151, PYTORCH_ROCM_ARCH=gfx942, etc.
• Multi-arch releases (new): one single build for PYTORCH_ROCM_ARCH=gfx1100;gfx1101;gfx1102;gfx1103;gfx1151;gfx1200;gfx1201;gfx900;gfx906;gfx908;gfx90a;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1150;gfx1152;gfx1153

See also:

• [Multi-arch] PyTorch packages build without flash attention when targets not supported by aotriton are included TheRock#4969
• Disable AOTriton for gfx900 in PyTorch builds  TheRock#4001

Analysis and remediation

If there is at least one supported architecture in the list then target filtering and fallback behavior appears to be working. If there are only unsupported architectures in the list then we hit the above error.

For per-family releases in TheRock we solved this by disabling flash attention and aotriton entirely if any unsupported architecture was included. I'm going to invert this for multi-arch releases to disable flash attention and aotriton only if all architectures are unsupported, now that we see how the filtering and fallback code paths are working.

If aotriton is patched to not error during generate.py (and produce a library where check_gpu() always returns failure) for target lists including only unsupported architectures then we can remove our downstream filtering in TheRock entirely.

Operating System

Linux and Windows

Steps to Reproduce

Build pytorch building with USE_FLASH_ATTENTION=ON and PYTORCH_ROCM_ARCH=gfx900 (or any set of archs with no aotriton support)

[ScottTodd]

Some raw notes from my coding agent debugging session with more specific code pointers:

debug findings:

● The failure is not in pytorch's flash attention compilation — it's in aotriton's configure step:

generate.py: error: argument --target_gpus: expected at least one argument
FAILED: aotriton/src/aotriton_runtime-stamp/aotriton_runtime-configure

What's happening:

1. PYTORCH_ROCM_ARCH=gfx900 is passed to aotriton via -DAOTRITON_TARGET_ARCH=gfx900
2. aotriton's gpu_targets.py silently filters gfx900 out (not supported), producing an empty target list
3. The empty list is passed to generate.py --target_gpus which fails because argparse requires at least one argument

This only happens when all archs are unsupported. The gfx908;gfx110X-all build succeeded because gfx110X targets are
supported by aotriton. The real multi-arch release builds (which always include gfx11xx/gfx120x/gfx94X) will always
have supported targets.

That said, the gfx900-only case is a regression from the old behavior (which would have disabled flash attention). A
couple of options:

1. Leave as-is — gfx900-only with flash attention is a nonsensical configuration; users can pass
   --no-enable-pytorch-flash-attention-linux if they need it. Multi-arch (the actual target) works.
2. Add a positive check — enable flash attention only when at least one supported arch is in the list. This handles
   the edge case but reintroduces an arch list we'd need to maintain.
3. Fix in aotriton — generate.py should handle an empty target list gracefully (skip shim generation, produce a stub
   library).

What's your preference?

me:

• 1 won't work, we're still building both single-arch and multi-arch releases concurrently.
• 2 is viable, flip our check to "enabled if at least one supported arch is included" or "disable if all archs are
  unsupported" instead of "disabled if at least one unsupported arch is included".
• How would 3 work? we'd enable flash attention in pytorch which enables the aotriton build but then aotriton would
  correctly filter down to 0 targets? what happens for those targets when compiled together with other supported targets? does pytorch fallback on its own to other attention implementations for just those targets?

response:

● Good question on option 3. Let me trace the full path to answer it precisely.

When aotriton is built with supported targets (e.g. gfx1100) but not gfx900:

1. Build time: aotriton produces libaotriton_v2.so with shim dispatch code only for supported targets. Kernel images
   are only generated/downloaded for supported targets (e.g. aotriton.images/amd-gfx11xx/).
2. Link time: pytorch links torch_hip against libaotriton_v2.so and compiles with -DUSE_FLASH_ATTENTION. This is
   arch-independent C++ — all archs get the same flash attention code compiled.
3. Runtime on gfx900: The SDP backend selection in sdp_utils.cpp:285-304 calls aotriton::v2::flash::check_gpu(stream).
   This checks if the current GPU has kernel images available. For gfx900 there are none, so it returns failure. pytorch
   then falls back through its priority order to math/CK backends.

So yes — in a multi-arch build with both gfx900 and gfx1100, flash attention code is compiled for all targets, but at
runtime gfx900 gracefully falls back while gfx1100 uses flash attention. That's the existing runtime dispatch working
as designed.

The issue with option 3 is specifically the gfx900-only build. If aotriton filters to zero targets, generate.py
crashes because it can't generate shim code for zero GPUs. If that were fixed (generate an empty/stub shim), then
libaotriton_v2.so would be a library where check_gpu() always fails → pytorch always falls back. That would work but
it's a change in the aotriton repo, not TheRock.

For multi-arch builds (the actual target of issue #4969), option 3 is already how things work — our change is
sufficient. The only gap is single-arch builds targeting unsupported-only families.

Given that you're still running single-arch releases, option 2 is the pragmatic fix. Let me implement it.

[xinyazhang]

The point is you should disable the build of aotriton rather than trying to silent this error by making a stub library, which only defers a problem detectable at compile-time to runtime and causes more problems in later phases.

[ScottTodd]

The new code I'm proposing in ROCm/TheRock#4986 is

    # aotriton supports a subset of GPU architectures. When at least one
    # target arch is supported, we enable flash attention and let aotriton's
    # build system (gpu_targets.py) filter to just the supported ones. The
    # runtime (check_gpu in sdp_utils.cpp) gracefully falls back to math/CK
    # backends on unsupported GPUs. We only disable flash attention when
    # *no* target arch is supported — otherwise aotriton's configure step
    # fails on the empty target list (https://github.com/ROCm/aotriton/issues/169).
    #
    # These prefixes match what aotriton's gpu_targets.py recognizes.
    # See also the image list in pytorch/cmake/External/aotriton.cmake.
    AOTRITON_SUPPORTED_ARCH_PREFIXES = ("gfx90a", "gfx942", "gfx950", "gfx11", "gfx12")
    # gfx1152/53: supported in aotriton 0.11.2b+ (https://github.com/ROCm/aotriton/pull/142),
    #   which is pinned by pytorch >= 2.11. Older versions don't include it.
    aotriton_unsupported_for_version = []
    if not is_pytorch_2_11_or_later:
        aotriton_unsupported_for_version = ["gfx1152", "gfx1153"]
    rocm_arch_list = env.get("PYTORCH_ROCM_ARCH", "").split(";")
    has_aotriton_supported_arch = any(
        arch.startswith(AOTRITON_SUPPORTED_ARCH_PREFIXES)
        and arch not in aotriton_unsupported_for_version
        for arch in rocm_arch_list
    )

            use_flash_attention = (
                "ON" if triton_requirement and has_aotriton_supported_arch else "OFF"
            )

We'd rather this filtering go in pytorch or aotriton itself, where specific knowledge about what is and isn't supported for a given version of each library is available. A developer (or script) building pytorch shouldn't be expected to know that USE_FLASH_ATTENTION is safe to set for gfx900;gfx942 but not for gfx900, especially since pytorch and aotriton already have some filtering and code fallback paths.

[xinyazhang]

We'd rather this filtering go in pytorch or aotriton itself

That's exactly what I'm thinking of. USE_FA/ME should be turned off if no supported arch, but it may take some time since the logic a little complicated.