From 171c51b786864f350f853292bc1e84f497b30ac7 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:13:37 +0800 Subject: [PATCH 1/2] [Fix] Discover flat CUDA includes for NVRTC --- .../test_tilelang_jit_nvrtc_include_paths.py | 56 +++++++++++++++++++ tilelang/jit/adapter/nvrtc/include_paths.py | 43 ++++++++++++++ tilelang/jit/adapter/nvrtc/libgen.py | 29 +++------- 3 files changed, 107 insertions(+), 21 deletions(-) create mode 100644 testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py create mode 100644 tilelang/jit/adapter/nvrtc/include_paths.py diff --git a/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py b/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py new file mode 100644 index 0000000000..3ae1f17037 --- /dev/null +++ b/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py @@ -0,0 +1,56 @@ +from pathlib import Path + +from tilelang.jit.adapter.nvrtc.include_paths import discover_cuda_include_paths + + +def _make_include_tree(cuda_home: Path, relative_paths: list[str]) -> list[str]: + paths = [] + for relative_path in relative_paths: + path = cuda_home / relative_path + path.mkdir(parents=True) + paths.append(str(path)) + return paths + + +def test_discovers_flat_pip_include_layout(tmp_path): + expected = _make_include_tree(tmp_path, ["include", "include/cccl"]) + + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == expected + + +def test_discovers_target_specific_system_include_layout(tmp_path): + expected = _make_include_tree( + tmp_path, + ["targets/x86_64-linux/include", "targets/x86_64-linux/include/cccl"], + ) + + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == expected + + +def test_discovers_flat_and_target_specific_include_layouts(tmp_path): + expected = _make_include_tree( + tmp_path, + [ + "include", + "include/cccl", + "targets/sbsa-linux/include", + "targets/sbsa-linux/include/cccl", + ], + ) + + assert discover_cuda_include_paths(str(tmp_path), machine="aarch64") == expected + + +def test_preserves_legacy_paths_when_cuda_layout_is_missing(tmp_path): + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == [ + str(tmp_path / "include"), + str(tmp_path / "targets/x86_64-linux/include"), + str(tmp_path / "targets/x86_64-linux/include/cccl"), + ] + + +def test_preserves_flat_windows_include_layout(tmp_path): + assert discover_cuda_include_paths(str(tmp_path), system="win32") == [ + str(tmp_path / "include"), + str(tmp_path / "include/cccl"), + ] diff --git a/tilelang/jit/adapter/nvrtc/include_paths.py b/tilelang/jit/adapter/nvrtc/include_paths.py new file mode 100644 index 0000000000..90ad7f0f72 --- /dev/null +++ b/tilelang/jit/adapter/nvrtc/include_paths.py @@ -0,0 +1,43 @@ +"""CUDA include-path discovery for NVRTC compilation.""" + +from __future__ import annotations + +import os.path as osp +import platform +import sys + + +def _target_include_path(cuda_home: str, machine: str | None = None) -> str: + machine = machine or platform.machine() + target_arch = "sbsa-linux" if machine in ("aarch64", "arm64") else "x86_64-linux" + return osp.join(cuda_home, "targets", target_arch, "include") + + +def discover_cuda_include_paths(cuda_home: str, machine: str | None = None, system: str | None = None) -> list[str]: + """Return NVRTC include paths supported by the CUDA layout at ``cuda_home``. + + CUDA pip packages may install headers in a flat ``include`` tree, while + system toolkits commonly use ``targets/-linux/include``. Include both + layouts when present so split or overlaid installations continue to work. + """ + flat_include = osp.join(cuda_home, "include") + system = system or sys.platform + if system.startswith("win32"): + return [flat_include, osp.join(flat_include, "cccl")] + + target_include = _target_include_path(cuda_home, machine) + include_paths = [] + + for include_path in (flat_include, target_include): + if osp.isdir(include_path): + include_paths.append(include_path) + cccl_include = osp.join(include_path, "cccl") + if osp.isdir(cccl_include): + include_paths.append(cccl_include) + + if include_paths: + return include_paths + + # Preserve the previous Linux search paths when probing an incomplete or + # not-yet-mounted toolkit. NVRTC will then report the missing headers. + return [flat_include, target_include, osp.join(target_include, "cccl")] diff --git a/tilelang/jit/adapter/nvrtc/libgen.py b/tilelang/jit/adapter/nvrtc/libgen.py index 032105221c..16d29ce2a6 100644 --- a/tilelang/jit/adapter/nvrtc/libgen.py +++ b/tilelang/jit/adapter/nvrtc/libgen.py @@ -18,8 +18,6 @@ import importlib import logging import os.path as osp -import platform -import sys import tempfile from types import ModuleType @@ -29,6 +27,7 @@ from tilelang.jit.adapter.libgen import LibraryGenerator from tilelang.jit.adapter.utils import is_cuda_target from tilelang.jit.adapter.nvrtc import is_nvrtc_available, NVRTC_UNAVAILABLE_MESSAGE +from tilelang.jit.adapter.nvrtc.include_paths import discover_cuda_include_paths logger = logging.getLogger(__name__) @@ -185,29 +184,13 @@ def compile_lib(self, timeout: float | None = None): cuda_home = CUDA_HOME if CUDA_HOME else "/usr/local/cuda" - # CUDA Toolkit include layout differs by platform: - # * Linux pip wheel ``nvidia-cuXX`` and system CUDA both expose - # per-target trees under ``targets/{arch}-linux/include``. - # * Windows pip wheel ``nvidia-cuXX`` and the system CUDA Toolkit - # ship a single flat ``include/`` directory — no ``targets/`` - # subtree exists. Hard-coding ``x86_64-linux`` there sends nvrtc - # to a non-existent path so headers like ``nvrtc_std.h`` fail to - # resolve. - cuda_include = osp.join(cuda_home, "include") - if sys.platform.startswith("win32"): - arch_include = cuda_include - else: - machine = platform.machine() - target_arch = "sbsa-linux" if machine in ("aarch64", "arm64") else "x86_64-linux" - arch_include = osp.join(cuda_home, "targets", target_arch, "include") + cuda_include_paths = discover_cuda_include_paths(cuda_home) __CUDACC_VER_MAJOR__ = get_nvrtc_version()[0] options = [ f"-I{tl_template_path}", f"-I{cutlass_path}", - f"-I{cuda_include}", - f"-I{arch_include}", - f"-I{arch_include}/cccl", + *(f"-I{include_path}" for include_path in cuda_include_paths), f"-D__CUDACC_VER_MAJOR__={__CUDACC_VER_MAJOR__}", ] @@ -220,7 +203,11 @@ def compile_lib(self, timeout: float | None = None): # forward declarations in ``cute/container/tuple.hpp`` under NVRTC # (cute uses variadic packs, cccl uses a single ``_Tp``). if __CUDACC_VER_MAJOR__ < 13: - options += [f"-I{arch_include}/cuda/std"] + options += [ + f"-I{include_path}/cuda/std" + for include_path in cuda_include_paths + if not include_path.endswith(osp.join("include", "cccl")) + ] if self.compile_flags: options += [item for flag in self.compile_flags for item in flag.split() if item not in options] From 998abf7f64db97db44df27222b2cb8831296d788 Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:26:23 +0800 Subject: [PATCH 2/2] [Test] Pin NVRTC layout cases to Linux --- .../python/jit/test_tilelang_jit_nvrtc_include_paths.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py b/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py index 3ae1f17037..1c19e7661c 100644 --- a/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py +++ b/testing/python/jit/test_tilelang_jit_nvrtc_include_paths.py @@ -15,7 +15,7 @@ def _make_include_tree(cuda_home: Path, relative_paths: list[str]) -> list[str]: def test_discovers_flat_pip_include_layout(tmp_path): expected = _make_include_tree(tmp_path, ["include", "include/cccl"]) - assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == expected + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == expected def test_discovers_target_specific_system_include_layout(tmp_path): @@ -24,7 +24,7 @@ def test_discovers_target_specific_system_include_layout(tmp_path): ["targets/x86_64-linux/include", "targets/x86_64-linux/include/cccl"], ) - assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == expected + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == expected def test_discovers_flat_and_target_specific_include_layouts(tmp_path): @@ -38,11 +38,11 @@ def test_discovers_flat_and_target_specific_include_layouts(tmp_path): ], ) - assert discover_cuda_include_paths(str(tmp_path), machine="aarch64") == expected + assert discover_cuda_include_paths(str(tmp_path), machine="aarch64", system="linux") == expected def test_preserves_legacy_paths_when_cuda_layout_is_missing(tmp_path): - assert discover_cuda_include_paths(str(tmp_path), machine="x86_64") == [ + assert discover_cuda_include_paths(str(tmp_path), machine="x86_64", system="linux") == [ str(tmp_path / "include"), str(tmp_path / "targets/x86_64-linux/include"), str(tmp_path / "targets/x86_64-linux/include/cccl"),