This is an automated email from the ASF dual-hosted git repository.
tqchen 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 9b574536b3 [BUGFIX] Fix runtime correctness and backend SDK
dependencies (#20715)
9b574536b3 is described below
commit 9b574536b3d03f1ccb611bc0786403834ab533c6
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 14 08:19:00 2026 -0400
[BUGFIX] Fix runtime correctness and backend SDK dependencies (#20715)
Use the resolved address family for Disco controller and client sockets
after socket initialization, preserving IPv4 and IPv6 connections.
Initialize masked CPU attention state and read its one-element rescaling
buffer at the scalar index. Seed the existing CUDA argsort expected
module with device metadata before allocation inference.
Keep backend SDK libraries private to their runtime implementations and
explicitly link required SDKs on direct consumers, including CUDA, HIP,
and CLML runtime components.
---
CMakeLists.txt | 5 +++++
cmake/modules/CUDA.cmake | 2 +-
cmake/modules/Metal.cmake | 2 +-
cmake/modules/OpenCL.cmake | 7 ++++++-
cmake/modules/ROCM.cmake | 2 +-
cmake/modules/Vulkan.cmake | 2 +-
cmake/modules/contrib/CLML.cmake | 5 -----
python/tvm/relax/frontend/nn/llm/_prefill_kernels.py | 3 ++-
src/runtime/extra/disco/distributed/socket_session.cc | 11 ++++++-----
tests/python/relax/test_backend_dispatch_sort_scan.py | 2 +-
10 files changed, 24 insertions(+), 17 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 7a8663ffb2..d7ed0fef4e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -546,6 +546,11 @@ target_link_libraries(tvm_runtime_extra PUBLIC tvm_runtime)
# If disco/cuda_ipc is included, link the CUDA DSO.
if(USE_CUDA)
target_link_libraries(tvm_runtime_extra PUBLIC tvm_runtime_cuda)
+ target_link_libraries(tvm_runtime_extra PRIVATE ${CUDA_CUDART_LIBRARY}
${CUDA_CUDA_LIBRARY})
+endif()
+
+if(USE_ROCM)
+ target_link_libraries(tvm_runtime_extra PRIVATE ${ROCM_HIPHCC_LIBRARY})
endif()
# CUTLASS fpA_intB_gemm and flash_attn are separate shared libs.
diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake
index e1731e8078..4e2cdd508c 100644
--- a/cmake/modules/CUDA.cmake
+++ b/cmake/modules/CUDA.cmake
@@ -84,7 +84,7 @@ if(USE_CUDA)
endif()
add_library(tvm_runtime_cuda SHARED $<TARGET_OBJECTS:tvm_runtime_cuda_objs>)
list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_cuda)
- target_link_libraries(tvm_runtime_cuda PUBLIC tvm_runtime
${CUDA_CUDART_LIBRARY} ${CUDA_CUDA_LIBRARY})
+ target_link_libraries(tvm_runtime_cuda PUBLIC tvm_runtime PRIVATE
${CUDA_CUDART_LIBRARY} ${CUDA_CUDA_LIBRARY})
tvm_configure_target_library(tvm_runtime_cuda RUNTIME_MODULE)
if(USE_NVTX)
diff --git a/cmake/modules/Metal.cmake b/cmake/modules/Metal.cmake
index 02526b4c9c..247f4fc5d1 100644
--- a/cmake/modules/Metal.cmake
+++ b/cmake/modules/Metal.cmake
@@ -29,7 +29,7 @@ if(USE_METAL)
endif()
add_library(tvm_runtime_metal SHARED
$<TARGET_OBJECTS:tvm_runtime_metal_objs>)
list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_metal)
- target_link_libraries(tvm_runtime_metal PUBLIC tvm_runtime ${METAL_LIB}
${FOUNDATION_LIB})
+ target_link_libraries(tvm_runtime_metal PUBLIC tvm_runtime PRIVATE
${METAL_LIB} ${FOUNDATION_LIB})
tvm_configure_target_library(tvm_runtime_metal RUNTIME_MODULE)
endif(USE_METAL)
# When USE_METAL=OFF the codegen-side fallback in
diff --git a/cmake/modules/OpenCL.cmake b/cmake/modules/OpenCL.cmake
index 392457da7a..cabfc3173f 100644
--- a/cmake/modules/OpenCL.cmake
+++ b/cmake/modules/OpenCL.cmake
@@ -45,9 +45,14 @@ if(USE_OPENCL)
endif()
add_library(tvm_runtime_opencl SHARED
$<TARGET_OBJECTS:tvm_runtime_opencl_objs>)
list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_opencl)
- target_link_libraries(tvm_runtime_opencl PUBLIC tvm_runtime ${_opencl_libs})
+ target_link_libraries(tvm_runtime_opencl PUBLIC tvm_runtime PRIVATE
${_opencl_libs})
tvm_configure_target_library(tvm_runtime_opencl RUNTIME_MODULE)
+ if(USE_CLML OR USE_CLML_GRAPH_EXECUTOR)
+ # CLML calls both OpenCL SDK and TVM OpenCL runtime APIs directly.
+ target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_opencl
${_opencl_libs})
+ endif()
+
if(USE_OPENCL_ENABLE_HOST_PTR)
add_definitions(-DOPENCL_ENABLE_HOST_PTR)
endif(USE_OPENCL_ENABLE_HOST_PTR)
diff --git a/cmake/modules/ROCM.cmake b/cmake/modules/ROCM.cmake
index ce9ad1414f..33abc75a2b 100644
--- a/cmake/modules/ROCM.cmake
+++ b/cmake/modules/ROCM.cmake
@@ -47,7 +47,7 @@ if(USE_ROCM)
endif()
add_library(tvm_runtime_rocm SHARED $<TARGET_OBJECTS:tvm_runtime_rocm_objs>)
list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_rocm)
- target_link_libraries(tvm_runtime_rocm PUBLIC tvm_runtime ${_rocm_libs})
+ target_link_libraries(tvm_runtime_rocm PUBLIC tvm_runtime PRIVATE
${_rocm_libs})
tvm_configure_target_library(tvm_runtime_rocm RUNTIME_MODULE)
endif(USE_ROCM)
diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake
index 2f34b63240..34d932b690 100644
--- a/cmake/modules/Vulkan.cmake
+++ b/cmake/modules/Vulkan.cmake
@@ -54,6 +54,6 @@ if(USE_VULKAN)
endif()
add_library(tvm_runtime_vulkan SHARED
$<TARGET_OBJECTS:tvm_runtime_vulkan_objs>)
list(APPEND TVM_RUNTIME_BACKEND_LIBS tvm_runtime_vulkan)
- target_link_libraries(tvm_runtime_vulkan PUBLIC tvm_runtime
${Vulkan_LIBRARY})
+ target_link_libraries(tvm_runtime_vulkan PUBLIC tvm_runtime PRIVATE
${Vulkan_LIBRARY})
tvm_configure_target_library(tvm_runtime_vulkan RUNTIME_MODULE)
endif(USE_VULKAN)
diff --git a/cmake/modules/contrib/CLML.cmake b/cmake/modules/contrib/CLML.cmake
index d3e80a2e84..dbe383159c 100644
--- a/cmake/modules/contrib/CLML.cmake
+++ b/cmake/modules/contrib/CLML.cmake
@@ -72,12 +72,7 @@ if(USE_CLML_GRAPH_EXECUTOR)
endif()
add_library(tvm_clml_objs OBJECT ${CLML_CONTRIB_SRC})
target_link_libraries(tvm_clml_objs PRIVATE tvm_runtime_extra_defs)
- # CLML depends on OpenCL runtime symbols — link the OpenCL DSO instead of
- # duplicating sources (which would cause duplicate registrations).
target_link_libraries(tvm_runtime_extra PRIVATE tvm_clml_objs
${EXTERN_CLML_COMPUTE_LIB})
- if(TARGET tvm_runtime_opencl)
- target_link_libraries(tvm_runtime_extra PRIVATE tvm_runtime_opencl)
- endif()
message(STATUS "Build with CLML graph runtime support: "
${EXTERN_CLML_COMPUTE_LIB})
diff --git a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
index dc5e77576a..4d31ee05c8 100644
--- a/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
+++ b/python/tvm/relax/frontend/nn/llm/_prefill_kernels.py
@@ -175,6 +175,7 @@ def _attention_prefill_cpu(
S_val[0] *= sm_scale * math.log2(math.exp(1))
# update m_val, d_val , O_local
+ new_m[0] = m_val[0]
if _causal_or_sliding_cross_mask(
causal,
row=q_idx,
@@ -195,7 +196,7 @@ def _attention_prefill_cpu(
m_val[0] = new_m[0]
factor[0] = T.exp2(S_val[0] - m_val[0])
for d_idx in T.serial(d):
- O_local[d_idx] = O_local[d_idx] *
scale_O[d_idx]
+ O_local[d_idx] = O_local[d_idx] *
scale_O[0]
for d_idx in T.serial(d):
O_local[d_idx] += V_local[d_idx] *
factor[0]
diff --git a/src/runtime/extra/disco/distributed/socket_session.cc
b/src/runtime/extra/disco/distributed/socket_session.cc
index 803119221d..0b67c19192 100644
--- a/src/runtime/extra/disco/distributed/socket_session.cc
+++ b/src/runtime/extra/disco/distributed/socket_session.cc
@@ -70,9 +70,10 @@ class SocketSessionObj : public BcastSessionObj {
num_workers_per_node_);
Socket::Startup();
- socket_.Create();
+ SockAddr addr{host.c_str(), port};
+ socket_.Create(addr.ss_family());
socket_.SetKeepAlive(true);
- socket_.Bind(SockAddr(host.c_str(), port));
+ socket_.Bind(addr);
socket_.Listen();
LOG(INFO) << "SocketSession controller listening on " << host << ":" <<
port;
@@ -214,10 +215,10 @@ class RemoteSocketSession {
public:
explicit RemoteSocketSession(const ffi::String& server_host, int server_port,
int num_local_workers) {
- socket_.Create();
- socket_.SetKeepAlive(true);
- SockAddr server_addr{server_host.c_str(), server_port};
Socket::Startup();
+ SockAddr server_addr{server_host.c_str(), server_port};
+ socket_.Create(server_addr.ss_family());
+ socket_.SetKeepAlive(true);
if (!socket_.Connect(server_addr)) {
TVM_FFI_THROW(InternalError) << "Failed to connect to server " <<
server_addr.AsString()
<< ", errno = " <<
Socket::GetLastErrorCode();
diff --git a/tests/python/relax/test_backend_dispatch_sort_scan.py
b/tests/python/relax/test_backend_dispatch_sort_scan.py
index f457ca554c..c70de46a14 100644
--- a/tests/python/relax/test_backend_dispatch_sort_scan.py
+++ b/tests/python/relax/test_backend_dispatch_sort_scan.py
@@ -277,7 +277,7 @@ def test_dispatch_argsort_cuda():
vdevices = [R.vdevice("cuda", 0)]
x = relax.Var("x", R.Tensor((2, 3), "float32", vdevices[0]))
y = relax.Var("y", R.Tensor((2, 3), "float32"))
- bb = relax.BlockBuilder()
+ bb = relax.BlockBuilder(tvm.IRModule(global_infos={"vdevice": vdevices}))
with target:
with bb.function("foo", (x,), {"global_symbol": "foo"}):
with bb.dataflow():
