fix(jit): isolate AITER extensions from HIP interposers

[JohnQinAMD]

Summary

Load AITER-owned native objects with per-object RTLD_DEEPBIND where it is available. This prevents a HIP compatibility stub that was loaded earlier into the ELF global scope from interposing incompatible HIP symbols in AITER JIT extensions.

Related to AITER issue #4585, which contains the standalone real-library reproducer, controlled loader results, and the full TileLang/vLLM root-cause analysis.

The loader covers all native-library paths currently used by AITER:

• pybind extensions imported by aiter.jit.core;
• standalone libraries loaded with ctypes;
• generated Python loaders; and
• nested libraries opened by the C++ SharedLibrary helper.

The implementation does not change the process-wide sys.setdlopenflags() state. AITER_DISABLE_DEEPBIND=1 provides a process-start opt-out for ASan, intentional LD_PRELOAD interposition, or tools that are incompatible with deep binding.

This PR is intended as a reference implementation and regression test, depends on how necessary of the protection is needed.

Problem

This is an ELF symbol-interposition failure caused by the interaction of a TileLang 0.1.10 HIP stub, TVM's global loading policy, and AITER extensions that are loaded later. It is not an invalid FMHA input. The vLLM regression, import chain, CI failures, and original AITER MoE reproducer are tracked in vLLM issue #51151.

What TileLang's HIP stub is intended to do

TileLang #1867 added libhip_stub.so so TileLang/TVM can be imported without eagerly linking the real HIP runtime. The stub exports drop-in HIP wrapper functions. On first use, it locates libamdhip64.so and resolves the real functions with dlopen/dlsym.

That implementation creates two separate cross-library problems in TileLang 0.1.10.

1. The private stub enters the process-global symbol scope

Importing TileLang imports its bundled TVM Python package. TVM deliberately opens libtvm_runtime.so with RTLD_GLOBAL:

_LIB_RUNTIME = libinfo.load_lib_ctypes(
    "tvm", "tvm_runtime", "RTLD_GLOBAL", extra_lib_paths=_extra_lib_paths
)

The packaged libtvm_runtime.so has a DT_NEEDED dependency on libhip_stub.so:

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

The global loading mode applies to that dependency group. Consequently, the stub's standard hip* wrapper names become eligible to satisfy symbol references from unrelated libraries loaded later in the process. In the pinned image, the installed stub exports 30 HIP functions, including hipMalloc, hipFree, hipGetDevice, hipModuleLaunchKernel, and:

$ readelf -Ws tilelang/lib/libhip_stub.so | grep hipGetDeviceProperties
FUNC GLOBAL ... hipGetDevicePropertiesR0600

An AITER JIT extension normally loaded afterward can therefore resolve a HIP reference to TileLang's stub even though the extension itself has a dependency on the real libamdhip64.so.

2. hipGetDeviceProperties crosses an incompatible ABI boundary

With current ROCm headers, the source-level API and type are macros:

#define hipGetDeviceProperties hipGetDevicePropertiesR0600
#define hipDeviceProp_t hipDeviceProp_tR0600

The wrapper in TileLang v0.1.10 hip.cc is written with those source-level names:

hipError_t hipGetDeviceProperties(hipDeviceProp_t* prop, int device_id) {
  return HIPDriverAPI::get()->hipGetDeviceProperties_(prop, device_id);
}

The preprocessor therefore compiles this as a wrapper exported under hipGetDevicePropertiesR0600 and taking an hipDeviceProp_tR0600*. However, the dispatch table resolves a hard-coded string:

LOOKUP(hipGetDeviceProperties_, "hipGetDeviceProperties")

Macro expansion does not happen inside a string literal. The R0600 wrapper therefore calls the legacy, unversioned function. The callee writes a different device-properties layout into memory that the caller interprets as hipDeviceProp_tR0600. The call can return without an obvious loader error, but fields used for device selection and kernel launch sizing are invalid.

How this reaches AITER and Kimi-K3

The failing order in vLLM is:

vLLM probes an optional TileLang kernel
  -> import tilelang
  -> TVM loads libtvm_runtime.so with RTLD_GLOBAL
  -> DT_NEEDED loads libhip_stub.so into global scope

AITER later imports module_fmha_v3_varlen_fwd.so
  -> its hipGetDevicePropertiesR0600 reference binds to TileLang's stub
  -> the stub forwards to legacy hipGetDeviceProperties
  -> AITER reads an R0600 structure populated with the wrong layout
  -> a valid FMHA call is rejected

LD_DEBUG=bindings confirms the critical binding in the Kimi-K3 process:

module_fmha_v3_varlen_fwd.so -> tilelang/lib/libhip_stub.so:
hipGetDevicePropertiesR0600 [hip_6.0]

A valid BF16/head-size-128 invocation then fails with the misleading error:

RuntimeError: invalid argument for fmha_v3_varlen_fwd

The same mechanism also explains the invalid configuration argument failures seen in AITER MoE paths: corrupted device properties are cached and later used to calculate a kernel launch grid. This was the initially reported failure in vLLM issue #51151, which shows the TileLang import chain and a minimal topk_softmax/moe_sorting reproducer.

The following controls use the same image, tensors, AITER operation, and real TileLang library. They isolate global visibility as the condition that changes the result:

State before the AITER extension is loaded	Result
Do not load libhip_stub.so	PASS
Load the same stub with RTLD_LOCAL	PASS
Load the same stub with RTLD_GLOBAL	FAIL: invalid argument
Load the stub globally, then promote libamdhip64.so globally	FAIL

Why vLLM #50879 exposed it

vLLM #50879 changed optional expert-parallel imports from eager to selected-branch imports. Before that change, the guarded Mori import happened early and opened the real HIP runtime with RTLD_GLOBAL before TileLang. That favorable load order accidentally masked the TileLang interposer. The resulting ROCm regression was reported and analyzed in vLLM issue #51151.

After Mori became lazy, Kimi-K3—which does not select Mori—could import TileLang first. The PR exposed an existing load-order dependency; it did not create the TileLang ABI bug. vLLM #51110 restores the guarded eager Mori import as a compatibility hotfix, but it only restores the masking order. vLLM #51159 instead defers the TileLang import on ROCm and fixes the reported vLLM startup path, but it also avoids the trigger by changing import timing rather than correcting TileLang's exported symbols or protecting AITER's loader.

Permanent TileLang correction

The originating corrections belong in TileLang:

1. resolve the ABI-selected symbol, for example by stringifying the expanded hipGetDeviceProperties macro rather than using an unversioned literal; and
2. avoid exporting standard HIP compatibility symbols through a globally loaded dependency, for example by using hidden visibility or TileLang-prefixed wrapper names.

This AITER change is independent defensive hardening. Even after TileLang fixes this specific stub, an AITER extension can encounter other global interposers. AITER's native objects should not accidentally bind their HIP dependencies to an unrelated optional backend solely because that backend was imported first.

Why promoting libamdhip64 with RTLD_GLOBAL is insufficient

The previous version of this PR promoted the real libamdhip64.so handle into the global scope. That does not fix a process in which the incompatible stub is already earlier in ELF global lookup order; reopening the real HIP runtime does not reorder that scope. The controlled reproduction still fails in this order:

RTLD_GLOBAL TileLang stub -> RTLD_GLOBAL libamdhip64 -> import AITER extension

Likewise, eagerly importing AITER before TileLang only changes load order. It can mask the problem, but it does not protect later or lazily compiled AITER extensions.

Fix

For a pybind extension, _deep_import() first resolves the extension path and opens that exact object with:

os.RTLD_NOW | os.RTLD_DEEPBIND

The normal Python import then reuses the already loaded object. AITER retains the ctypes.CDLL handle for process lifetime. This is scoped to the AITER object and avoids a temporary process-global flag change, which would be racy with unrelated imports in other threads.

The same policy is applied directly to AITER's other ctypes and C++ dlopen paths. On platforms where RTLD_DEEPBIND is unavailable, the added flag is zero and the existing loader behavior is preserved.

Set the following before importing AITER to disable the policy:

export AITER_DISABLE_DEEPBIND=1

This opt-out is necessary for AddressSanitizer because ASan explicitly rejects objects opened with RTLD_DEEPBIND. It is also useful when HIP interposition is deliberate. Disabling it restores the original loader behavior and therefore also re-exposes AITER to this TileLang failure.

Reproduction

This deterministic reproducer explicitly places TileLang's real HIP stub in the
global ELF scope before AITER loads its FMHA extension. It does not depend on
vLLM's current optional-kernel import order.

Image: vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263
Digest: sha256:d22922d540810d90c5a3eafe91d3b4a62c2b881f3e990d0b7180b2875a5d176d
AITER 0.1.19, TileLang 0.1.10, ROCm 7.2.3, gfx950

IMAGE=vllm/vllm-openai-rocm:nightly-cb8104839c141609d99f1254459ef3a4f1bd4263
DIGEST=sha256:d22922d540810d90c5a3eafe91d3b4a62c2b881f3e990d0b7180b2875a5d176d

docker run --rm -i \
  --device=/dev/kfd --device=/dev/dri --group-add video --ipc=host \
  --entrypoint python3 "${IMAGE}@${DIGEST}" - <<'PY'
import ctypes
import importlib.util
import math
from pathlib import Path

import torch

spec = importlib.util.find_spec("tilelang")
root = Path(next(iter(spec.submodule_search_locations)))
ctypes.CDLL(str(root / "lib" / "libhip_stub.so"), 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 = torch.tensor([0, 128], device="cuda", dtype=torch.int32)
fmha_v3_varlen_fwd(
    q, q, q, cu, cu, 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

The existing loader fails with:

RuntimeError: invalid argument for fmha_v3_varlen_fwd

With this PR, the same invocation prints PASS. Setting
AITER_DISABLE_DEEPBIND=1 restores the failure, confirming that the result
comes from loader isolation. Full controls and loader evidence are in
AITER issue #4585.

Validation

The regression test covers pybind, standalone ctypes, generated-library,
MoE sorting, and Kimi-K3 FMHA loading paths.

pytest -q -p no:cacheprovider op_tests/test_jit_loader.py
6 passed

ruff check --no-cache <changed Python files>
All checks passed!

black --check <changed Python files>
All done!

The C++ loader change also passes a syntax-only compile in the same container.

Kimi-K3 model A/B

The real Kimi-K3 MXFP4 checkpoint was tested on 8x MI355X with TP8 and a 1M
maximum context. The test used the immutable image above, vLLM source at
d31de3c421c0281a959ac1a3cbe1a7f354bae179, and a controlled
Torch -> global TileLang stub -> AITER load order.

Existing loader:
  encoder profiling fails with invalid argument for fmha_v3_varlen_fwd

This PR:
  gfx950 BF16 FMHA loads
  19/19 decode graphs complete
  API health and a real chat completion return HTTP 200
  invalid argument/configuration errors: 0

The pinned nightly has an unrelated Gluon API mismatch: Triton expects
block_bases, while the tested AITER MLA source passes cga_layout. The
end-to-end run therefore used VLLM_ROCM_USE_AITER_MLA=0. The affected ViT
FMHA path remained AITER-backed and exercised the loader changed by this PR.

No documentation file is included in this PR.

Related

• AITER issue #4585
• vLLM #50879
• vLLM #51110
• vLLM #51151
• vLLM #51159
• TileLang #1867

Prepared with OpenAI Codex assistance.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4566 --add-label <label>

[Copilot]

Pull request overview

This PR hardens AITER’s JIT extension loading by ensuring the HIP runtime (libamdhip64.so) is promoted into the process-global dynamic symbol scope before any AITER pybind extension is imported, avoiding import-order–dependent initialization failures.

Changes:

• Add a one-time, process-lifetime loader _load_hip_runtime_global() in aiter/jit/core.py that dlopens libamdhip64.so with RTLD_GLOBAL.
• Call the loader immediately before the first JIT extension import in get_module_custom_op.
• Add unit tests validating (1) global load occurs once and doesn’t mutate Python’s dlopen flags, and (2) HIP promotion happens before extension import.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
aiter/jit/core.py	Promotes HIP runtime to RTLD_GLOBAL once before importing the first JIT extension module.
op_tests/test_jit_loader.py	Adds regression tests for one-time global HIP load and correct loader/import ordering.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[zufayu]

Q: does RTLD_DEEPBIND affect HIP API tracing?
It disables symbol interposition, which is how some profilers observe HIP calls — callback-based tools are fine, interposing ones would silently stop seeing AITER. Could you check rocprofv3 --hip-trace on one op with AITER_DISABLE_DEEPBIND=1 vs 0 and confirm the API counts match? If they don't, worth documenting next to the ASan note.

Nit: RTLD_NOW here is an unrelated behavior change
ctypes.DEFAULT_MODE is 0, so this line goes lazy → eager independently of RTLD_DEEPBIND — undeclared, and the PR's other three sites keep their existing mode. Eager resolution turns a never-called dangling reference into a load-time undefined symbol. Suggest mode=os.RTLD_LAZY | _RTLD_DEEPBIND.

[JohnQinAMD]

Q: does RTLD_DEEPBIND affect HIP API tracing? It disables symbol interposition, which is how some profilers observe HIP calls — callback-based tools are fine, interposing ones would silently stop seeing AITER. Could you check rocprofv3 --hip-trace on one op with AITER_DISABLE_DEEPBIND=1 vs 0 and confirm the API counts match? If they don't, worth documenting next to the ASan note.

@zufayu Thanks for raising this. I tested both the pybind and ctypes loader paths with rocprofv3 --hip-trace.

Test environment:

• Docker
• MI355X / gfx950
• ROCm 7.2.3
• rocprofv3 1.1.0
• Python 3.12.13
• PR commit e6bc3ade18

From a container where this PR checkout has been built, the test can be reproduced with:

export HIP_VISIBLE_DEVICES=0
export PYTHONPATH="$PWD${PYTHONPATH:+:$PYTHONPATH}"
TRACE_ROOT=$(mktemp -d /tmp/aiter-rocprof.XXXXXX)

profile_op() {
    local name="$1"
    local code="$2"

    for mode in 0 1; do
        AITER_DISABLE_DEEPBIND="$mode" \
            rocprofv3 \
                --log-level error \
                --hip-trace \
                --stats \
                --output-format csv \
                --output-directory \
                    "$TRACE_ROOT/$name/deepbind-$mode" \
                --output-file "$name" \
                -- \
                python -c "$code"
    done
}

profile_op topk_pybind '
import torch
import aiter

gate = torch.randn((128, 128), device="cuda", dtype=torch.float32)
weights = torch.empty((128, 8), device="cuda", dtype=torch.float32)
indices = torch.empty((128, 8), device="cuda", dtype=torch.int32)
tokens = torch.empty((128, 8), device="cuda", dtype=torch.int32)

aiter.topk_softmax(weights, indices, tokens, gate, False)
torch.cuda.synchronize()
'

profile_op topk_ctypes '
import torch
from aiter.ops.moe_op import topk_softmax_asm

gate = torch.randn((128, 128), device="cuda", dtype=torch.float32)
weights = torch.empty((128, 8), device="cuda", dtype=torch.float32)
indices = torch.empty((128, 8), device="cuda", dtype=torch.int32)
tokens = torch.empty((128, 8), device="cuda", dtype=torch.int32)

topk_softmax_asm(weights, indices, tokens, gate, False)
torch.cuda.synchronize()
'

python - "$TRACE_ROOT" <<'PY'
import csv
import sys
from pathlib import Path

root = Path(sys.argv[1])

for op in ("topk_pybind", "topk_ctypes"):
    results = {}
    for mode in (0, 1):
        path = (
            root
            / op
            / f"deepbind-{mode}"
            / f"{op}_hip_api_stats.csv"
        )
        with path.open() as f:
            results[mode] = {
                row["Name"]: int(row["Calls"])
                for row in csv.DictReader(f)
            }

    assert results[0] == results[1], {
        "deepbind_enabled": results[0],
        "deepbind_disabled": results[1],
    }
    print(f"PASS {op}: {len(results[0])} HIP APIs, counts match")
    for name in (
        "hipLaunchKernel",
        "hipModuleLaunchKernel",
        "hipGetDevicePropertiesR0600",
        "hipDeviceSynchronize",
    ):
        if name in results[0]:
            print(f"  {name}: {results[0][name]}")
PY

AITER_DISABLE_DEEPBIND=0 enables the protection, while 1 disables it.

The result was:

PASS topk_pybind: 18 HIP APIs, counts match
  hipLaunchKernel: 2
  hipGetDevicePropertiesR0600: 4
  hipDeviceSynchronize: 1

PASS topk_ctypes: 20 HIP APIs, counts match
  hipLaunchKernel: 1
  hipModuleLaunchKernel: 1
  hipGetDevicePropertiesR0600: 5
  hipDeviceSynchronize: 1

The complete API-name/count dictionaries matched in both comparisons; only durations varied between runs. Therefore, rocprofv3's callback-based HIP tracing is unaffected by RTLD_DEEPBIND, and no rocprofv3 warning is needed next to the ASan note. AITER_DISABLE_DEEPBIND=1 remains available for tools that intentionally depend on ELF symbol interposition.

[JohnQinAMD]

Nit: RTLD_NOW here is an unrelated behavior change.
ctypes.DEFAULT_MODE is 0, so this line goes lazy → eager independently of RTLD_DEEPBIND — undeclared, and the PR's other three sites keep their existing mode. Eager resolution turns a never-called dangling reference into a load-time undefined symbol. Suggest mode=os.RTLD_LAZY | _RTLD_DEEPBIND.

@zufayu Good catch — agreed. RTLD_NOW was an unrelated behavior change and could make an otherwise unused unresolved symbol fail eagerly at library load time.
Fixed in e6bc3ad by changing the mode to:

os.RTLD_LAZY | _RTLD_DEEPBIND

This preserves the previous lazy symbol-resolution behavior while adding only the intended RTLD_DEEPBIND isolation, consistent with the other loader sites.

[zufayu]

The fix is technically sound, but it changes dynamic-linking semantics for every AITER user, on by default, to work around a third-party bug. Two things follow.

1. Scope @valarLip . Is this a temporary shim until TileLang opt_unit_test #1867 lands, or a permanent defense? If temporary, it should carry a TODO and a removal path. If permanent, the default deserves a decision of its own — an AITER_ENABLE_DEEPBIND=1 opt-in, or auto-enabling only when an interposing stub is detected in the global scope, would keep the blast radius proportional to the problem.
2. One open technical question for
   Karatzas, Andreas . RTLD_DEEPBIND binds malloc/free inside AITER's .so to libc, while the rest of the process uses whatever is LD_PRELOADed — tcmalloc/jemalloc are common in serving deployments. If any AITER .so allocates memory that is freed across the module boundary, or vice versa, that is cross-allocator heap corruption, on by default, with symptoms that are hard to attribute. Profiler interposition and ASan are already handled (rocprofv3 A/B shows matching HIP API counts; ASan has the opt-out), but this one has not been answered.