This is an automated email from the ASF dual-hosted git repository.
jinhongyii pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fb00a9bdb4 [Fix][Support] Use sbsa-linux CUDA include dir on ARM64
Linux (#20222)
fb00a9bdb4 is described below
commit fb00a9bdb4e1b439eaf1781ec487276bbdc9cf64
Author: Bohan Hou <[email protected]>
AuthorDate: Fri Aug 28 19:10:13 2026 -0400
[Fix][Support] Use sbsa-linux CUDA include dir on ARM64 Linux (#20222)
### Problem
`_compile_cuda_nvrtc` builds the architecture-specific CUDA include
directory as
`targets/{platform.machine()}-{platform.system().lower()}/include`, i.e.
`targets/aarch64-linux` on ARM64.
CUDA toolkits for ARM64 **servers** (SBSA — GH200/GB200 and similar)
install their headers under `targets/sbsa-linux`;
`targets/aarch64-linux` only exists in the embedded/L4T toolkits. On an
SBSA machine the computed path therefore does not exist, so both the
architecture-specific include directory and the CCCL include directory
derived from it are silently dropped from the NVRTC options. When the
toolkit has no top-level `include/` either, compilation fails with
"Cannot find CUDA headers".
### Fix
Add `_find_cuda_target_include()`, which probes `sbsa-linux` first on
ARM64 Linux (`aarch64`/`arm64`) and falls back to `aarch64-linux` for
embedded/L4T layouts. Other platforms keep the previous
`{machine}-{system}` behavior. The helper returns `None` when no
`targets/` layout is present, which is also a small cleanup of the CCCL
lookup (it no longer re-runs `os.path.isdir` on the same directory).
This mirrors the CUDA target selection used by other projects on SBSA
hosts, e.g. [Triton-distributed's
`build_scm.sh`](https://github.com/ByteDance-Seed/Triton-distributed/blob/8260bc34398c2b8f36dc840fd22f741ca9294584/build_scm.sh#L21).
### Testing
New `tests/python/support/test_nvcc.py` covers the SBSA, L4T, and x86_64
layouts plus the no-`targets/` case against a fake toolkit tree; it runs
on any host since the platform is monkeypatched.
```
python -m pytest tests/python/support/test_nvcc.py -q # 5 passed
pre-commit run --files python/tvm/support/nvcc.py
tests/python/support/test_nvcc.py
```
---
python/tvm/support/nvcc.py | 33 +++++++++++++++-----
tests/python/support/test_nvcc.py | 63 +++++++++++++++++++++++++++++++++++++++
2 files changed, 88 insertions(+), 8 deletions(-)
diff --git a/python/tvm/support/nvcc.py b/python/tvm/support/nvcc.py
index de83e70468..cdb602629c 100644
--- a/python/tvm/support/nvcc.py
+++ b/python/tvm/support/nvcc.py
@@ -292,6 +292,28 @@ def _compile_cuda_nvcc(
return data
+def _find_cuda_target_include(cuda_path):
+ """Find the architecture-specific ``targets/<triple>/include`` directory.
+
+ CUDA ships ARM64 server toolkits under ``targets/sbsa-linux``; only the
+ embedded/L4T toolkits use ``targets/aarch64-linux``, so probe both.
+
+ Returns None when no architecture-specific include directory exists.
+ """
+ machine = platform.machine()
+ system = platform.system().lower()
+ if system == "linux" and machine.lower() in ("aarch64", "arm64"):
+ triples = ["sbsa-linux", f"{machine}-{system}"]
+ else:
+ triples = [f"{machine}-{system}"]
+
+ for triple in triples:
+ include_dir = os.path.join(cuda_path, "targets", triple, "include")
+ if os.path.isdir(include_dir):
+ return include_dir
+ return None
+
+
def _compile_cuda_nvrtc(
code, target_format=None, arch=None, options=None, path_target=None,
use_nvshmem=False
):
@@ -481,18 +503,13 @@ namespace std {
include_paths.append(standard_include)
# Check architecture-specific include directory
- arch_include = os.path.join(
- cuda_path,
- "targets",
- f"{platform.machine()}-{platform.system().lower()}",
- "include",
- )
- if os.path.isdir(arch_include):
+ arch_include = _find_cuda_target_include(cuda_path)
+ if arch_include:
include_paths.append(arch_include)
# Check for CCCL include directory (required for cuda/std/cstdint and
type_traits)
# CCCL provides standard library functionality for device code
- cccl_include = os.path.join(arch_include, "cccl") if
os.path.isdir(arch_include) else None
+ cccl_include = os.path.join(arch_include, "cccl") if arch_include else None
if cccl_include and os.path.isdir(cccl_include):
include_paths.append(cccl_include)
diff --git a/tests/python/support/test_nvcc.py
b/tests/python/support/test_nvcc.py
new file mode 100644
index 0000000000..540e4b7162
--- /dev/null
+++ b/tests/python/support/test_nvcc.py
@@ -0,0 +1,63 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for functions in tvm/python/tvm/support/nvcc.py."""
+
+import os
+
+import pytest
+
+import tvm.testing
+from tvm.support import nvcc
+
+
+def _make_cuda_root(root, triples):
+ """Create a fake CUDA toolkit exposing the given ``targets/<triple>``
dirs."""
+ for triple in triples:
+ os.makedirs(os.path.join(root, "targets", triple, "include"))
+ return str(root)
+
+
[email protected](
+ "machine,available,expected",
+ [
+ # ARM64 server toolkits ship the headers under "sbsa-linux".
+ ("aarch64", ["sbsa-linux", "aarch64-linux"], "sbsa-linux"),
+ # Embedded/L4T toolkits only provide "aarch64-linux".
+ ("aarch64", ["aarch64-linux"], "aarch64-linux"),
+ ("arm64", ["sbsa-linux"], "sbsa-linux"),
+ ("x86_64", ["x86_64-linux"], "x86_64-linux"),
+ ],
+)
+def test_find_cuda_target_include(tmp_path, monkeypatch, machine, available,
expected):
+ """The architecture-specific include dir matches the installed toolkit
layout."""
+ monkeypatch.setattr(nvcc.platform, "machine", lambda: machine)
+ monkeypatch.setattr(nvcc.platform, "system", lambda: "Linux")
+ cuda_path = _make_cuda_root(tmp_path, available)
+ assert nvcc._find_cuda_target_include(cuda_path) == os.path.join(
+ cuda_path, "targets", expected, "include"
+ )
+
+
+def test_find_cuda_target_include_absent(tmp_path, monkeypatch):
+ """Toolkits without a ``targets/`` layout report no architecture-specific
dir."""
+ monkeypatch.setattr(nvcc.platform, "machine", lambda: "aarch64")
+ monkeypatch.setattr(nvcc.platform, "system", lambda: "Linux")
+ assert nvcc._find_cuda_target_include(str(tmp_path)) is None
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
