feat(wheel): self-contained wheels that run on a stock torch+cpu

CMakeLists.txt

@@ -16,6 +16,15 @@ cmake_minimum_required(VERSION 3.18 FATAL_ERROR)
 
 set(ACCELERATOR "cuda" CACHE STRING "Accelerator platform: cuda, metax, ascend, tsingmicro, dcu, gcu, or musa")
 
+# Directory inside the wheel holding a bundled forked libtorch, when the backend
+# ships one (see scripts/bundle_*_libtorch.sh). "lib" means "no separate bundle
+# dir": the CUDA backend drops its extra .so straight into torch_fl/lib/.
+# Set per backend below; consumed by cmake/FlagosRpath.cmake.
+set(FLAGOS_BUNDLE_LIBDIR "lib")
+# Vendor driver/runtime dirs to embed in RUNPATH. Never bundled into the wheel
+# (a machine with the card has the driver), mirroring MetaX's /opt/maca.
+set(FLAGOS_VENDOR_RPATH_DIRS "")
+
 # Kernel build options (can be overridden via -D flags from setup.py)
 option(CUDA_KERNEL "Build CUDA kernel implementations" ON)
 option(FLAGGEMS_KERNEL "Build FlagGems C++ kernel wrappers (requires liboperators.so)" ON)
@@ -54,6 +63,9 @@ elseif(ACCELERATOR STREQUAL "metax")
   elseif(NOT FLAGGEMS_KERNEL)
     set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE)
   endif()
+  # Self-contained wheel: forked libtorch bundled by
+  # scripts/bundle_maca_libtorch.sh (the maca runtime stays on the target).
+  set(FLAGOS_BUNDLE_LIBDIR "lib_maca")
   project(TORCH_FLAGOS CXX C)
 elseif(ACCELERATOR STREQUAL "tsingmicro")
   project(TORCH_FLAGOS CXX C)
@@ -67,6 +79,10 @@ elseif(ACCELERATOR STREQUAL "dcu")
   # No nvcc and no CUDA language: nothing outside backends/metax/ is a .cu.
   set(CUDA_KERNEL OFF CACHE BOOL "Build CUDA kernel implementations" FORCE)
   set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE)
+  # Self-contained wheel: DTK's forked libtorch (libtorch_cpu.so carries hip
+  # symbols and needs libgalaxyhip.so.5) is bundled here by
+  # scripts/bundle_dcu_libtorch.sh. The DTK driver itself stays on the target.
+  set(FLAGOS_BUNDLE_LIBDIR "lib_dcu")
   project(TORCH_FLAGOS CXX C)
 elseif(ACCELERATOR STREQUAL "gcu")
   project(TORCH_FLAGOS CXX C)
@@ -80,6 +96,15 @@ elseif(ACCELERATOR STREQUAL "musa")
   set(FLAGGEMS_KERNEL OFF CACHE BOOL "Build FlagGems kernel implementations" FORCE)
   project(TORCH_FLAGOS CXX C)
 else()
+  # PPU rides this branch: it has no accelerator value of its own because it is
+  # built against PPU_SDK/CUDA_SDK, so the CUDA boxing kernels apply unchanged
+  # (see the README's PPU section). It still needs its own bundle dir, since that
+  # libtorch is a local PPU build and not the upstream one. Detected by the SDK
+  # env vars, matching setup.py::_vendor_supplies_triton().
+  if(DEFINED ENV{PPU_SDK} OR DEFINED ENV{PPU_HOME})
+    set(FLAGOS_BUNDLE_LIBDIR "lib_ppu")
+    message(STATUS "PPU SDK detected: bundling libtorch into torch_fl/lib_ppu")
+  endif()
   # project(... CUDA) runs nvcc compiler-id before find_package(CUDAToolkit); conda
   # layouts need CUDAToolkit_ROOT (headers under targets/x86_64-linux/include).
   if(NOT CUDAToolkit_ROOT)
@@ -222,7 +247,23 @@ elseif(ACCELERATOR STREQUAL "dcu")
   set(CUDA_RUNTIME_LIB cuda_runtime_compat)
 
   # Embed the compat + DTK library paths so LD_LIBRARY_PATH is not needed.
-  list(APPEND CMAKE_INSTALL_RPATH "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib")
+  # Measured layout of the deps the forked libtorch actually pulls. They are all
+  # on the DTK container's LD_LIBRARY_PATH, which a clean env does NOT inherit --
+  # so the RUNPATH has to name them or a self-contained wheel gets "not found":
+  #   cuda/cuda-*/lib64      libcudart.so.12 (shim over libgalaxyhip)
+  #   lib                    libgalaxyhip.so.5 libhipnn librocfft.so.0
+  #                          librocrand.so.1 librocsparse.so.1 libunwind.so.8
+  #                          libMIOpen-recommend.so
+  #   hip/lib, lib64         hip runtime
+  #   aillvm/lib             libomp.so
+  #   .hyhal/rocm_smi/lib    librocm_smi64.so.2
+  # Keep in sync with VENDOR_RPATH in scripts/bundle_dcu_libtorch.sh.
+  set(_dcu_rpath_dirs
+      "${DCU_CUDA_ROOT}/lib64" "${DTK_ROOT}/lib" "${DTK_ROOT}/hip/lib"
+      "${DTK_ROOT}/lib64" "${DTK_ROOT}/aillvm/lib" "${DTK_ROOT}/llvm/lib"
+      "${DTK_ROOT}/.hyhal/rocm_smi/lib" "/opt/hyhal/lib")
+  list(APPEND CMAKE_INSTALL_RPATH ${_dcu_rpath_dirs})
+  set(FLAGOS_VENDOR_RPATH_DIRS ${_dcu_rpath_dirs})
 
   # Pre-create torch::cudart so PyTorch's cuda.cmake is skipped entirely
   # (guard: if(TARGET torch::cudart) return()).
@@ -358,6 +399,12 @@ elseif(ACCELERATOR STREQUAL "metax")
   endif()
   message(STATUS "MetaX SDK path: ${METAX_PATH}")
 
+  # The maca runtime (mcblas/mcdnn/...) comes from the target's /opt/maca driver
+  # install; it is never bundled. /opt/maca is kept as an explicit fallback in
+  # case the build machine's METAX_PATH is a non-default location.
+  set(FLAGOS_VENDOR_RPATH_DIRS
+      "${METAX_PATH}/lib" "${METAX_PATH}/lib64" "/opt/maca/lib" "/opt/maca/lib64")
+
   add_library(cuda_runtime_compat INTERFACE IMPORTED)
   set_target_properties(cuda_runtime_compat PROPERTIES
     INTERFACE_INCLUDE_DIRECTORIES "${METAX_PATH}/tools/cu-bridge/include;${METAX_PATH}/include;${METAX_PATH}/include/mcr"
@@ -513,11 +560,15 @@ else()
 endif()
 
 if(PYTORCH_INSTALL_DIR)
+  # Needed for the *build* link step (and for an in-place/editable install, where
+  # the built .so stay in the source tree). flagos_set_portable_rpath() strips it
+  # from the installed targets' RUNPATH so the wheel is relocatable.
   list(APPEND CMAKE_INSTALL_RPATH "${PYTORCH_INSTALL_DIR}/lib")
 endif()
 
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 include(${PROJECT_SOURCE_DIR}/cmake/TorchPythonTargets.cmake)
+include(${PROJECT_SOURCE_DIR}/cmake/FlagosRpath.cmake)
 
 add_subdirectory(${PROJECT_SOURCE_DIR}/csrc/runtime/accelerator)
 add_subdirectory(${PROJECT_SOURCE_DIR}/csrc)

cmake/FlagosRpath.cmake

@@ -0,0 +1,76 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed 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.
+
+# Portable RUNPATH for the installed native libs.
+#
+# The problem: CMakeLists.txt appends "${PYTORCH_INSTALL_DIR}/lib" to
+# CMAKE_INSTALL_RPATH so an in-place build finds the active torch wheel. That is
+# the *build machine's* interpreter path (e.g.
+# /nfs/.../envs/torch-fl-210/lib/python3.12/site-packages/torch/lib), which does
+# not exist on a target machine -- and worse, on a machine that does have that
+# path it silently wins over the interpreter actually running, so a py3.12 env
+# can load another env's libtorch.
+#
+# Dropping it is safe: libtorch_fl.so is only ever loaded through
+# `import torch_fl` -> torch_fl._C -> libtorch_bindings.so, which happens after
+# torch_fl/__init__.py has already run `import torch`. libc10 / libtorch_cpu are
+# mapped by then, so the loader satisfies those DT_NEEDED entries by soname from
+# the already-loaded set. This is why the MetaX self-contained wheel works today.
+#
+# Anything else the backend put on CMAKE_INSTALL_RPATH (vendor driver dirs, the
+# FlagGems liboperators dir) is preserved, so in-place builds keep working.
+#
+# flagos_set_portable_rpath(<target> [EXTRA_DIRS <dir>...])
+function(flagos_set_portable_rpath _target)
+  cmake_parse_arguments(_ARG "" "" "EXTRA_DIRS" ${ARGN})
+
+  if(NOT UNIX OR APPLE)
+    return()
+  endif()
+  if(NOT TARGET ${_target})
+    return()
+  endif()
+
+  # $ORIGIN is torch_fl/lib/ for every installed target. A sibling bundle dir
+  # (torch_fl/lib_maca, lib_dcu, lib_ppu) is reached via $ORIGIN/../<name>.
+  set(_rpath "$ORIGIN" "$ORIGIN/lib")
+  if(FLAGOS_BUNDLE_LIBDIR AND NOT FLAGOS_BUNDLE_LIBDIR STREQUAL "lib")
+    list(APPEND _rpath "$ORIGIN/../${FLAGOS_BUNDLE_LIBDIR}")
+  endif()
+
+  foreach(_dir IN LISTS _ARG_EXTRA_DIRS FLAGOS_VENDOR_RPATH_DIRS)
+    if(_dir)
+      list(APPEND _rpath "${_dir}")
+    endif()
+  endforeach()
+
+  # Inherit whatever the backend branches appended, minus the build machine's
+  # torch/lib and minus the $ORIGIN entries already placed above.
+  foreach(_dir IN LISTS CMAKE_INSTALL_RPATH)
+    if(NOT _dir MATCHES "^\\$ORIGIN"
+       AND NOT (PYTORCH_INSTALL_DIR AND _dir STREQUAL "${PYTORCH_INSTALL_DIR}/lib"))
+      list(APPEND _rpath "${_dir}")
+    endif()
+  endforeach()
+
+  list(REMOVE_DUPLICATES _rpath)
+  string(REPLACE ";" ":" _rpath_str "${_rpath}")
+
+  set_target_properties(${_target} PROPERTIES
+    INSTALL_RPATH "${_rpath_str}"
+    # Stop CMake appending the imported torch target's build-machine link dir.
+    INSTALL_RPATH_USE_LINK_PATH OFF
+    BUILD_WITH_INSTALL_RPATH ON
+  )
+endfunction()

csrc/CMakeLists.txt

@@ -325,29 +325,18 @@ if(ACCELERATOR STREQUAL "metax")
     message(WARNING "mcblas not found, MetaX mm will fallback to naive kernel")
   endif()
 
-  # Self-contained wheel: the installed libtorch_fl.so must NOT bake the build
-  # machine's absolute torch/lib path into its RUNPATH (that path does not exist
-  # on target machines). Resolve the forked libtorch bundled under
-  # torch_fl/lib_maca/ via $ORIGIN, and the maca runtime (mcblas/mcdnn/...) from
-  # the target's /opt/maca driver install. Disabling INSTALL_RPATH_USE_LINK_PATH
-  # stops CMake from appending the imported torch target's build-machine link
-  # dir. Mirrors scripts/bundle_maca_libtorch.sh's patchelf pass.
-  set(_metax_rpath
-      "$ORIGIN:$ORIGIN/../lib_maca:${METAX_PATH}/lib:${METAX_PATH}/lib64:/opt/maca/lib:/opt/maca/lib64")
-  # Setting INSTALL_RPATH explicitly discards the entries the top-level
-  # CMakeLists appended to CMAKE_INSTALL_RPATH, including FlagGems'
-  # liboperators.so dir. Re-add it so the C++ FlagGems path (FLAGGEMS_KERNEL,
-  # kFlagOs) resolves liboperators.so without LD_LIBRARY_PATH.
-  if(FLAGGEMS_KERNEL AND _flaggems_libdir)
-    set(_metax_rpath "${_metax_rpath}:${_flaggems_libdir}")
-  endif()
-  set_target_properties(${LIBRARY_NAME} PROPERTIES
-    INSTALL_RPATH "${_metax_rpath}"
-    INSTALL_RPATH_USE_LINK_PATH OFF
-    BUILD_WITH_INSTALL_RPATH ON
-  )
 endif()
 
+# Self-contained wheel: the installed libtorch_fl.so must NOT bake the build
+# machine's absolute torch/lib path into its RUNPATH -- that path does not exist
+# on a target machine, and on one where it does exist it silently overrides the
+# interpreter actually running. Resolve a bundled forked libtorch via $ORIGIN
+# instead. Everything else the backend put on CMAKE_INSTALL_RPATH is kept,
+# including FlagGems' liboperators.so dir, so the C++ FlagGems path
+# (FLAGGEMS_KERNEL, kFlagOs) still resolves without LD_LIBRARY_PATH. Mirrors the
+# patchelf pass in scripts/bundle_*_libtorch.sh.
+flagos_set_portable_rpath(${LIBRARY_NAME})
+
 install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}

csrc/aten/backends/flagos/python_op_caller.cc

@@ -610,6 +610,15 @@ at::Generator GetFlagosDefaultCudaGenerator(int64_t device_index) {
   }
   py::gil_scoped_acquire gil;
   py::module_ torch_cuda = py::module_::import("torch.cuda");
+  // torch.cuda.default_generators is populated by CUDA lazy-init, not at import.
+  // A self-contained wheel front-ends a stock torch+cpu whose CUDA state nothing
+  // else touches, so the tuple is still empty here and indexing it would raise
+  // IndexError -- which surfaces as a failure of every device-side RNG op
+  // (randn/rand/normal_) while plain factories (empty/zeros/ones) work fine.
+  // Force lazy-init first; it is idempotent and cheap once initialized.
+  if (py::len(torch_cuda.attr("default_generators")) == 0) {
+    torch_cuda.attr("init")();
+  }
   py::object gens = torch_cuda.attr("default_generators");
   py::object py_gen = gens[py::cast(device_index)];
   // torch.Generator -> at::Generator via THPGenerator unpack.

csrc/aten/copy_ops.cc

@@ -26,8 +26,16 @@
 
 // On the CUDA-family backends (including MetaX boxing) the flagos device shares
 // the vendor's CUDA streams, so the current stream is readable from c10::cuda.
+//
+// USE_DCU is excluded for the same reason runtime/guard.h:22 excludes it: the
+// DCU wheel is hipified, so <c10/cuda/CUDAStream.h> resolves through DTK's
+// CUDA-compat shim against USE_ROCM-hipified torch headers and expands to hip*
+// symbols the shim never declares (`'hipStreamCaptureStatus' was not declared`),
+// and c10::cuda::getCurrentCUDAStream would not link there anyway -- DTK exports
+// c10::hip with zero c10::cuda symbols. DCU still shares the vendor's streams,
+// so it needs *some* barrier; see SyncCurrentStreamBeforeBlockingCopy below.
 #if !defined(USE_ASCEND) && !defined(USE_TSINGMICRO) && !defined(USE_GCU) && \
-    !defined(USE_MUSA)
+    !defined(USE_MUSA) && !defined(USE_DCU)
 #define FLAGOS_COPY_HAS_CUDA_STREAM 1
 #include <c10/cuda/CUDAStream.h>
 #endif
@@ -53,6 +61,13 @@ inline void SyncCurrentStreamBeforeBlockingCopy() {
   if (stream.stream() != nullptr) {
     stream.synchronize();
   }
+#elif defined(USE_DCU)
+  // DCU shares the vendor's streams, so it needs this barrier just as much as
+  // the other CUDA-family backends -- but it cannot ask which stream is current
+  // (no c10::cuda symbols in a hipified wheel; see the guard above). Fall back to
+  // a device-wide sync: a superset of the per-stream wait, so still correct,
+  // just coarser. Only reached on the blocking-copy paths.
+  ::DeviceSynchronize();
 #endif
 }
 

csrc/runtime/accelerator/CMakeLists.txt

@@ -123,6 +123,11 @@ else()
   target_link_libraries(${LIBRARY_NAME} PRIVATE CUDA::cudart)
 endif()
 
+# libflagos.so needs no torch symbols at all (only the vendor runtime), so the
+# inherited ${PYTORCH_INSTALL_DIR}/lib RUNPATH entry is pure build-machine
+# leakage. Vendor driver dirs are preserved.
+flagos_set_portable_rpath(${LIBRARY_NAME})
+
 install(TARGETS ${LIBRARY_NAME}
     ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
     LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}