[JIT] AITER extensions can bind HIP symbols from an earlier global TileLang stub

Summary

AITER native extensions are currently sensitive to process library load order. If TileLang's HIP stub is loaded with RTLD_GLOBAL before an AITER extension, the AITER extension can resolve a HIP runtime symbol from the stub instead of from libamdhip64.so.

This is the same class of ROCm regression reported in vLLM issue #51151. That issue records the vLLM import chain, CI failures, and an AITER topk_softmax/moe_sorting reproducer. This ticket adds a smaller AITER FMHA reproducer and identifies the specific TileLang device-properties ABI mismatch.

One visible result is that a valid fmha_v3_varlen_fwd invocation on MI355X is rejected with:

RuntimeError: invalid argument for fmha_v3_varlen_fwd

The same bad binding can corrupt cached device properties used by AITER MoE paths and later surface as hipErrorInvalidConfiguration / invalid configuration argument, as observed in vLLM issue #51151.

The same tensors and AITER operation pass if the TileLang stub is not loaded, or if the exact same stub is loaded with RTLD_LOCAL.

This is a real-library reproducer: it uses the TileLang and AITER packages already installed in the vLLM ROCm nightly image. It does not build or inject a synthetic interposer.

Environment

• Image tag: vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263
• Immutable image digest: sha256:d22922d540810d90c5a3eafe91d3b4a62c2b881f3e990d0b7180b2875a5d176d
• Image build date: 2026-08-04
• vLLM: 0.26.1rc1.dev306+gcb8104839.rocm723
• AITER: 0.1.19
• TileLang: 0.1.10
• PyTorch HIP runtime: 7.2.53211
• GPU: AMD Instinct MI355X (gfx950)

Minimal reproduction

docker run --rm -i \
  --entrypoint python3 \
  --device=/dev/kfd \
  --device=/dev/dri \
  --group-add video \
  --ipc=host \
  --shm-size=16g \
  vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263@sha256:d22922d540810d90c5a3eafe91d3b4a62c2b881f3e990d0b7180b2875a5d176d - <<'PY'
import ctypes
import importlib.util
import math
from pathlib import Path

import torch

# Locate the installed TileLang library without importing TileLang itself.
spec = importlib.util.find_spec("tilelang")
assert spec is not None and spec.submodule_search_locations is not None
tilelang_root = Path(next(iter(spec.submodule_search_locations)))
hip_stub = tilelang_root / "lib" / "libhip_stub.so"
print(f"loading {hip_stub} with RTLD_GLOBAL")
ctypes.CDLL(str(hip_stub), mode=ctypes.RTLD_GLOBAL)

from aiter.ops.mha import fmha_v3_varlen_fwd

q = torch.randn((128, 8, 128), device="cuda", dtype=torch.bfloat16)
cu_seqlens = torch.tensor([0, 128], device="cuda", dtype=torch.int32)

fmha_v3_varlen_fwd(
    q,
    q,
    q,
    cu_seqlens,
    cu_seqlens,
    128,
    128,
    0,
    0.0,
    1.0 / math.sqrt(128),
    0.0,
    False,
    False,
    -1,
    -1,
    False,
    False,
    1,
)
torch.cuda.synchronize()
print("PASS")
PY

This loads the exact libhip_stub.so shipped by TileLang instead of a synthetic interposer. Loading it directly keeps the reproducer focused on the condition AITER observes. In the original vLLM process, the same global state is reached through import tilelang and the bundled TVM runtime, as documented in vLLM issue #51151.

Actual result

RuntimeError: invalid argument for fmha_v3_varlen_fwd

Expected result

The valid FMHA invocation should complete successfully and should not depend on whether another optional backend was imported first.

Controlled results

Using the same image, tensors, AITER operation, and TileLang library:

Setup before importing AITER	Result
Do not load libhip_stub.so	PASS
Load libhip_stub.so with RTLD_LOCAL	PASS
Load libhip_stub.so with RTLD_GLOBAL	FAIL: invalid argument
Load the stub globally, then promote libamdhip64.so globally	FAIL

The last control is important: loading the real HIP runtime with RTLD_GLOBAL after the stub is already global does not repair the existing ELF global lookup order.

Dynamic-loader evidence

With LD_DEBUG=bindings, glibc reports the AITER FMHA extension binding the versioned device-properties symbol to TileLang's stub:

binding file .../aiter/jit/module_fmha_v3_varlen_fwd.so [0]
  to .../tilelang/lib/libhip_stub.so [0]:
  normal symbol `hipGetDevicePropertiesR0600' [hip_6.0]

The call then fails with invalid argument for fmha_v3_varlen_fwd.

readelf also shows that the TileLang stub exports this symbol:

FUNC GLOBAL ... hipGetDevicePropertiesR0600

The reason the stub becomes global is visible in the packaged libraries and TVM loader:

$ readelf -d tilelang/lib/libtvm_runtime.so | grep hip_stub
NEEDED  Shared library: [libhip_stub.so]

TileLang imports its bundled TVM package, whose tvm/base.py opens libtvm_runtime.so with RTLD_GLOBAL. That mode applies to the dependency group, so libhip_stub.so also enters the global lookup scope. In the pinned image, the stub exports 30 standard hip* functions, including hipMalloc, hipFree, hipGetDevice, hipModuleLaunchKernel, and hipGetDevicePropertiesR0600.

Root cause

This is ELF symbol interposition, not an invalid FMHA input. TileLang 0.1.10 has two relevant cross-library problems.

1. A private compatibility stub exports standard HIP names globally

TileLang #1867 added libhip_stub.so to allow lazy loading of the real HIP runtime. The stub is a DT_NEEDED dependency of the globally loaded TVM runtime and exports drop-in hip* wrappers. AITER later imports its native extension. Although that extension depends on the real HIP runtime, a matching definition in the earlier global object can preempt the definition in the extension's dependency group.

2. The device-properties wrapper resolves the wrong ABI

Current ROCm headers map both the source-level function and structure type to the ROCm 6 ABI:

#define hipGetDeviceProperties hipGetDevicePropertiesR0600
#define hipDeviceProp_t hipDeviceProp_tR0600

TileLang's wrapper is written using these source-level names, so preprocessing makes it an exported hipGetDevicePropertiesR0600 function accepting an hipDeviceProp_tR0600*. However, its dispatch table uses:

LOOKUP(hipGetDeviceProperties_, "hipGetDeviceProperties")

The C preprocessor does not expand a macro inside a string literal. The R0600 wrapper therefore calls the legacy unversioned function. That function writes a different device-properties layout into storage interpreted by AITER as hipDeviceProp_tR0600.

The resulting chain is:

import tilelang
  -> TVM loads libtvm_runtime.so with RTLD_GLOBAL
  -> DT_NEEDED loads libhip_stub.so into global scope

load an AITER extension
  -> hipGetDevicePropertiesR0600 binds to TileLang's wrapper
  -> wrapper calls legacy hipGetDeviceProperties
  -> AITER consumes invalid R0600 device properties
  -> FMHA rejects valid inputs, or a later kernel gets an invalid launch grid

TileLang should correct both the versioned wrapper/ABI handling and the global exposure of standard HIP names. However, an AITER extension can encounter other global interposers too. Relying on optional-backend import order leaves AITER fragile even after this particular TileLang bug is fixed.

Why this became visible in vLLM

vLLM PR #50879 moved optional expert-parallel imports into their selected runtime branches. This also made the guarded Mori import lazy. Previously, importing Mori opened the real HIP runtime with RTLD_GLOBAL before TileLang, creating a favorable symbol order that accidentally masked the problem.

Without that side effect, vLLM can probe TileLang first and expose the existing interposition bug. vLLM issue #51151 documents this regression and the concrete vLLM import chain. #50879 changed when the bug becomes visible; it did not create the TileLang ABI mismatch.

vLLM PR #51110 restores the guarded eager Mori import and can mask the symptom for that startup path. vLLM PR #51159 instead defers TileLang import on ROCm and fixes the reported vLLM startup path. Both are import-order compatibility fixes, not general AITER loader fixes or corrections to TileLang's exported symbols.

Suggested TileLang direction

The originating corrections belong in TileLang:

1. Resolve the ABI-selected function name. For example, stringify the expanded hipGetDeviceProperties macro instead of hard-coding the unversioned string.
2. Avoid exposing standard HIP wrapper names through a globally loaded dependency. Hidden visibility or TileLang-prefixed wrapper symbols would prevent the stub from interposing on unrelated libraries.

These changes fix the concrete TileLang defect. They do not remove the value of making AITER's native loader resilient to other global interposers.

Suggested AITER direction

Could the AITER team make native extension loading independent of earlier global HIP interposers?

One possible Linux/glibc design is to load each AITER native object with local, per-object dependency preference (for example RTLD_NOW | RTLD_DEEPBIND) before executing/importing it. An equivalent maintainer-preferred linker or loader design would also be suitable.

The protection needs to be centralized and consistently cover all AITER paths that load native objects:

• Python/pybind JIT extension imports;
• Python ctypes loading of standalone JIT libraries;
• generated Python AOT/JIT library loading;
• nested C++ dlopen calls used by generated wrappers.

If RTLD_DEEPBIND is selected, an opt-out such as AITER_DISABLE_DEEPBIND=1 would be useful for ASan, LD_PRELOAD allocators, profilers, or intentional interposition. A process-wide sys.setdlopenflags() change should preferably be avoided because it affects unrelated concurrent imports.

Please also add a regression test that places a HIP-symbol interposer in global scope before loading an AITER extension and verifies that the extension still uses the intended HIP implementation.

Ownership and reference implementation

AITER PR #4566 is provided as an investigation, reference implementation, and regression test. It is not intended to prescribe the final design or require the code to be merged verbatim. The AITER team should feel free to:

• implement the loader fix differently;
• reuse only the reproducer or regression tests;
• reuse selected parts of the reference implementation; or
• coordinate the exact compatibility contract with TileLang.

The earlier experiment in AITER PR #4566 promoted the real HIP runtime with RTLD_GLOBAL. That is insufficient once the stub is already earlier in global lookup scope; the updated proposal uses per-object isolation instead.

Related work

• AITER PR #4566
• vLLM PR #50879
• vLLM PR #51110
• vLLM issue #51151
• vLLM PR #51159
• TileLang PR #1867
• TileLang v0.1.10 stub source
• Bundled TVM RTLD_GLOBAL loader

Start with the minimal Python reproducer and AITER PR #4566, then trace the native-loading paths named here: pybind JIT imports, ctypes standalone libraries, generated AOT/JIT loaders, and nested C++ dlopen calls. Done means a centralized loader protects each path from earlier global HIP interposers and a regression test passes with a global HIP-symbol interposer loaded first.