[ROCm] Defer `tilelang` import through its import `from vllm.tilelang_utils import tilelang` and relaxed `has_tilelang` by fxmarty-amd · Pull Request #51159 · vllm-project/vllm

fxmarty-amd · 2026-08-05T12:25:31Z

Purpose

This PR introduces vllm.tilelang_utils and prevents direct tilelang imports.

Avoid importing TileLang during ROCm module import, because importing TileLang can load wrongful/bugged TVM and HIP stub symbols into the global process scope before AITER loads its JIT modules.

This changes _tilelang_jit so ROCm applies tilelang.jit lazily on first kernel call, while non-ROCm keeps eager JIT decoration.

This PR is AI-assisted.

Proper fixes in followups need to be:

a PR in tilelang or TVM to fix the wrongful symbol load.
a PR in vLLM so that _has_module does not import modules [Feature]: vllm.utils.import_util._has_module should not import modules #51162

Test Plan

python tools/pre_commit/check_forbidden_imports.py --test-regex
pytest tests/kernels/test_mhc_tilelang_jit.py -v

Added new test to verify that CUDA applies tilelang.jit eagerly, ROCm defers tilelang.jit until first call, and the ROCm wrapper caches the compiled kernel.

VLLM_ROCM_USE_AITER=1 HIP_VISIBLE_DEVICES=1 pytest tests/quantization/test_online.py::test_online_quantization[True-fp8_per_tensor-None-Fp8PerTensorOnlineLinearMethod-Fp8PerTensorOnlineMoEMethod] -s -vvvvv

that is failing on main following #50879, and which surfaces the issue #51151

Test Result

test_mhc_tilelang_jit.py pass on MI350.

tests/quantization/test_online.py as well on MI300 (previous failing case)

python tools/pre_commit/check_forbidden_imports.py --test-regex passes as well.

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

AndreasKaratzas · 2026-08-05T12:34:24Z

/ci run

github-actions · 2026-08-05T12:34:40Z

✅ Triggered Buildkite CI #82487 for commit 467d3390f9b2.

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

Isotr0py · 2026-08-05T13:32:24Z

+def _tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
+    """Apply `tilelang.jit`, deferring until first call on ROCm.
+
+    ROCm defers JIT decoration so importing this module does not require
+    TileLang immediately. CUDA keeps the eager decoration behavior.
+    """
+    if not current_platform.is_rocm():
+        _ensure_tilelang_imported()
+        return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)
+
+    compiled_kernel: Callable[..., Any] | None = None
+
+    @functools.wraps(kernel_function)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        nonlocal compiled_kernel
+        if compiled_kernel is None:
+            _ensure_tilelang_imported()
+            compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
+                kernel_function
+            )
+        return compiled_kernel(*args, **kwargs)
+
+    return wrapper


How about making a tilelang_utils similar to triton_utils?

vllm/vllm/triton_utils/__init__.py

Lines 1 to 26 in a9b39d6

# SPDX-License-Identifier: Apache-2.0

# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from typing import TYPE_CHECKING

from vllm.triton_utils.importing import (

HAS_TRITON,

TritonLanguagePlaceholder,

TritonPlaceholder,

)

if TYPE_CHECKING or HAS_TRITON:

import triton

import triton.language as tl

import triton.language.extra.libdevice as tldevice

from triton.experimental import gluon

from triton.experimental.gluon import language as gl

from triton.language.core import _aggregate as aggregate # noqa: E501

else:

triton = TritonPlaceholder()

tl = TritonLanguagePlaceholder()

tldevice = TritonLanguagePlaceholder()

gluon = TritonLanguagePlaceholder()

gl = TritonLanguagePlaceholder()

aggregate = TritonLanguagePlaceholder()

from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor

Fixed in fa9cc9e, thank you!

Also made from vllm.tilelang_utils import tilelang import pattern compulsory, similar to triton.

Co-authored-by: Claude Signed-off-by: Felix Marty <Felix.Marty@amd.com>

AndreasKaratzas · 2026-08-07T16:30:27Z

/ci run

github-actions · 2026-08-07T16:30:49Z

✅ Triggered Buildkite CI #82879 for commit e0903f35e0b2.

AndreasKaratzas · 2026-08-07T21:25:59Z

/ci retry

Isotr0py

Anyway, let's land this to fix the broken model first.

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

…om libhip_stub.so Signed-off-by: Felix Marty <Felix.Marty@amd.com>

mergify · 2026-08-12T20:28:44Z

Documentation preview: https://vllm--51159.org.readthedocs.build/en/51159/

fxmarty-amd · 2026-08-12T20:29:14Z

@fxmarty-amd @rasmith can you try launching deepseek v4 as it is the only model that is relying on tilelang. Check the server logs for any traces of tilelang jit build to find out if we still have active path that is using tilelang.

@tjtanaa @Fangzhou-Ai actually, I do not quite get where self.use_fused_mhc -> _forward_fused_post_pre -> hc_pre/mhc_fused_post_pre -> torch.ops.vllm.mhc_fused_post_pre_tilelang is really used in #43950 - it looks to me the tilelang path is never exercised on tested MI300/MI325/MI350 code paths ?

This is guarded on:

vllm/vllm/models/deepseek_v4/amd/model.py

Lines 398 to 400 in 34735ac

    
           self.use_fused_mhc = HAS_TILELANG_MHC and not ( 
        
               HAS_AITER_MHC and self.hidden_size % 256 == 0 
        
           )

where

vllm/vllm/model_executor/layers/mhc.py

Line 29 in 34735ac

HAS_AITER_MHC = is_aiter_found_and_supported()

which simply checks if we are on CDNA > 2 (>=MI300).

Moreover:

deepseek-ai/DeepSeek-V4-Flash has hidden_size = 4096 => falls into aiter
deepseek-ai/DeepSeek-V4-Pro has hidden_size = 7168 => falls into aiter

@tjtanaa this looks safe to me after ensuring the vLLM CI / AMD CI runs on it. But from what I found there is no registered DS-v4 test that exercises a tilelang path.

Fangzhou-Ai · 2026-08-12T21:55:58Z

@fxmarty-amd In realistic we never use the tilelang mHC, AITER mHC is always preferred.

tjtanaa · 2026-08-13T14:19:04Z

/ci run

github-actions · 2026-08-13T14:19:09Z

✅ @fxmarty-amd, CI is now available for this PR.

/ci run starts a CI build.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/ci cancel cancels scheduled or running CI builds for this PR branch.

github-actions · 2026-08-13T14:19:22Z

✅ Triggered Buildkite CI #83743 for commit 97f38dffc690.

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd · 2026-08-13T15:08:57Z

/ci run

github-actions · 2026-08-13T15:09:15Z

✅ Triggered Buildkite CI #83748 for commit a9a188702c7c.

fxmarty-amd · 2026-08-13T16:45:26Z

buildkite/amd-ci looks sane to me (although it has some skipped test groups)

…_utils import tilelang` and relaxed `has_tilelang` (vllm-project#51159) Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>

…_utils import tilelang` and relaxed `has_tilelang` (vllm-project#51159) Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

defer tilelang import on rocm

bcfd259

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

claude Bot reviewed Aug 5, 2026

View reviewed changes

mergify Bot added the rocm Related to AMD ROCm label Aug 5, 2026

github-project-automation Bot added this to AMD Aug 5, 2026

github-project-automation Bot moved this to Todo in AMD Aug 5, 2026

add test

467d339

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd requested review from AndreasKaratzas, WoosukKwon, mgoin, tlrmchlsmth, yewentao256 and zyongye as code owners August 5, 2026 12:26

fxmarty-amd added 2 commits August 5, 2026 07:38

linting

341f234

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

mypy

073c6af

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

Isotr0py reviewed Aug 5, 2026

View reviewed changes

This was referenced Aug 5, 2026

[ROCm][Bugfix] Restore guarded eager Mori EP import #51110

Closed

fix(jit): isolate AITER extensions from HIP interposers ROCm/aiter#4566

Open

[JIT] AITER extensions can bind HIP symbols from an earlier global TileLang stub ROCm/aiter#4585

Open

lazypool mentioned this pull request Aug 6, 2026

[Feature]: vllm.utils.import_util._has_module should not import modules #51162

Open

1 task done

Extract tilelang_utils module for TileLang import/JIT guards

fa9cc9e

Co-authored-by: Claude Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd requested a review from hmellor as a code owner August 6, 2026 10:11

fxmarty-amd changed the title ~~[ROCm] Defer tilelang import~~ [ROCm] Defer tilelang import through its import from vllm.tilelang_utils import tilelang Aug 6, 2026

Keep ENABLE_PDL in tilelang_kernels.py

69bdfc9

Co-authored-by: Claude Signed-off-by: Felix Marty <Felix.Marty@amd.com>

fxmarty-amd changed the title ~~[ROCm] Defer tilelang import through its import from vllm.tilelang_utils import tilelang~~ [ROCm] Defer tilelang import through its import from vllm.tilelang_utils import tilelang and relaxed has_tilelang Aug 6, 2026

Merge branch 'main' into tilelang-lazy-import

e0903f3

Isotr0py approved these changes Aug 12, 2026

View reviewed changes

fxmarty-amd added 2 commits August 12, 2026 14:49

reset dockerfile-stages-dependency.png

64fb5c3

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

add test verifying tilelang is not imported, nor symbols populated fr…

712fadd

…om libhip_stub.so Signed-off-by: Felix Marty <Felix.Marty@amd.com>

tjtanaa approved these changes Aug 13, 2026

View reviewed changes

tjtanaa added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 13, 2026

Merge branch 'main' into tilelang-lazy-import

97f38df

skip jit monitor tilelang tests

a9a1887

Signed-off-by: Felix Marty <Felix.Marty@amd.com>

AndreasKaratzas approved these changes Aug 13, 2026

View reviewed changes

AndreasKaratzas merged commit c5b7c06 into vllm-project:main Aug 13, 2026
90 checks passed

github-project-automation Bot moved this from Todo to Done in AMD Aug 13, 2026

Suppressor72 mentioned this pull request Aug 24, 2026

[Performance]: MTP first repeat misses prefix cache on a hybrid Mamba/GDN model #53504

Closed

migarci2 mentioned this pull request Aug 27, 2026

[Core] Make optional module discovery side-effect-free #54100

Open

simondanielsson mentioned this pull request Sep 16, 2026

[Bugfix][ROCm][GLM-5.3-Flash] Apply deferred tilelang.jit already on attribute access #57192

Merged

4 tasks

AndreasKaratzas mentioned this pull request Oct 18, 2026

[CI][ROCm] Run ROCm-capable tests that are currently skipped on AMD #62298

Open

2 of 6 tasks

This was referenced Oct 28, 2026

adapt(amd): AMD ROCm (W7900D/gfx1100) backend on vLLM 0.32 empty physics31415926/vllm-plugin-FL#5

Closed

adapt(amd): AMD ROCm (W7900D/gfx1100) backend on vLLM 0.32 flagos-ai/vllm-plugin-FL#659

Draft

	# SPDX-License-Identifier: Apache-2.0
	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
	from typing import TYPE_CHECKING

	from vllm.triton_utils.importing import (
	HAS_TRITON,
	TritonLanguagePlaceholder,
	TritonPlaceholder,
	)

	if TYPE_CHECKING or HAS_TRITON:
	import triton
	import triton.language as tl
	import triton.language.extra.libdevice as tldevice
	from triton.experimental import gluon
	from triton.experimental.gluon import language as gl
	from triton.language.core import _aggregate as aggregate # noqa: E501
	else:
	triton = TritonPlaceholder()
	tl = TritonLanguagePlaceholder()
	tldevice = TritonLanguagePlaceholder()
	gluon = TritonLanguagePlaceholder()
	gl = TritonLanguagePlaceholder()
	aggregate = TritonLanguagePlaceholder()

	from vllm.triton_utils.tensor_descriptor import use_tensor_descriptor

Conversation

fxmarty-amd commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

AndreasKaratzas commented Aug 5, 2026

Uh oh!

github-actions Bot commented Aug 5, 2026

Uh oh!

Isotr0py Aug 5, 2026

Choose a reason for hiding this comment

Uh oh!

fxmarty-amd Aug 6, 2026

Choose a reason for hiding this comment

Uh oh!

AndreasKaratzas commented Aug 7, 2026

Uh oh!

github-actions Bot commented Aug 7, 2026

Uh oh!

AndreasKaratzas commented Aug 7, 2026

Uh oh!

Isotr0py left a comment

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Aug 12, 2026

Uh oh!

fxmarty-amd commented Aug 12, 2026

Uh oh!

Fangzhou-Ai commented Aug 12, 2026

Uh oh!

tjtanaa commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

fxmarty-amd commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

fxmarty-amd commented Aug 13, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

fxmarty-amd commented Aug 5, 2026 •

edited

Loading