[Feature]: AITER path should allow for experimentation

[dimon777]

Suggestion Description

Hello,

I am opening this issue which stemmed from this VLLM thread Someone please upstream this gfx1201/RDNA4 FP8 Patch into vllm-rocm. I've built a docker image to validate OP contribution and confirming performance improvements for R9700 (RDNA4 arch). Building VLLM from the most recent code base however no longer works because AITER changed behavior to hard fail on:

import aiter.ops.triton.utils.arch_info as arch_info;

regardless of the value of VLLM_ROCM_USE_AITER (see below).

I know that support for consumer GPUs is not planned in AITER but I'd like to propose two things:

1. Do not hard fail AITER initialization on Python import for unsupported (yet) architectures. While PEP8 doesn't specifically prohibits implementing logic inside imports such behavior, it is a very questionable practice. We have classes, init(), configure() for a reason. Please do it downstream, please throw a warning, but don't fail hard there.
2. When the core development team builds AITER, please give it a small thought that other people want to contribute to it to improve it and eventually converge it to at least professional line of GPUs. So, please leave for us a path to "hack" it like we did in the above VLLM issue.

cc: schung-amd, Rob-P-Smith

Thank you.

root@06e7efa7d4f5:/app# export VLLM_ROCM_USE_AITER=0
root@06e7efa7d4f5:/app# python3
Python 3.12.12 (main, Oct 10 2025, 08:52:57) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import aiter.ops.triton.utils.arch_info as arch_info;
Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/utils/_triton/arch_info.py", line 13, in get_arch
    triton.runtime.driver.active.get_current_target().arch
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/triton/runtime/driver.py", line 30, in __getattr__
    return getattr(self._initialize_obj(), name)
                   ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/triton/runtime/driver.py", line 26, in _initialize_obj
    self._obj = self._init_fn()
                ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/triton/runtime/driver.py", line 11, in _create_driver
    raise RuntimeError(f"{len(active_drivers)} active drivers ({active_drivers}). There should only be one.")
RuntimeError: 0 active drivers ([]). There should only be one.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.12/dist-packages/aiter/__init__.py", line 63, in <module>
    from .ops.attention import *
  File "/usr/local/lib/python3.12/dist-packages/aiter/ops/attention.py", line 16, in <module>
    from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype
  File "/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/utils/types.py", line 33, in <module>
    e5m2_dtype, e4m3_dtype = get_fp8_dtypes()
                             ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/utils/types.py", line 14, in get_fp8_dtypes
    if arch_info.get_arch() in ("gfx950"):
       ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/utils/_triton/arch_info.py", line 16, in get_arch
    from jax._src.lib import gpu_triton as triton_kernel_call_lib
ModuleNotFoundError: No module named 'jax'

Operating System

Debian 12

GPU

R9700

ROCm Component

No response

[hurricane-dorian]

The devs are quite responsive to the community despite their employers official positions on what supported. I'm sure the behavior change was done for good reasons and if this case was overlook it was because they really didn't know how projects like LLVM are using their code.

[carlushuang]

cc @wuhuikx @zejunchen-zejun

[zejunchen-zejun]

Hi, @dimon777

Thank you for using AITER and AMD GPUs. First of all, the env flag VLLM_ROCM_USE_AITER is only used in VLLM, instead of AITER. It will not have impact on AITER behavior. Then, we need to check if AITER has plan on desktop GPUs.
This issue is related to how Triton handles unsupported GPU architectures during its initialization and based on that, AITER will import JAX, which is a new framework like Torch. Maybe you can install the JAX and check if the issue is gone.

Thank you.

[hurricane-dorian]

"aiter.ops.triton.utils.triton.arch_info" seems to have been moved to "aiter.ops.triton.utils._triton.arch_info"
a very easy to miss underscore has been added to the functions name "_triton.arch_info"

[hurricane-dorian]

There is a PR which adds basic definitions for gfx1201 and gfx908. Any chance it will be merged anytime soon?
add gfx1201 and gfx908 definitions #918