diff --git a/README.md b/README.md
index 196aa9a..1842655 100644
--- a/README.md
+++ b/README.md
@@ -7,11 +7,11 @@
 
 ## 编译与安装
 
-`FastPT` 的 wheel 会构建集成的 `torch-mocker`，因此构建环境必须同时具备 DTK、已安装的目标 HIP torch package 和同版本 HIP torch 源码。`TORCH_PATH` 必须指向干净的 HIP torch 安装目录，不能是已经被 FastPT `-C` 覆盖过的 torch。
+`FastPT` 的 wheel 会构建集成的 `torch-mocker`，因此构建环境必须具备 DTK、已安装的目标 HIP torch package，以及 CMake、GNU Make 和可用的编译器。使用的 torch 应为干净的目标 HIP torch package，不能是已经被 FastPT `-C` 覆盖过的 torch。
 
 构建前确认：
 
-- FastPT 分支、已安装 torch 和 `HIP_TORCH_PATH` 对应同一 PyTorch 版本；
+- FastPT 分支与已安装 torch 对应同一 PyTorch 版本；
 - 已安装 `wheel`、CMake、GNU Make 和可用的 DTK 编译环境；
 - 当前环境未安装过 `FastPT wheel`。检查torch的安装路径下， torch/lib 不存在 libtorch_cuda.so、libc10_cuda.so、libtorch_mocker.so，并确认`torch/utils/cpp_extension.py` 不含 USE_FASTPT_CUDA 等 FastPT 路由。
 
@@ -21,16 +21,11 @@
 cd /path/to/fastpt
 source /opt/dtk/cuda/env.sh
 
-# 必须是干净的目标 HIP torch package，
-export TORCH_PATH="$(python3 -c 'import os, torch; print(os.path.dirname(torch.__file__))')"
-# /path/to/pytorch: 为HCU适配的pytorch的源码，注意版本分支。
-export HIP_TORCH_PATH=/path/to/pytorch
-
 git submodule update --init --recursive
 rm -rf build dist fastpt.egg-info torch-mocker/build
 
-# FASTPT_CTEST可以控制CTest部分的编译，默认为 1：构建 torch-mocker 和已注册的 CTest 可执行文件，不需要Ctest可以设置为0.
-python3 setup.py bdist_wheel
+# FASTPT_CTEST 默认为 1；设置为 0 可跳过 CTest 可执行文件构建。
+FASTPT_CTEST=1 python3 setup.py bdist_wheel
 ```
 
 构建产物位于 `dist/fastpt-*.whl`。若仅需要构建发布 wheel，且已通过后文的独立 CTest，可设置 `FASTPT_CTEST=0` 跳过测试可执行文件构建。
@@ -44,7 +39,7 @@ cd /path/to/fastpt
 FASTPT_HIPIFY_ONLY=1 python3 setup.py bdist_wheel
 ```
 
-该模式不需要设置 `TORCH_PATH` 或 `HIP_TORCH_PATH`，但仍需要可导入的匹配 HCU/HIP torch，以生成包版本元数据。它不构建或打包 `libc10_cuda.so`、`libtorch_cuda.so`、`libtorch_mocker.so`、torch overlay 和 `fastpt -C/-E` 环境脚本，只能用于源码/CMake 转码, `CUDAExtension`、`CppExtension`的扩展编译，不能用于基于GPUFusion的不转码 编译或运行环境初始化。
+该模式不需要设置 `TORCH_PATH`，但仍需要可导入的匹配 HCU/HIP torch，以生成包版本元数据。它不构建或打包 `libc10_cuda.so`、`libtorch_cuda.so`、`libtorch_mocker.so`、torch overlay 和 `fastpt -C/-E` 环境脚本，只能用于源码/CMake 转码, `CUDAExtension`、`CppExtension`的扩展编译，不能用于基于GPUFusion的不转码 编译或运行环境初始化。
 
 ### 安装与验证
 
@@ -59,17 +54,18 @@ fastpt.hipify.file --help
 
 ### CTest 测试
 
-`FASTPT_CTEST`不设置时默认为1，会编译CTest测试文件，可通过以下方式进行`torch-mocker`部分的测试：
+`FASTPT_CTEST` 不设置时默认为 1，会构建 `torch-mocker` 的 CTest 可执行文件。wheel 构建后可使用以下命令运行测试：
 
 ```bash
-export LD_LIBRARY_PATH="$PWD/torch-mocker/build:$TORCH_PATH/lib:${LD_LIBRARY_PATH:-}"
+TORCH_LIB_DIR="$(python3 -c 'import os, torch; print(os.path.join(os.path.dirname(torch.__file__), "lib"))')"
+export LD_LIBRARY_PATH="$PWD/torch-mocker/build:$TORCH_LIB_DIR:${LD_LIBRARY_PATH:-}"
 ctest --test-dir torch-mocker/build --output-on-failure -j1
 ```
 
-使用 `ctest --test-dir torch-mocker/build_ctest -N` 可列出已注册测试。定位单个失败时，使用测试名筛选，例如：
+使用 `ctest --test-dir torch-mocker/build -N` 可列出已注册测试。定位单个失败时，使用测试名筛选，例如：
 
 ```bash
-ctest --test-dir torch-mocker/build_ctest --output-on-failure \
+ctest --test-dir torch-mocker/build --output-on-failure \
   -R '^torch_mocker\.c10\.cuda_functions$'
 ```
 
diff --git a/setup.py b/setup.py
index 002b02c..7f593c7 100644
--- a/setup.py
+++ b/setup.py
@@ -18,8 +18,11 @@
 try:
     import torch
     torch_install_path = os.path.dirname(torch.__file__)
+    # CMake consumes TORCH_PATH from its environment. Derive it from the
+    # imported torch package so callers do not need to export it manually.
+    os.environ["TORCH_PATH"] = torch_install_path
 except ImportError as e:
-    raise("Torch is required")
+    raise RuntimeError("Torch is required") from e
 
 torch_version = torch.__version__
 
@@ -49,12 +52,9 @@ def run(self):
                 "from the fastpt repository root before building."
             )
 
-        HIP_TORCH_PATH = os.environ.get("HIP_TORCH_PATH", "")
-        if len(HIP_TORCH_PATH) == 0:
-            raise("Environment variable HIP_TORCH_PATH is required, please set the torch source dir such as: export HIP_TORCH_PATH=/path/to/pytorch ")
-        torch_lib_path = os.environ.get("TORCH_PATH", "")
-        if len(torch_lib_path) == 0:
-            raise("Environment variable TORCH_PATH is required, please set the torch source dir such as: export TORCH_PATH=/path/to/site-packages/torch ")
+        torch_lib_path = torch_install_path
+        if not os.path.isdir(torch_lib_path):
+            raise RuntimeError(f"Torch installation directory does not exist: {torch_lib_path}")
         cmake_args = os.environ.get("TORCH_MOCKER_CMAKE_ARGS", "")
         make_args = "-j8 VERBOSE=1"
         make_args_env = os.environ.get("TORCH_MOCKER_MAKE_ARGS", "")
diff --git a/torch-mocker/CMakeLists.txt b/torch-mocker/CMakeLists.txt
index 83db68b..75d7384 100644
--- a/torch-mocker/CMakeLists.txt
+++ b/torch-mocker/CMakeLists.txt
@@ -6,6 +6,36 @@ project(Torch-Mocker-Test LANGUAGES CXX CUDA)
 set(CMAKE_CXX_STANDARD 20)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)
 find_package(CUDA REQUIRED)
+find_package(Python3 COMPONENTS Interpreter REQUIRED)
+
+set(TORCH_PATH "" CACHE PATH "Path to the installed torch package")
+if(NOT TORCH_PATH)
+    if(DEFINED ENV{TORCH_PATH} AND NOT "$ENV{TORCH_PATH}" STREQUAL "")
+        set(TORCH_PATH "$ENV{TORCH_PATH}")
+    else()
+        execute_process(
+            COMMAND "${Python3_EXECUTABLE}" -c
+                    "import os, torch; print(os.path.dirname(torch.__file__))"
+            RESULT_VARIABLE torch_path_result
+            OUTPUT_VARIABLE TORCH_PATH
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            ERROR_VARIABLE torch_path_error
+        )
+        if(NOT torch_path_result EQUAL 0)
+            message(FATAL_ERROR
+                "Unable to determine the installed torch path: ${torch_path_error}")
+        endif()
+    endif()
+endif()
+
+get_filename_component(TORCH_PATH "${TORCH_PATH}" ABSOLUTE)
+set(TORCH_PATH "${TORCH_PATH}" CACHE PATH "Path to the installed torch package" FORCE)
+if(NOT EXISTS "${TORCH_PATH}/include")
+    message(FATAL_ERROR
+        "Invalid TORCH_PATH '${TORCH_PATH}': torch headers were not found")
+endif()
+set(ENV{TORCH_PATH} "${TORCH_PATH}")
+message(STATUS "Using torch package: ${TORCH_PATH}")
 
 # fmt is a transitive dependency of torch_cuda.  Configure the vendored
 # project before this legacy build sets nvcc as the host C++ compiler.
@@ -51,12 +81,12 @@ include(GNUInstallDirs)
 find_package(Python3 COMPONENTS Development REQUIRED)
 include_directories(${Python3_INCLUDE_DIRS})
 link_directories(${Python3_LIBRARY_DIRS})
-find_package(Python3 COMPONENTS Interpreter REQUIRED)
 file(GLOB_RECURSE Torch_CUDA_SRCS "aten/src/ATen/*.cpp")
 list(FILTER Torch_CUDA_SRCS EXCLUDE REGEX "aten/src/ATen/cuda/nvrtc_stub/*ATenNVRTC.cpp")
 file(GLOB NVRTC_SRCS "aten/src/ATen/cuda/nvrtc_stub/ATenNVRTC.cpp")
 file(GLOB_RECURSE Torch_CUDA_SRCS2 "torch/csrc/*.cpp" "torch/csrc/*.cu")
 list(APPEND Torch_CUDA_SRCS ${Torch_CUDA_SRCS} ${Torch_CUDA_SRCS2})
+
 file(GLOB C10_CUDA_SRCS1 "c10/cuda/*.cpp")
 file(GLOB C10_CUDA_SRCS2 "c10/cuda/impl/*.cpp")
 file(GLOB C10_CUDA_SRCS3 "c10/util/*.cpp")
@@ -67,12 +97,10 @@ file(GLOB Torch_MOCKER_VLOG "torch_mocker/utils/*.cpp")
 # build libcaffe2_nvrtc.so
 #add_library(caffe2_nvrtc SHARED ${NVRTC_SRCS})
 #target_link_libraries(caffe2_nvrtc PRIVATE hiprtc)
-#target_include_directories(caffe2_nvrtc PRIVATE $ENV{HIP_TORCH_PATH})
 #target_compile_definitions(caffe2_nvrtc PRIVATE "-D__HIP_PLATFORM_AMD__")
 
 add_library(torch_mocker  SHARED ${Torch_MOCKER_SRCS})
 add_library(torch_cuda SHARED ${Torch_CUDA_SRCS})
-target_include_directories(torch_cuda PRIVATE $ENV{HIP_TORCH_PATH})
 add_library(c10_cuda  SHARED ${C10_CUDA_SRCS} ${Torch_MOCKER_VLOG})
 
 # Encode DTK MAJOR.MINOR.PATCH as MAJOR * 10000 + MINOR * 100 + PATCH.
@@ -119,11 +147,11 @@ add_dependencies(torch_cuda torch_mocker c10_cuda)
 include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include)
 include_directories(BEFORE ${CMAKE_CURRENT_SOURCE_DIR}/include/aten/src)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/aten/src)
-include_directories(SYSTEM $ENV{TORCH_PATH}/include)
+include_directories(SYSTEM ${TORCH_PATH}/include)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/torch_mocker)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/cudnn-frontend/include)
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/cudnn-frontend/include/cudnn_frontend/thirdparty)
-include_directories(SYSTEM $ENV{TORCH_PATH}/include/torch/csrc/api/include)
+include_directories(SYSTEM ${TORCH_PATH}/include/torch/csrc/api/include)
 
 add_definitions(-DTORCH_USE_CUDA_DSA)
 add_definitions(-DPYTORCH_C10_DRIVER_API_SUPPORTED)
@@ -132,24 +160,22 @@ add_definitions(-DPYTORCH_C10_DRIVER_API_SUPPORTED)
 #add_definitions("-DC10_USING_CUSTOM_GENERATED_MACROS")
 target_compile_definitions(torch_mocker PRIVATE "-D__CUDA_HIP_PLATFORM_AMD__")
 target_compile_definitions(torch_mocker PRIVATE "-D__HIP_PLATFORM_AMD__")
-target_link_libraries(torch_mocker PUBLIC $ENV{TORCH_PATH}/lib/libtorch_hip.so)
+target_link_libraries(torch_mocker PUBLIC ${TORCH_PATH}/lib/libtorch_hip.so)
 target_compile_definitions(torch_cuda PUBLIC "-DC10_USING_CUSTOM_GENERATED_MACROS")
 #target_include_directories(torch_cuda PRIVATE $ENV{TORCH_PATH}/include)
-target_include_directories(c10_cuda PRIVATE $ENV{HIP_TORCH_PATH})
 target_compile_definitions(c10_cuda PRIVATE "-D__CUDA_HIP_PLATFORM_AMD__")
 target_compile_definitions(c10_cuda PRIVATE "-D__HIP_PLATFORM_AMD__")
 target_compile_definitions(c10_cuda PUBLIC "-DC10_USING_CUSTOM_GENERATED_MACROS")
 
 target_link_libraries(torch_cuda PUBLIC torch_mocker)
-message("$ENV{TORCH_PATH}/lib/libc10_hip.so")
-target_link_libraries(c10_cuda PUBLIC $ENV{TORCH_PATH}/lib/libc10_hip.so cublasLt torch_mocker)
+message("${TORCH_PATH}/lib/libc10_hip.so")
+target_link_libraries(c10_cuda PUBLIC ${TORCH_PATH}/lib/libc10_hip.so cublasLt torch_mocker)
 
-target_include_directories(torch_cuda PRIVATE $ENV{HIP_TORCH_PATH})
 target_compile_definitions(torch_cuda PRIVATE "-D__HIP_PLATFORM_AMD__")
 target_compile_definitions(torch_cuda PUBLIC "-DC10_USING_CUSTOM_GENERATED_MACROS")
 if(USE_GLOG)
     target_compile_definitions(torch_cuda PUBLIC "-DC10_USE_GLOG")
-    set(TORCH_LIBS_DIR $ENV{TORCH_PATH}/../torch.libs)
+    set(TORCH_LIBS_DIR ${TORCH_PATH}/../torch.libs)
     if(EXISTS ${TORCH_LIBS_DIR})
         file(GLOB GLOG_LIB "${TORCH_LIBS_DIR}/libglog-*.so*")
         if(NOT GLOG_LIB)
@@ -163,14 +189,14 @@ if(USE_GLOG)
     endif()
 endif()
 target_link_libraries(torch_cuda PUBLIC torch_mocker)
-target_link_libraries(torch_cuda PUBLIC $ENV{TORCH_PATH}/lib/libtorch_hip.so)
-target_link_libraries(torch_cuda PUBLIC $ENV{TORCH_PATH}/lib/libtorch_python.so)
+target_link_libraries(torch_cuda PUBLIC ${TORCH_PATH}/lib/libtorch_hip.so)
+# target_link_libraries(torch_cuda PUBLIC ${TORCH_PATH}/lib/libtorch_python.so)
 target_link_directories(torch_cuda PUBLIC ${CMAKE_INSTALL_PREFIX}/${CMAKE_INSTALL_LIBDIR})
 target_link_libraries(torch_cuda PUBLIC cuda nvrtc cudnn_cnn_infer cublas cublasLt cufft cusparse cudnn cudnn_cnn_infer cudnn_adv_infer nvToolsExt fmt)
 #target_link_libraries(torch_cuda PUBLIC cuda nvrtc cudnn_cnn_infer cublas cublasLt cufft cusparse cudnn cudnn_cnn_infer cudnn_adv_infer)
 target_link_libraries(torch_cuda PUBLIC c10_cuda)
-target_link_libraries(torch_cuda PUBLIC ${Python3_LIBRARIES})
-message("$ENV{TORCH_PATH}/lib/libtorch_cuda.so")
+# target_link_libraries(torch_cuda PUBLIC ${Python3_LIBRARIES})
+message("${TORCH_PATH}/lib/libtorch_cuda.so")
 
 set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++20")
 include(CTest)
diff --git a/torch-mocker/c10/util/StringUtil.cpp b/torch-mocker/c10/util/StringUtil.cpp
index f1339b3..026feaa 100644
--- a/torch-mocker/c10/util/StringUtil.cpp
+++ b/torch-mocker/c10/util/StringUtil.cpp
@@ -2,4 +2,221 @@
 // SPDX-License-Identifier: BSD-3-Clause
 // Modified by Hygon Information Technology Co., Ltd., 2026.
 
-#include<c10/util/StringUtil.cpp>
+#include <c10/util/StringUtil.h>
+
+#include <string>
+
+#ifndef _WIN32
+#include <codecvt>
+#include <locale>
+#else
+#include <c10/util/Unicode.h>
+#endif
+
+namespace c10 {
+
+namespace detail {
+
+std::string StripBasename(const std::string& full_path) {
+#ifdef _WIN32
+  const std::string separators("/\\");
+#else
+  const std::string separators("/");
+#endif
+  size_t pos = full_path.find_last_of(separators);
+  if (pos != std::string::npos) {
+    return full_path.substr(pos + 1, std::string::npos);
+  } else {
+    return full_path;
+  }
+}
+
+std::string ExcludeFileExtension(const std::string& file_name) {
+  const char sep = '.';
+  auto end_index = file_name.find_last_of(sep) == std::string::npos
+      ? -1
+      : file_name.find_last_of(sep);
+  return file_name.substr(0, end_index);
+}
+
+// Narrows the wstr argument and then passes it to _str.
+// Assumes that the input (wide) text is encoded as UTF-16.
+static std::ostream& _strFromWide(
+    std::ostream& ss,
+    const std::wstring& wString);
+
+#ifndef _WIN32
+
+C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wdeprecated-declarations")
+// TODO (huydhn) https://en.cppreference.com/w/cpp/header/codecvt has been
+// deprecated in C++17 but there is no alternative yet, so I just ack it
+static std::ostream& _strFromWide(
+    std::ostream& ss,
+    const std::wstring& wString) {
+  std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
+  return _str(ss, converter.to_bytes(wString));
+}
+C10_DIAGNOSTIC_POP()
+
+#else // #ifndef _WIN32
+// The WIN32 implementation of wstring_convert leaks memory; see
+// https://github.com/microsoft/STL/issues/443
+
+static std::ostream& _strFromWide(
+    std::ostream& ss,
+    const std::wstring& wString) {
+  return _str(ss, u16u8(wString));
+}
+
+#endif // _WIN32
+
+std::ostream& _str(std::ostream& ss, const wchar_t* wCStr) {
+  return _strFromWide(ss, std::wstring(wCStr));
+}
+std::ostream& _str(std::ostream& ss, const wchar_t& wChar) {
+  return _strFromWide(ss, std::wstring(1, wChar));
+}
+std::ostream& _str(std::ostream& ss, const std::wstring& wString) {
+  return _strFromWide(ss, wString);
+}
+
+} // namespace detail
+
+std::ostream& operator<<(std::ostream& out, const SourceLocation& loc) {
+  out << loc.function << " at " << loc.file << ':' << loc.line;
+  return out;
+}
+
+size_t ReplaceAll(std::string& s, std::string_view from, std::string_view to) {
+  if (from.empty()) {
+    return 0;
+  }
+
+  size_t numReplaced = 0;
+  std::string::size_type last_pos = 0u;
+  std::string::size_type cur_pos = 0u;
+  std::string::size_type write_pos = 0u;
+  const std::string_view input(s);
+
+  if (from.size() >= to.size()) {
+    // If the replacement string is not larger than the original, we
+    // can do the replacement in-place without allocating new storage.
+    char* s_data = &s[0];
+
+    while ((cur_pos = s.find(from.data(), last_pos, from.size())) !=
+           std::string::npos) {
+      ++numReplaced;
+      // Append input between replaced sub-strings
+      if (write_pos != last_pos) {
+        std::copy(s_data + last_pos, s_data + cur_pos, s_data + write_pos);
+      }
+      write_pos += cur_pos - last_pos;
+      // Append the replacement sub-string
+      std::copy(to.begin(), to.end(), s_data + write_pos);
+      write_pos += to.size();
+      // Start search from next character after `from`
+      last_pos = cur_pos + from.size();
+    }
+
+    // Append any remaining input after replaced sub-strings
+    if (write_pos != last_pos) {
+      std::copy(s_data + last_pos, s_data + input.size(), s_data + write_pos);
+      write_pos += input.size() - last_pos;
+      s.resize(write_pos);
+    }
+    return numReplaced;
+  }
+
+  // Otherwise, do an out-of-place replacement in a temporary buffer
+  std::string buffer;
+
+  while ((cur_pos = s.find(from.data(), last_pos, from.size())) !=
+         std::string::npos) {
+    ++numReplaced;
+    // Append input between replaced sub-strings
+    buffer.append(input.begin() + last_pos, input.begin() + cur_pos);
+    // Append the replacement sub-string
+    buffer.append(to.begin(), to.end());
+    // Start search from next character after `from`
+    last_pos = cur_pos + from.size();
+  }
+  if (numReplaced == 0) {
+    // If nothing was replaced, don't modify the input
+    return 0;
+  }
+  // Append any remaining input after replaced sub-strings
+  buffer.append(input.begin() + last_pos, input.end());
+  s = std::move(buffer);
+  return numReplaced;
+}
+
+template <>
+std::optional<int64_t> tryToNumber<int64_t>(const std::string& symbol) {
+  return tryToNumber<int64_t>(symbol.c_str());
+}
+
+template <>
+std::optional<int64_t> tryToNumber<int64_t>(const char* symbol) {
+  // TODO Using strtoll for portability. Consider using std::from_chars in the
+  // future. According to https://libcxx.llvm.org/Status/Cxx17.html,
+  // std::from_chars is not supported until clang 20. We will need MSVC to also
+  // fully support std::from_chars.
+  if (!symbol) {
+    return std::nullopt;
+  }
+  char* end = nullptr;
+  errno = 0;
+  int64_t value = strtoll(symbol, &end, 0);
+  if (errno != 0) {
+    errno = 0;
+    return std::nullopt;
+  }
+  if (*end != '\0' || end == symbol) {
+    return std::nullopt;
+  }
+  return value;
+}
+
+template <>
+std::optional<double> tryToNumber<double>(const std::string& symbol) {
+  return tryToNumber<double>(symbol.c_str());
+}
+
+template <>
+std::optional<double> tryToNumber<double>(const char* symbol) {
+  // TODO Using strtod for portability. Consider using std::from_chars in the
+  // future. According to https://libcxx.llvm.org/Status/Cxx17.html,
+  // std::from_chars is not supported until clang 20. We will need MSVC to also
+  // fully support std::from_chars.
+  if (!symbol) {
+    return std::nullopt;
+  }
+  char* end = nullptr;
+  errno = 0;
+  double value = strtod(symbol, &end);
+  if (errno != 0) {
+    errno = 0;
+    return std::nullopt;
+  }
+  if (*end != '\0' || end == symbol) {
+    return std::nullopt;
+  }
+  return value;
+}
+
+std::vector<std::string_view> split(std::string_view target, char delimiter) {
+  std::vector<std::string_view> atoms;
+  std::string_view buffer = target;
+  while (!buffer.empty()) {
+    auto i = buffer.find(delimiter);
+    if (i == std::string_view::npos) {
+      atoms.push_back(buffer);
+      buffer.remove_prefix(buffer.size());
+    } else {
+      atoms.push_back(buffer.substr(0, i));
+      buffer.remove_prefix(i + 1);
+    }
+  }
+  return atoms;
+}
+} // namespace c10
diff --git a/torch-mocker/c10/util/flags_use_no_gflags.cpp b/torch-mocker/c10/util/flags_use_no_gflags.cpp
index aaacb6a..bd64c96 100644
--- a/torch-mocker/c10/util/flags_use_no_gflags.cpp
+++ b/torch-mocker/c10/util/flags_use_no_gflags.cpp
@@ -2,4 +2,209 @@
 // SPDX-License-Identifier: BSD-3-Clause
 // Modified by Hygon Information Technology Co., Ltd., 2026.
 
-#include<c10/util/flags_use_no_gflags.cpp>
+#include <c10/macros/Macros.h>
+#include <c10/util/Flags.h>
+
+#include <cstdlib>
+#include <iostream>
+#include <sstream>
+#include <string>
+
+#ifndef C10_USE_GFLAGS
+
+namespace c10 {
+
+using std::string;
+
+C10_DEFINE_REGISTRY(C10FlagsRegistry, C10FlagParser, const string&)
+
+namespace {
+bool gCommandLineFlagsParsed = false;
+// Since flags is going to be loaded before logging, we would
+// need to have a stringstream to hold the messages instead of directly
+// using caffe logging.
+std::stringstream& GlobalInitStream() {
+  static std::stringstream ss;
+  return ss;
+}
+const char* gUsageMessage = "(Usage message not set.)";
+} // namespace
+
+C10_EXPORT void SetUsageMessage(const string& str) {
+  static string usage_message_safe_copy = str;
+  gUsageMessage = usage_message_safe_copy.c_str();
+}
+
+C10_EXPORT const char* UsageMessage() {
+  return gUsageMessage;
+}
+
+C10_EXPORT bool ParseCommandLineFlags(int* pargc, char*** pargv) {
+  if (*pargc == 0)
+    return true;
+  char** argv = *pargv;
+  bool success = true;
+  GlobalInitStream() << "Parsing commandline arguments for c10." << '\n';
+  // write_head is the location we write the unused arguments to.
+  int write_head = 1;
+  for (int i = 1; i < *pargc; ++i) {
+    string arg(argv[i]);
+
+    if (arg.find("--help") != string::npos) {
+      // Print the help message, and quit.
+      std::cout << UsageMessage() << '\n';
+      std::cout << "Arguments: " << '\n';
+      for (const auto& help_msg : C10FlagsRegistry()->HelpMessage()) {
+        std::cout << "    " << help_msg.first << ": " << help_msg.second
+                  << '\n';
+      }
+      exit(0);
+    }
+    // If the arg does not start with "--", we will ignore it.
+    if (arg[0] != '-' || arg[1] != '-') {
+      GlobalInitStream()
+          << "C10 flag: commandline argument does not match --name=var "
+             "or --name format: "
+          << arg << ". Ignoring this argument." << '\n';
+      argv[write_head++] = argv[i];
+      continue;
+    }
+
+    string key;
+    string value;
+    size_t prefix_idx = arg.find('=');
+    if (prefix_idx == string::npos) {
+      // If there is no equality char in the arg, it means that the
+      // arg is specified in the next argument.
+      key = arg.substr(2, arg.size() - 2);
+      ++i;
+      if (i == *pargc) {
+        GlobalInitStream()
+            << "C10 flag: reached the last commandline argument, but "
+               "I am expecting a value for "
+            << arg;
+        success = false;
+        break;
+      }
+      value = string(argv[i]);
+    } else {
+      // If there is an equality character, we will basically use the value
+      // after the "=".
+      key = arg.substr(2, prefix_idx - 2);
+      value = arg.substr(prefix_idx + 1, string::npos);
+    }
+    // If the flag is not registered, we will ignore it.
+    if (!C10FlagsRegistry()->Has(key)) {
+      GlobalInitStream() << "C10 flag: unrecognized commandline argument: "
+                         << arg << '\n';
+      success = false;
+      break;
+    }
+    std::unique_ptr<C10FlagParser> parser(
+        C10FlagsRegistry()->Create(key, value));
+    if (!parser->success()) {
+      GlobalInitStream() << "C10 flag: illegal argument: " << arg << '\n';
+      success = false;
+      break;
+    }
+  }
+  *pargc = write_head;
+  gCommandLineFlagsParsed = true;
+  // TODO: when we fail commandline flag parsing, shall we continue, or
+  // shall we just quit loudly? Right now we carry on the computation, but
+  // since there are failures in parsing, it is very likely that some
+  // downstream things will break, in which case it makes sense to quit loud
+  // and early.
+  if (!success) {
+    std::cerr << GlobalInitStream().str();
+  }
+  // Clear the global init stream.
+  GlobalInitStream().str(std::string());
+  return success;
+}
+
+C10_EXPORT bool CommandLineFlagsHasBeenParsed() {
+  return gCommandLineFlagsParsed;
+}
+
+template <>
+C10_EXPORT bool C10FlagParser::Parse<string>(
+    const string& content,
+    string* value) {
+  *value = content;
+  return true;
+}
+
+template <>
+C10_EXPORT bool C10FlagParser::Parse<int>(const string& content, int* value) {
+  try {
+    *value = std::atoi(content.c_str());
+    return true;
+  } catch (...) {
+    GlobalInitStream() << "C10 flag error: Cannot convert argument to int: "
+                       << content << '\n';
+    return false;
+  }
+}
+
+template <>
+C10_EXPORT bool C10FlagParser::Parse<int64_t>(
+    const string& content,
+    int64_t* value) {
+  try {
+    static_assert(sizeof(long long) == sizeof(int64_t));
+#ifdef __ANDROID__
+    // Android does not have std::atoll.
+    *value = atoll(content.c_str());
+#else
+    *value = std::atoll(content.c_str());
+#endif
+    return true;
+  } catch (...) {
+    GlobalInitStream() << "C10 flag error: Cannot convert argument to int: "
+                       << content << '\n';
+    return false;
+  }
+}
+
+template <>
+C10_EXPORT bool C10FlagParser::Parse<double>(
+    const string& content,
+    double* value) {
+  try {
+    *value = std::atof(content.c_str());
+    return true;
+  } catch (...) {
+    GlobalInitStream() << "C10 flag error: Cannot convert argument to double: "
+                       << content << '\n';
+    return false;
+  }
+}
+
+template <>
+C10_EXPORT bool C10FlagParser::Parse<bool>(const string& content, bool* value) {
+  if (content == "false" || content == "False" || content == "FALSE" ||
+      content == "0") {
+    *value = false;
+    return true;
+  } else if (
+      content == "true" || content == "True" || content == "TRUE" ||
+      content == "1") {
+    *value = true;
+    return true;
+  } else {
+    GlobalInitStream()
+        << "C10 flag error: Cannot convert argument to bool: " << content
+        << '\n'
+        << "Note that if you are passing in a bool flag, you need to "
+           "explicitly specify it, like --arg=True or --arg True. Otherwise, "
+           "the next argument may be inadvertently used as the argument, "
+           "causing the above error."
+        << '\n';
+    return false;
+  }
+}
+
+} // namespace c10
+
+#endif // C10_USE_GFLAGS
diff --git a/torch-mocker/include/torch/csrc/autograd/functions/comm.h b/torch-mocker/include/torch/csrc/autograd/functions/comm.h
index c6aa1fe..35034c3 100644
--- a/torch-mocker/include/torch/csrc/autograd/functions/comm.h
+++ b/torch-mocker/include/torch/csrc/autograd/functions/comm.h
@@ -12,7 +12,7 @@
 #ifndef USE_ROCM
 #include <c10/cuda/CUDAStream.h>
 #else
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 #endif
 #include <optional>
 
@@ -22,23 +22,13 @@
 namespace torch::autograd {
 
 struct TORCH_CUDA_CU_API Scatter : public Node {
-  #ifndef USE_ROCM
-  explicit Scatter(
+    explicit Scatter(
       std::vector<at::Device> devices,
       std::optional<std::vector<int64_t>> chunk_sizes = std::nullopt,
       int64_t dim = 0,
       std::optional<std::vector<std::optional<at::cuda::CUDAStream>>> streams =
           std::nullopt,
       bool unsqueeze_scalars = false);
-  #else
-  explicit Scatter(
-      std::vector<at::Device> devices,
-      std::optional<std::vector<int64_t>> chunk_sizes = std::nullopt,
-      int64_t dim = 0,
-      std::optional<std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>> streams =
-          std::nullopt,
-      bool unsqueeze_scalars = false);
-  #endif
   ~Scatter() override;
 
   variable_list apply(variable_list&& inputs) override;
@@ -46,11 +36,7 @@ struct TORCH_CUDA_CU_API Scatter : public Node {
   std::vector<at::Device> devices_;
   std::optional<std::vector<int64_t>> chunk_sizes_;
   int64_t dim_;
-  #ifndef USE_ROCM
   std::optional<std::vector<std::optional<at::cuda::CUDAStream>>> streams_;
-  #else
-  std::optional<std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>> streams_;
-  #endif
   bool unsqueeze_scalars_;
 };
 
diff --git a/torch-mocker/include/torch/csrc/cuda/CUDAPluggableAllocator.h b/torch-mocker/include/torch/csrc/cuda/CUDAPluggableAllocator.h
index 58e07b2..4b8d659 100644
--- a/torch-mocker/include/torch/csrc/cuda/CUDAPluggableAllocator.h
+++ b/torch-mocker/include/torch/csrc/cuda/CUDAPluggableAllocator.h
@@ -14,7 +14,7 @@
 #else
 #include <c10/hip/HIPGraphsC10Utils.h>
 #include <c10/hip/HIPMacros.h>
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 
 #include <c10/hip/HIPCachingAllocator.h>
 #endif
diff --git a/torch-mocker/include/torch/csrc/cuda/Stream.h b/torch-mocker/include/torch/csrc/cuda/Stream.h
index 007b4ed..62b6dc1 100644
--- a/torch-mocker/include/torch/csrc/cuda/Stream.h
+++ b/torch-mocker/include/torch/csrc/cuda/Stream.h
@@ -8,19 +8,14 @@
 #ifndef USE_ROCM
 #include <c10/cuda/CUDAStream.h>
 #else
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 #endif
 #include <torch/csrc/Stream.h>
 #include <torch/csrc/python_headers.h>
 
 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
 struct THCPStream : THPStream {
-#ifndef USE_ROCM
   at::cuda::CUDAStream cuda_stream;
-#else
-  at::hip::HIPStreamMasqueradingAsCUDA cuda_stream;
-#endif
-
 };
 extern PyObject* THCPStreamClass;
 
diff --git a/torch-mocker/include/torch/csrc/cuda/comm.h b/torch-mocker/include/torch/csrc/cuda/comm.h
index bff061b..99e8fe0 100644
--- a/torch-mocker/include/torch/csrc/cuda/comm.h
+++ b/torch-mocker/include/torch/csrc/cuda/comm.h
@@ -37,26 +37,16 @@ TORCH_CUDA_CU_API std::vector<at::Tensor>& scatter_out(
     const at::Tensor& tensor,
     std::vector<at::Tensor>& out_tensors,
     int64_t dim = 0,
-    #ifndef USE_ROCM
     const std::optional<std::vector<std::optional<at::cuda::CUDAStream>>>&
         streams = std::nullopt);
-    #else
-    const std::optional<std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>>&
-        streams = std::nullopt);
-    #endif
 
 TORCH_CUDA_CU_API std::vector<at::Tensor> scatter(
     const at::Tensor& tensor,
     at::IntArrayRef devices,
     const std::optional<std::vector<int64_t>>& chunk_sizes = std::nullopt,
     int64_t dim = 0,
-    #ifndef USE_ROCM
     const std::optional<std::vector<std::optional<at::cuda::CUDAStream>>>&
         streams = std::nullopt);
-    #else
-    const std::optional<std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>>&
-        streams = std::nullopt);
-    #endif
 
 TORCH_CUDA_CU_API at::Tensor& gather_out(
     at::TensorList tensors,
diff --git a/torch-mocker/include/torch/csrc/cuda/nccl.h b/torch-mocker/include/torch/csrc/cuda/nccl.h
index e71edb3..93361ff 100644
--- a/torch-mocker/include/torch/csrc/cuda/nccl.h
+++ b/torch-mocker/include/torch/csrc/cuda/nccl.h
@@ -120,11 +120,8 @@ TORCH_CUDA_CPP_API void check_inputs(
 } // namespace detail
 
 using comm_list = std::vector<ncclComm_t>;
-#ifndef USE_ROCM
 using stream_list = std::vector<std::optional<at::cuda::CUDAStream>>;
-#else
-using stream_list = std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>;
-#endif
+
 TORCH_CUDA_CPP_API std::uint64_t version();
 TORCH_CUDA_CPP_API const char* version_suffix();
 
@@ -170,27 +167,20 @@ TORCH_CUDA_CPP_API void reduce_scatter(
     int32_t op = static_cast<int>(ncclRedOp::Sum),
     const stream_list& streams = {},
     const comm_list& user_comms = {});
-#ifndef USE_ROCM
+
 TORCH_CUDA_CPP_API void scatter(
     const std::vector<at::Tensor>& inputs,
     at::Tensor& outputs,
     ncclComm_t comm,
     at::cuda::CUDAStream& stream,
     int32_t root = 0);
-#else
-TORCH_CUDA_CPP_API void scatter(
-    const std::vector<at::Tensor>& inputs,
-    at::Tensor& outputs,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA& stream,
-    int32_t root = 0);
-#endif
+
 TORCH_CUDA_CPP_API void all_gather(
     const std::vector<at::Tensor>& inputs,
     std::vector<at::Tensor>& outputs,
     const stream_list& streams = {},
     const comm_list& user_comms = {});
-#ifndef USE_ROCM
+
 TORCH_CUDA_CPP_API void gather(
     const at::Tensor& inputs,
     std::vector<at::Tensor>& outputs,
@@ -234,49 +224,5 @@ TORCH_CUDA_CPP_API void recv(
     ncclComm_t comm,
     at::cuda::CUDAStream stream,
     int src);
-#else
-TORCH_CUDA_CPP_API void gather(
-    const at::Tensor& inputs,
-    std::vector<at::Tensor>& outputs,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA& stream,
-    int32_t root = 0);
-
-TORCH_CUDA_CPP_API void all2all_single_equal_split(
-    at::Tensor& input,
-    at::Tensor& output,
-    int size,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA& stream);
-
-TORCH_CUDA_CPP_API void all2all_single_unequal_split(
-    void* sendbuff,
-    const size_t* sendcounts,
-    const size_t* senddispls,
-    void* recvbuff,
-    const size_t* recvcounts,
-    const size_t* recvdispls,
-    size_t size,
-    c10::ScalarType type,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA& stream);
 
-TORCH_CUDA_CPP_API void all2all(
-    std::vector<at::Tensor>& outputTensors,
-    std::vector<at::Tensor>& inputTensors,
-    ncclComm_t _comm,
-    at::hip::HIPStreamMasqueradingAsCUDA& stream);
-
-TORCH_CUDA_CPP_API void send(
-    const at::Tensor& input,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA stream,
-    int dst);
-
-TORCH_CUDA_CPP_API void recv(
-    at::Tensor& output,
-    ncclComm_t comm,
-    at::hip::HIPStreamMasqueradingAsCUDA stream,
-    int src);
-#endif
 } // namespace torch::cuda::nccl
diff --git a/torch-mocker/include/torch/csrc/cuda/utils.h b/torch-mocker/include/torch/csrc/cuda/utils.h
index bdad6de..174dfea 100644
--- a/torch-mocker/include/torch/csrc/cuda/utils.h
+++ b/torch-mocker/include/torch/csrc/cuda/utils.h
@@ -7,15 +7,11 @@
 #ifndef USE_ROCM
 #include <c10/cuda/CUDAStream.h>
 #else
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 #endif
 #include <torch/csrc/utils/python_numbers.h>
 
 #include <vector>
 
-#ifndef USE_ROCM
 std::vector<std::optional<at::cuda::CUDAStream>>
-#else
-std::vector<std::optional<at::hip::HIPStreamMasqueradingAsCUDA>>
-#endif
 THPUtils_PySequence_to_CUDAStreamList(PyObject* obj);
diff --git a/torch-mocker/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp b/torch-mocker/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
index f63e813..0dc1957 100644
--- a/torch-mocker/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
+++ b/torch-mocker/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp
@@ -8,7 +8,7 @@
 #ifndef USE_ROCM
 #include <c10/cuda/CUDAStream.h>
 #else
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 #endif
 #include <torch/csrc/distributed/c10d/Store.hpp>
 #include <torch/csrc/distributed/c10d/Work.hpp>
@@ -64,7 +64,6 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target {
   at::Tensor allReduce(const at::Tensor& input, AllReduceAlgo algo);
 
  private:
-  #ifndef USE_ROCM
   at::Tensor oneShotAllReduce(
       const at::Tensor& input,
       at::cuda::CUDAStream& stream);
@@ -72,15 +71,6 @@ class TORCH_API IntraNodeComm : public c10::intrusive_ptr_target {
   at::Tensor twoShotAllReduce(
       const at::Tensor& input,
       at::cuda::CUDAStream& stream);
-  #else
-  at::Tensor oneShotAllReduce(
-      const at::Tensor& input,
-      at::hip::HIPStreamMasqueradingAsCUDA& stream);
-
-  at::Tensor twoShotAllReduce(
-      const at::Tensor& input,
-      at::hip::HIPStreamMasqueradingAsCUDA& stream);
-  #endif
 
   c10::intrusive_ptr<Store> store_;
   size_t rank_;
diff --git a/torch-mocker/include/torch/csrc/jit/cuda/cuda.h b/torch-mocker/include/torch/csrc/jit/cuda/cuda.h
index bad00d2..101e6c7 100644
--- a/torch-mocker/include/torch/csrc/jit/cuda/cuda.h
+++ b/torch-mocker/include/torch/csrc/jit/cuda/cuda.h
@@ -7,7 +7,7 @@
 #include <c10/cuda/CUDAStream.h>
 #else
 #include <ATen/hip/HIPEvent.h>
-#include <ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h>
+#include <c10/hip/HIPStream.h>
 #endif
 #include <c10/core/Device.h>
 #include <torch/custom_class.h>
@@ -186,23 +186,24 @@ TORCH_LIBRARY(cuda, m) {
 }
 #else
 class CUDAEvent;
-// This class is a wrapper around c10::hip::HIPStreamMasqueradingAsCUDA.
+// This class is a wrapper around c10::cuda::CUDAStream backed by HIP.
 // It is needed because TorchBind does not support all of the argument types
-// for c10::hip::HIPStreamMasqueradingAsCUDA. For more details, please refer to
-// ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h.
-class HIPStreamMasqueradingAsCUDA final : public CustomClassHolder {
+// for c10::cuda::CUDAStream. For more details, please refer to
+// c10/hip/HIPStream.h.
+class CUDAStream final : public CustomClassHolder {
  public:
-  HIPStreamMasqueradingAsCUDA(
+  CUDAStream(
       std::optional<c10::Device> device = std::nullopt,
       int64_t priority = 0) {
     c10::DeviceIndex device_index =
-        device.has_value() ? device->index() : c10::hip::current_device();
-    stream_ = std::make_unique<c10::hip::HIPStreamMasqueradingAsCUDA>(
-        c10::hip::getStreamFromPoolMasqueradingAsCUDA(static_cast<int>(priority), device_index));
+        device.has_value() ? device->index() : c10::cuda::current_device();
+    stream_ = std::make_unique<c10::cuda::CUDAStream>(
+        c10::cuda::getStreamFromPool(
+            static_cast<int>(priority), device_index));
   }
 
-  HIPStreamMasqueradingAsCUDA(c10::hip::HIPStreamMasqueradingAsCUDA s) {
-    stream_ = std::make_unique<c10::hip::HIPStreamMasqueradingAsCUDA>(s);
+  CUDAStream(c10::cuda::CUDAStream s) {
+    stream_ = std::make_unique<c10::cuda::CUDAStream>(s);
   }
 
   bool query() {
@@ -218,7 +219,7 @@ class HIPStreamMasqueradingAsCUDA final : public CustomClassHolder {
 
   void waitEvent(const c10::intrusive_ptr<CUDAEvent>& event);
 
-  void waitStream(const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream);
+  void waitStream(const c10::intrusive_ptr<CUDAStream>& stream);
 
   /// Get the CUDA device index that this stream is associated with.
   int64_t device_index() const {
@@ -237,11 +238,11 @@ class HIPStreamMasqueradingAsCUDA final : public CustomClassHolder {
   }
 
  private:
-  std::unique_ptr<c10::hip::HIPStreamMasqueradingAsCUDA> stream_;
+  std::unique_ptr<c10::cuda::CUDAStream> stream_;
   friend class CUDAEvent;
 };
 
-// This class is a wrapper around at::hip::HIPStreamMasqueradingAsCUDA.
+// This class is a wrapper around at::cuda::CUDAStream backed by HIP.
 // It is needed because TorchBind does not support all of the argument types
 // for at::cuda::CUDAEvent. For more details, please refer to
 // aten/src/ATen/cuda/CUDAEvent.h.
@@ -281,21 +282,21 @@ class CUDAEvent final : public CustomClassHolder {
     return event_->query();
   }
 
-  void record(const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream);
+  void record(const c10::intrusive_ptr<CUDAStream>& stream);
 
   void synchronize() {
     event_->synchronize();
   }
-  void wait(const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream);
+  void wait(const c10::intrusive_ptr<CUDAStream>& stream);
 
  private:
-  void recordInternal(HIPStreamMasqueradingAsCUDA* stream);
+  void recordInternal(CUDAStream* stream);
   std::unique_ptr<at::cuda::CUDAEvent> event_;
 
-  friend class HIPStreamMasqueradingAsCUDA;
+  friend class CUDAStream;
 };
 
-inline c10::intrusive_ptr<CUDAEvent> HIPStreamMasqueradingAsCUDA::recordEvent(
+inline c10::intrusive_ptr<CUDAEvent> CUDAStream::recordEvent(
     c10::intrusive_ptr<CUDAEvent> event) {
   if (!event) {
     event = c10::make_intrusive<CUDAEvent>();
@@ -305,31 +306,34 @@ inline c10::intrusive_ptr<CUDAEvent> HIPStreamMasqueradingAsCUDA::recordEvent(
   return event;
 }
 
-inline void HIPStreamMasqueradingAsCUDA::waitEvent(const c10::intrusive_ptr<CUDAEvent>& event) {
+inline void CUDAStream::waitEvent(
+    const c10::intrusive_ptr<CUDAEvent>& event) {
   event->event_->block(*stream_);
 }
 
-inline void HIPStreamMasqueradingAsCUDA::waitStream(
-    const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream) {
+inline void CUDAStream::waitStream(
+    const c10::intrusive_ptr<CUDAStream>& stream) {
   auto ev = c10::make_intrusive<CUDAEvent>();
   stream->recordEvent(ev);
   waitEvent(ev);
 }
 
-inline void CUDAEvent::record(const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream) {
+inline void CUDAEvent::record(
+    const c10::intrusive_ptr<CUDAStream>& stream) {
   event_->record(*stream->stream_);
 }
 
-inline void CUDAEvent::recordInternal(HIPStreamMasqueradingAsCUDA* stream) {
+inline void CUDAEvent::recordInternal(CUDAStream* stream) {
   event_->record(*stream->stream_);
 }
 
-inline void CUDAEvent::wait(const c10::intrusive_ptr<HIPStreamMasqueradingAsCUDA>& stream) {
+inline void CUDAEvent::wait(
+    const c10::intrusive_ptr<CUDAStream>& stream) {
   event_->block(*stream->stream_);
 }
 
 TORCH_LIBRARY(cuda, m) {
-  auto stream_class = m.class_<torch::jit::HIPStreamMasqueradingAsCUDA>("Stream").def(
+  auto stream_class = m.class_<torch::jit::CUDAStream>("Stream").def(
       torch::init<std::optional<c10::Device>, int64_t>(),
       "",
       {torch::arg("device") = std::nullopt, torch::arg("priority") = 0});
@@ -340,14 +344,14 @@ TORCH_LIBRARY(cuda, m) {
        torch::arg("blocking") = false,
        torch::arg("interprocess") = false});
 
-  stream_class.def("query", &HIPStreamMasqueradingAsCUDA::query)
-      .def("record_event", &HIPStreamMasqueradingAsCUDA::recordEvent)
-      .def("synchronize", &HIPStreamMasqueradingAsCUDA::synchronize)
-      .def("wait_event", &HIPStreamMasqueradingAsCUDA::waitEvent)
-      .def("wait_stream", &HIPStreamMasqueradingAsCUDA::waitStream)
-      .def("device_index", &HIPStreamMasqueradingAsCUDA::device_index)
-      .def_property("device", &HIPStreamMasqueradingAsCUDA::device)
-      .def("id", &HIPStreamMasqueradingAsCUDA::id);
+  stream_class.def("query", &CUDAStream::query)
+      .def("record_event", &CUDAStream::recordEvent)
+      .def("synchronize", &CUDAStream::synchronize)
+      .def("wait_event", &CUDAStream::waitEvent)
+      .def("wait_stream", &CUDAStream::waitStream)
+      .def("device_index", &CUDAStream::device_index)
+      .def_property("device", &CUDAStream::device)
+      .def("id", &CUDAStream::id);
 
   event_class.def("elapsed_time", &CUDAEvent::elapsedTime)
       .def("query", &CUDAEvent::query)
diff --git a/torch-mocker/setup.py b/torch-mocker/setup.py
index 3d941e3..b12b54a 100644
--- a/torch-mocker/setup.py
+++ b/torch-mocker/setup.py
@@ -15,9 +15,12 @@
     import torch
 
     torch_install_path = os.path.dirname(torch.__file__)
+    # CMake consumes TORCH_PATH from its environment. Derive it from the
+    # imported torch package so callers do not need to export it manually.
+    os.environ["TORCH_PATH"] = torch_install_path
 except ImportError as e:
     print("Torch is required, ", e)
-    raise("Torch is required")
+    raise RuntimeError("Torch is required") from e
 
 include_ends=['.h', '.cuh', '.hpp']
 
@@ -78,13 +81,10 @@ def run(self):
             shutil.copy(src, dest_path)
             print(f"Copying {src} to build directory: {dest_path}")
 
-        # 原始构建逻辑保持不变...
-        HIP_TORCH_PATH = os.environ.get("HIP_TORCH_PATH", "")
-        if len(HIP_TORCH_PATH)==0:
-            raise("Enevironment para HIP_TORCH_PATH is required, please set the torch source dir such as: export HIP_TORCH_PATH=/path/to/pytorch ")
-        torch_lib_path = os.environ.get("TORCH_PATH", "")
-        if len(torch_lib_path)==0:
-            raise("Enevironment para TORCH_PATH is required, please set the torch source dir such as: export TORCH_PATH=/path/to/site-packages/torch ")
+        # CMake inherits TORCH_PATH from the torch package imported above.
+        torch_lib_path = torch_install_path
+        if not os.path.isdir(torch_lib_path):
+            raise RuntimeError(f"Torch installation directory does not exist: {torch_lib_path}")
         cmake_args = os.environ.get("TORCH_MOCKER_CMAKE_ARGS", "")
         make_args = "-j8 VERBOSE=1"
         make_args_env = os.environ.get("TORCH_MOCKER_MAKE_ARGS", "")
@@ -233,4 +233,3 @@ def run(self):
               },
     zip_safe=False,
 )
-
diff --git a/torch-mocker/test_py/CMakeLists.txt b/torch-mocker/test_py/CMakeLists.txt
index c09d0a3..d354f39 100644
--- a/torch-mocker/test_py/CMakeLists.txt
+++ b/torch-mocker/test_py/CMakeLists.txt
@@ -6,10 +6,10 @@ file(GLOB_RECURSE TEST_PYTHON_SRCS "*.cpp")
 add_library(test_THCP MODULE ${TEST_PYTHON_SRCS})
 target_compile_definitions(test_THCP PRIVATE "-DUSE_CUDA")
 target_link_libraries(test_THCP ${Python3_LIBRARIES})
-target_link_libraries(test_THCP $ENV{TORCH_PATH}/lib/libc10.so)
-target_link_libraries(test_THCP $ENV{TORCH_PATH}/lib/libtorch_cpu.so)
-target_link_libraries(test_THCP $ENV{TORCH_PATH}/lib/libtorch.so)
+target_link_libraries(test_THCP ${TORCH_PATH}/lib/libc10.so)
+target_link_libraries(test_THCP ${TORCH_PATH}/lib/libtorch_cpu.so)
+target_link_libraries(test_THCP ${TORCH_PATH}/lib/libtorch.so)
 target_link_libraries(test_THCP torch_cuda c10_cuda torch_mocker)
 target_link_libraries(test_THCP cufft cusparse cudnn cudnn_cnn_infer cudnn_adv_infer)
 target_link_libraries(test_THCP cublas cublasLt)
-set_target_properties(test_THCP PROPERTIES PREFIX "" SUFFIX ".so")
\ No newline at end of file
+set_target_properties(test_THCP PROPERTIES PREFIX "" SUFFIX ".so")
diff --git a/torch-mocker/test_py/test_python_c.cpp b/torch-mocker/test_py/test_python_c.cpp
deleted file mode 100644
index 2ab0158..0000000
--- a/torch-mocker/test_py/test_python_c.cpp
+++ /dev/null
@@ -1,72 +0,0 @@
-// Copyright (c) 2026 Hygon Information Technology Co., Ltd.
-// SPDX-License-Identifier: BSD-3-Clause
-
-// #include <pybind11/pybind11.h>
-// #include <torch/csrc/THP.h>
-// #include <torch/csrc/Stream.h>
-// #include <torch/csrc/cuda/Stream.h>
-// #include <torch/csrc/utils.h>
-
-// /*******注册THPUtils_PySequence_to_CUDAStreamList**********/
-// static PyObject* py_THPUtils_PySequence_to_CUDAStreamList(PyObject* self, PyObject* args) {
-//     PyObject* obj;
-//     if (!PyArg_ParseTuple(args, "O", &obj)) {
-//         return nullptr;
-//     }
-//     try {
-//         auto streams = THPUtils_PySequence_to_CUDAStreamList(obj);
-//         PyObject* py_list = PyList_New(streams.size());
-//         if (!py_list) {
-//             return nullptr;
-//         }
-//         for (size_t i = 0; i < streams.size(); ++i) {
-//             PyObject* py_stream = nullptr;
-//             if (streams[i].has_value()) {
-//                 py_stream = PyLong_FromVoidPtr(reinterpret_cast<void*>(streams[i]->stream()));
-//             } else {
-//                 Py_INCREF(Py_None);
-//                 py_stream = Py_None;
-//             }
-//             PyList_SET_ITEM(py_list, i, py_stream);
-//         }
-//         return py_list;
-//     } catch (const std::runtime_error& e) {
-//         PyErr_SetString(PyExc_RuntimeError, e.what());
-//         return nullptr;
-//     }
-// }
-// static PyMethodDef UtilMethods[] = {
-//     {"to_cuda_stream_list", (PyCFunction)py_THPUtils_PySequence_to_CUDAStreamList, METH_VARARGS, "Convert a Python sequence to CUDAStream list"},
-//     {nullptr, nullptr, 0, nullptr}
-// };
-
-// void utils_init(PyObject* module) {
-//     for (PyMethodDef* method = UtilMethods; method->ml_name != nullptr; ++method) {
-//         PyObject* func = PyCFunction_NewEx(method, nullptr, nullptr);
-//         if (func == nullptr || PyModule_AddObject(module, method->ml_name, func) < 0) {
-//             Py_XDECREF(func);
-//             throw python_error();
-//         }
-//     }
-// }
-
-// /*******注册THPUtils_PySequence_to_CUDAStreamList**********/
-
-// static struct PyModuleDef stream_module = {
-//     PyModuleDef_HEAD_INIT,
-//     "test_THCP",
-//     "Test THCP Module",
-//     -1,
-//     nullptr
-// };
-
-// PyMODINIT_FUNC PyInit_test_THCP(void) {
-//     PyObject* module = PyModule_Create(&stream_module);
-//     if (!module) {
-//         return nullptr;
-//     }
-//     THPStream_init(module);//先初始化THPStreamClass
-//     THCPStream_init(module);
-//     utils_init(module);
-//     return module;
-// }
\ No newline at end of file
diff --git a/torch-mocker/tests/CMakeLists.txt b/torch-mocker/tests/CMakeLists.txt
index b672ce6..cbd1a9b 100644
--- a/torch-mocker/tests/CMakeLists.txt
+++ b/torch-mocker/tests/CMakeLists.txt
@@ -9,10 +9,10 @@ set(TORCH_MOCKER_TEST_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}")
 function(torch_mocker_add_test_runner target)
     cuda_add_executable(${target} ${TORCH_MOCKER_TEST_SOURCE_DIR}/support/TestMain.cpp)
     target_link_libraries(${target}
-        $ENV{TORCH_PATH}/lib/libc10.so
-        $ENV{TORCH_PATH}/lib/libtorch_cpu.so
-        $ENV{TORCH_PATH}/lib/libtorch.so
-        $ENV{TORCH_PATH}/lib/libcaffe2_nvrtc.so
+        ${TORCH_PATH}/lib/libc10.so
+        ${TORCH_PATH}/lib/libtorch_cpu.so
+        ${TORCH_PATH}/lib/libtorch.so
+        ${TORCH_PATH}/lib/libcaffe2_nvrtc.so
         torch_cuda c10_cuda torch_mocker
         cufft cusparse cudnn cudnn_cnn_infer cudnn_adv_infer numa
         cublas nccl cublasLt pthread fmt)
diff --git a/torch-mocker/torch/csrc/cuda/Stream.cpp b/torch-mocker/torch/csrc/cuda/Stream.cpp
deleted file mode 100644
index 77f1452..0000000
--- a/torch-mocker/torch/csrc/cuda/Stream.cpp
+++ /dev/null
@@ -1,236 +0,0 @@
-// Copyright (c) 2026 Hygon Information Technology Co., Ltd.
-// SPDX-License-Identifier: BSD-3-Clause
-// Modified by Hygon Information Technology Co., Ltd., 2026.
-
-#include <pybind11/pybind11.h>
-#include <torch/csrc/Device.h>
-#include <torch/csrc/THP.h>
-#include <torch/csrc/cuda/Module.h>
-#include <torch/csrc/cuda/Stream.h>
-#include <torch/csrc/utils/pybind.h>
-#include <torch/csrc/utils/python_numbers.h>
-
-#ifndef USE_ROCM
-#include <c10/cuda/CUDAGuard.h>
-
-#include <cuda_runtime_api.h>
-#else
-#include <c10/hip/HIPGuard.h>
-
-#include <hip/hip_runtime_api.h>
-#endif
-#include <structmember.h>
-#include <utils/mocker_logging.h>
-
-PyObject* THCPStreamClass = nullptr;
-
-static PyObject* THCPStream_pynew(
-    PyTypeObject* type,
-    PyObject* args,
-    PyObject* kwargs) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-
-  const auto current_device = c10::cuda::current_device();
-
-  int priority = 0;
-  int64_t stream_id = 0;
-  int64_t device_index = 0;
-  int64_t device_type = 0;
-  uint64_t stream_ptr = 0;
-
-  // NOLINTNEXTLINE(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays)
-  constexpr const char* kwlist[] = {
-      "priority",
-      "stream_id",
-      "device_index",
-      "device_type",
-      "stream_ptr",
-      nullptr};
-  if (!PyArg_ParseTupleAndKeywords(
-          args,
-          kwargs,
-          "|iLLLK",
-          // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
-          const_cast<char**>(kwlist),
-          &priority,
-          &stream_id,
-          &device_index,
-          &device_type,
-          &stream_ptr)) {
-    return nullptr;
-  }
-
-  THPObjectPtr ptr(type->tp_alloc(type, 0));
-  if (!ptr) {
-    return nullptr;
-  }
-
-  if (stream_ptr) {
-    TORCH_CHECK(
-        priority == 0, "Priority was explicitly set for a external stream")
-  }
-  at::cuda::CUDAStream stream = (stream_id || device_index || device_type)
-      ? at::cuda::CUDAStream::unpack3(
-            stream_id,
-            static_cast<c10::DeviceIndex>(device_index),
-            static_cast<c10::DeviceType>(device_type))
-      : stream_ptr ? at::cuda::getStreamFromExternal(
-                         // NOLINTNEXTLINE(performance-no-int-to-ptr)
-#ifndef USE_ROCM
-                         reinterpret_cast<cudaStream_t>(stream_ptr),
-#else
-                         reinterpret_cast<hipStream_t>(stream_ptr),
-#endif
-                         current_device)
-                   : at::cuda::getStreamFromPool(priority);
-
-  THCPStream* self = (THCPStream*)ptr.get();
-  self->stream_id = static_cast<int64_t>(stream.id());
-  // NOLINTNEXTLINE(bugprone-signed-char-misuse)
-  self->device_index = static_cast<int64_t>(stream.device_index());
-  self->device_type = static_cast<int64_t>(stream.device_type());
-  new (&self->cuda_stream) at::cuda::CUDAStream(stream);
-
-  return (PyObject*)ptr.release();
-  END_HANDLE_TH_ERRORS
-}
-
-static void THCPStream_dealloc(THCPStream* self) {
-  if (Vlog()) LOG_INFO();
-  self->cuda_stream.~CUDAStream();
-  Py_TYPE(self)->tp_free((PyObject*)self);
-}
-
-static PyObject* THCPStream_get_cuda_stream(THCPStream* self, void* unused) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-  return PyLong_FromVoidPtr(self->cuda_stream.stream());
-  END_HANDLE_TH_ERRORS
-}
-
-static PyObject* THCPStream_get_priority(THCPStream* self, void* unused) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-  return THPUtils_packInt64(self->cuda_stream.priority());
-  END_HANDLE_TH_ERRORS
-}
-
-static PyObject* THCPStream_priority_range(
-    PyObject* _unused,
-    PyObject* noargs) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-  auto [least_priority, greatest_priority] =
-      at::cuda::CUDAStream::priority_range();
-  return Py_BuildValue("(ii)", least_priority, greatest_priority);
-  END_HANDLE_TH_ERRORS
-}
-
-static PyObject* THCPStream_query(PyObject* _self, PyObject* noargs) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-  auto self = (THCPStream*)_self;
-  return PyBool_FromLong(self->cuda_stream.query());
-  END_HANDLE_TH_ERRORS
-}
-
-static PyObject* THCPStream_synchronize(PyObject* _self, PyObject* noargs) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS {
-    pybind11::gil_scoped_release no_gil;
-    auto self = (THCPStream*)_self;
-    self->cuda_stream.synchronize();
-  }
-  Py_RETURN_NONE;
-  END_HANDLE_TH_ERRORS
-}
-
-static PyObject* THCPStream_eq(PyObject* _self, PyObject* _other) {
-  if (Vlog()) LOG_INFO();
-  HANDLE_TH_ERRORS
-  auto self = (THCPStream*)_self;
-  auto other = (THCPStream*)_other;
-  return PyBool_FromLong(self->cuda_stream == other->cuda_stream);
-  END_HANDLE_TH_ERRORS
-}
-
-// NOLINTNEXTLINE(*-c-arrays*, *-global-variables)
-static struct PyMemberDef THCPStream_members[] = {{nullptr}};
-
-// NOLINTNEXTLINE(*-c-arrays*, *-global-variables)
-static struct PyGetSetDef THCPStream_properties[] = {
-    {"cuda_stream",
-     (getter)THCPStream_get_cuda_stream,
-     nullptr,
-     nullptr,
-     nullptr},
-    {"priority", (getter)THCPStream_get_priority, nullptr, nullptr, nullptr},
-    {nullptr}};
-
-// NOLINTNEXTLINE(*-c-arrays*, *-global-variables)
-static PyMethodDef THCPStream_methods[] = {
-    {"query", THCPStream_query, METH_NOARGS, nullptr},
-    {"synchronize", THCPStream_synchronize, METH_NOARGS, nullptr},
-    {"priority_range",
-     THCPStream_priority_range,
-     METH_STATIC | METH_NOARGS,
-     nullptr},
-    {"__eq__", THCPStream_eq, METH_O, nullptr},
-    {nullptr}};
-
-PyTypeObject THCPStreamType = {
-    PyVarObject_HEAD_INIT(nullptr, 0)
-    "torch._C._CudaStreamBase", /* tp_name */
-    sizeof(THCPStream), /* tp_basicsize */
-    0, /* tp_itemsize */
-    (destructor)THCPStream_dealloc, /* tp_dealloc */
-    0, /* tp_vectorcall_offset */
-    nullptr, /* tp_getattr */
-    nullptr, /* tp_setattr */
-    nullptr, /* tp_reserved */
-    nullptr, /* tp_repr */
-    nullptr, /* tp_as_number */
-    nullptr, /* tp_as_sequence */
-    nullptr, /* tp_as_mapping */
-    nullptr, /* tp_hash  */
-    nullptr, /* tp_call */
-    nullptr, /* tp_str */
-    nullptr, /* tp_getattro */
-    nullptr, /* tp_setattro */
-    nullptr, /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
-    nullptr, /* tp_doc */
-    nullptr, /* tp_traverse */
-    nullptr, /* tp_clear */
-    nullptr, /* tp_richcompare */
-    0, /* tp_weaklistoffset */
-    nullptr, /* tp_iter */
-    nullptr, /* tp_iternext */
-    THCPStream_methods, /* tp_methods */
-    THCPStream_members, /* tp_members */
-    THCPStream_properties, /* tp_getset */
-    nullptr, /* tp_base */
-    nullptr, /* tp_dict */
-    nullptr, /* tp_descr_get */
-    nullptr, /* tp_descr_set */
-    0, /* tp_dictoffset */
-    nullptr, /* tp_init */
-    nullptr, /* tp_alloc */
-    THCPStream_pynew, /* tp_new */
-};
-
-void THCPStream_init(PyObject* module) {
-  if (Vlog()) LOG_INFO();
-  Py_INCREF(THPStreamClass);
-  THCPStreamType.tp_base = THPStreamClass;
-  THCPStreamClass = (PyObject*)&THCPStreamType;
-  if (PyType_Ready(&THCPStreamType) < 0) {
-    throw python_error();
-  }
-  Py_INCREF(&THCPStreamType);
-  if (PyModule_AddObject(
-          module, "_CudaStreamBase", (PyObject*)&THCPStreamType) < 0) {
-    throw python_error();
-  }
-}
diff --git a/torch-mocker/torch/share/cmake/Caffe2/Caffe2Targets.cmake b/torch-mocker/torch/share/cmake/Caffe2/Caffe2Targets.cmake
index e2a005d..37acff3 100644
--- a/torch-mocker/torch/share/cmake/Caffe2/Caffe2Targets.cmake
+++ b/torch-mocker/torch/share/cmake/Caffe2/Caffe2Targets.cmake
@@ -147,7 +147,7 @@ add_library(torch_cuda SHARED IMPORTED)
 set_target_properties(torch_cuda PROPERTIES
   INTERFACE_COMPILE_DEFINITIONS "USE_C10D_NCCL"
   INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include"
-  INTERFACE_LINK_LIBRARIES "c10_cuda;torch_cpu_library;${_IMPORT_PREFIX}/lib/libtorch_python.so"
+  INTERFACE_LINK_LIBRARIES "c10_cuda;torch_cpu_library"
 )
 
 add_library(torch_hip_library INTERFACE IMPORTED)