InfiniTensor · voltjia · Aug 11, 2026 · Aug 11, 2026
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -276,6 +276,14 @@ if(WITH_TORCH OR WITH_LINKED)
             -Wl,--no-as-needed ${TORCH_CUDA_LIB} ${C10_CUDA_LIB} -Wl,--as-needed)
     endif()
 
+    if(WITH_METAX)
+        find_library(C10_CUDA_LIB c10_cuda HINTS ${_torch_lib_dirs} REQUIRED)
+        find_library(MACA_TORCH_RUNTIME_LIB runtime_cu
+            HINTS "$ENV{MACA_PATH}/lib" REQUIRED)
+        list(APPEND TORCH_LIBRARIES
+            ${C10_CUDA_LIB} ${MACA_TORCH_RUNTIME_LIB})
+    endif()
+
     # `auditwheel`-repaired `torch` wheels bundle transitive dependencies
     # (e.g. `libgfortran-<hash>.so`, `libopenblasp-<hash>.so`) in a sibling
     # `torch.libs/` directory that `library_paths()` does not return.  When

diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py
@@ -74,6 +74,13 @@
     "kHygon",
 )
 
+_C10_DEVICE_TYPES = (
+    "kNvidia",
+    "kCambricon",
+    "kMetax",
+    "kMoore",
+)
+
 # YAML scalar-type tokens → C++ types. Reference types (e.g. `const Scalar&`)
 # are not used so the generated signatures match the existing hand-written
 # ones, which pass by value to keep pybind11 binding generation simple.
@@ -1270,7 +1277,6 @@ def _generate_torch_method_source(name: str, op: Op) -> str:
     op_type = _op_cpp_type(name)
     conversion_lines = []
     out_device_index = f"{op.out_params[0].api_name}.device().index()"
-    conversion_lines.append(f"  const auto device_index = {out_device_index};")
 
     def _optional_aten_type(param: Param) -> str:
         return _NULLOPT_BY_TYPE[param.aten_type].removesuffix("{}")
@@ -1446,6 +1452,10 @@ def _render_arg(schema_index, p):
         op_type=op_type,
         op_call_signature=_format_signature(op),
         tensor_conversions="\n".join(conversion_lines),
+        out_device_index=out_device_index,
+        c10_condition=" ||\n      ".join(
+            f"kDev == Device::Type::{dev}" for dev in _C10_DEVICE_TYPES
+        ),
         # The generated call expression resolves the right kernel via C++
         # overload resolution from the argument types we pass.
         aten_call=aten_call,
@@ -1456,6 +1466,12 @@ def _render_arg(schema_index, p):
 def _generate_torch_source(name: str, ops: list[Op]) -> str:
     op_type = _op_cpp_type(name)
     methods = "\n\n".join(_generate_torch_method_source(name, op) for op in ops)
+    c10_includes = "\n".join(
+        f"#ifdef WITH_{dev.removeprefix('k').upper()}\n"
+        f'#include "torch/{dev.removeprefix("k").lower()}/c10.h"\n'
+        f"#endif"
+        for dev in _C10_DEVICE_TYPES
+    )
     # Guard each explicit instantiation by the matching `WITH_<DEV>` macro
     # so a build that only enables a subset of devices does not pay the
     # ATen template-instantiation cost (and memory pressure) for the
@@ -1471,6 +1487,7 @@ def _generate_torch_source(name: str, ops: list[Op]) -> str:
     return _TORCH_SOURCE_TEMPLATE.format(
         name=name,
         methods=methods,
+        c10_includes=c10_includes,
         instantiations=instantiations,
     )
 
@@ -1524,16 +1541,30 @@ class Operator<{op_type}, kDev, {slot}> : public {op_type} {{
 _TORCH_METHOD_TEMPLATE = """\
 template <Device::Type kDev>
 void Operator<{op_type}, kDev, {slot}>::operator()({op_call_signature}) const {{
+  const auto device_index = {out_device_index};
+  const auto run = [&] {{
 {tensor_conversions}
 
-  {aten_call};
+    {aten_call};
+  }};
+
+  if constexpr ({c10_condition}) {{
+    const typename C10<kDev>::StreamGuard stream_guard{{
+        C10<kDev>::GetStreamFromExternal(stream_, device_index)}};
+    run();
+    return;
+  }}
+
+  run();
 }}
 """
 
 
 _TORCH_SOURCE_TEMPLATE = """\
 #include "torch/{name}/{name}.h"
 
+#include "torch/c10.h"
+{c10_includes}
 #include "torch/tensor_.h"
 
 namespace infini::ops {{

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -693,7 +693,7 @@ if(_infini_ops_linked_uses_torch)
     list(APPEND TORCH_SOURCES ${INFINI_OPS_LINKED_SOURCES})
 endif()
 
-if(WITH_CAMBRICON AND _infini_ops_linked_uses_torch)
+if(WITH_CAMBRICON AND TORCH_SOURCES)
     execute_process(
         COMMAND "${_TORCH_PYTHON}" -c
             "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-mlu').locate_file('')).resolve(); include = root / 'torch_mlu/csrc'; lib = include / 'lib'; print(include, lib, sep=';')"
@@ -808,32 +808,39 @@ if(TORCH_SOURCES)
 
         set(_torch_vendor_include_flags "")
         if(WITH_METAX)
-            list(APPEND _torch_vendor_include_flags "-I${MACA_PATH}/include")
+            list(APPEND _torch_vendor_include_flags
+                "-I${MACA_PATH}/include"
+                "-I${MACA_PATH}/include/mcr"
+                "-I${MACA_PATH}/tools/cu-bridge/include")
         endif()
         if(WITH_MOORE)
             list(APPEND _torch_vendor_include_flags "-I${MUSA_ROOT}/include")
-            if(WITH_LINKED AND _infini_ops_linked_uses_torch)
-                execute_process(
-                    COMMAND "${_TORCH_PYTHON}" -c
-                        "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-musa').locate_file('')).resolve(); compat = root / 'torch_musa/share/generated_cuda_compatible/include'; print(root, compat, sep=';')"
-                    OUTPUT_VARIABLE _torch_musa_include_dirs
-                    OUTPUT_STRIP_TRAILING_WHITESPACE
-                    RESULT_VARIABLE _torch_musa_include_result
-                )
-                if(NOT _torch_musa_include_result EQUAL 0 OR
-                   NOT _torch_musa_include_dirs)
+            execute_process(
+                COMMAND "${_TORCH_PYTHON}" -c
+                    "import importlib.metadata, pathlib; root = pathlib.Path(importlib.metadata.distribution('torch-musa').locate_file('')).resolve(); compat = root / 'torch_musa/share/generated_cuda_compatible/include'; lib = root / 'torch_musa/lib'; print(root, compat, lib, sep=';')"
+                OUTPUT_VARIABLE _torch_musa_paths
+                OUTPUT_STRIP_TRAILING_WHITESPACE
+                RESULT_VARIABLE _torch_musa_paths_result
+            )
+            if(NOT _torch_musa_paths_result EQUAL 0 OR NOT _torch_musa_paths)
+                message(FATAL_ERROR
+                    "Could not determine `torch_musa` include and library directories.")
+            endif()
+            list(GET _torch_musa_paths 2 _torch_musa_library_dir)
+            list(REMOVE_AT _torch_musa_paths 2)
+            foreach(_dir IN LISTS _torch_musa_paths)
+                if(NOT IS_DIRECTORY "${_dir}")
                     message(FATAL_ERROR
-                        "Could not determine `torch_musa` include directories.")
+                        "Could not find `torch_musa` include directory `${_dir}`.")
                 endif()
-                foreach(_dir IN LISTS _torch_musa_include_dirs)
-                    if(NOT IS_DIRECTORY "${_dir}")
-                        message(FATAL_ERROR
-                            "Could not find `torch_musa` include directory `${_dir}`.")
-                    endif()
-                    list(APPEND _torch_include_flags
-                        "-isystem" "${_dir}")
-                endforeach()
-            endif()
+                list(APPEND _torch_include_flags "-isystem" "${_dir}")
+            endforeach()
+
+            find_library(TORCH_MUSA_LIB musa_python
+                HINTS "${_torch_musa_library_dir}"
+                NO_DEFAULT_PATH REQUIRED)
+            target_link_libraries(infiniops PUBLIC "${TORCH_MUSA_LIB}")
+            list(APPEND TORCH_RUNTIME_DIRS "${_torch_musa_library_dir}")
         endif()
 
         # Vendor-specific defines required by forked `torch` headers.

diff --git a/src/linked/torch/c10.h b/src/linked/torch/c10.h
diff --git a/src/linked/torch/cambricon/ops/rms_norm/apex.h b/src/linked/torch/cambricon/ops/rms_norm/apex.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_RMS_NORM_APEX_H_
 #define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_RMS_NORM_APEX_H_
 
-#include "linked/torch/cambricon/c10.h"
 #include "linked/torch/ops/rms_norm.h"
+#include "torch/cambricon/c10.h"
 
 namespace infini::ops::linked::torch::cambricon {
 

diff --git a/src/linked/torch/cambricon/ops/silu_and_mul/apex.h b/src/linked/torch/cambricon/ops/silu_and_mul/apex.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_SILU_AND_MUL_APEX_H_
 #define INFINI_OPS_LINKED_TORCH_CAMBRICON_OPS_SILU_AND_MUL_APEX_H_
 
-#include "linked/torch/cambricon/c10.h"
 #include "linked/torch/ops/silu_and_mul.h"
+#include "torch/cambricon/c10.h"
 
 namespace infini::ops::linked::torch::cambricon {
 

diff --git a/src/linked/torch/metax/ops/rms_norm/vllm.h b/src/linked/torch/metax/ops/rms_norm/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_RMS_NORM_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_METAX_OPS_RMS_NORM_VLLM_H_
 
-#include "linked/torch/metax/c10.h"
 #include "linked/torch/ops/rms_norm.h"
+#include "torch/metax/c10.h"
 
 namespace infini::ops::linked::torch::metax {
 

diff --git a/src/linked/torch/metax/ops/silu_and_mul/vllm.h b/src/linked/torch/metax/ops/silu_and_mul/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_METAX_OPS_SILU_AND_MUL_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_METAX_OPS_SILU_AND_MUL_VLLM_H_
 
-#include "linked/torch/metax/c10.h"
 #include "linked/torch/ops/silu_and_mul.h"
+#include "torch/metax/c10.h"
 
 namespace infini::ops::linked::torch::metax {
 

diff --git a/src/linked/torch/moore/ops/rms_norm/vllm.h b/src/linked/torch/moore/ops/rms_norm/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_RMS_NORM_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_MOORE_OPS_RMS_NORM_VLLM_H_
 
-#include "linked/torch/moore/c10.h"
 #include "linked/torch/ops/rms_norm.h"
+#include "torch/moore/c10.h"
 
 namespace infini::ops::linked::torch::moore {
 

diff --git a/src/linked/torch/moore/ops/silu_and_mul/vllm.h b/src/linked/torch/moore/ops/silu_and_mul/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_MOORE_OPS_SILU_AND_MUL_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_MOORE_OPS_SILU_AND_MUL_VLLM_H_
 
-#include "linked/torch/moore/c10.h"
 #include "linked/torch/ops/silu_and_mul.h"
+#include "torch/moore/c10.h"
 
 namespace infini::ops::linked::torch::moore {
 

diff --git a/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/awq_marlin_repack/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_AWQ_MARLIN_REPACK_VLLM_H_
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/awq_marlin_repack.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_varlen_func/flash_attn.h
@@ -3,8 +3,8 @@
 
 #include <ATen/core/Generator.h>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/flash_attn_varlen_func.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h b/src/linked/torch/nvidia/ops/flash_attn_with_kvcache/flash_attn.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
 #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_FLASH_ATTN_WITH_KVCACHE_FLASH_ATTN_H_
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/flash_attn_with_kvcache.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h b/src/linked/torch/nvidia/ops/get_cutlass_moe_mm_data/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GET_CUTLASS_MOE_MM_DATA_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GET_CUTLASS_MOE_MM_DATA_VLLM_H_
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/get_cutlass_moe_mm_data.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h b/src/linked/torch/nvidia/ops/gptq_marlin_repack/vllm.h
@@ -1,8 +1,8 @@
 #ifndef INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_
 #define INFINI_OPS_LINKED_TORCH_NVIDIA_OPS_GPTQ_MARLIN_REPACK_VLLM_H_
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/gptq_marlin_repack.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/grouped_topk/vllm.h b/src/linked/torch/nvidia/ops/grouped_topk/vllm.h
@@ -3,8 +3,8 @@
 
 #include <utility>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/grouped_topk.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_gemm/vllm.h
@@ -3,8 +3,8 @@
 
 #include <optional>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/moe_wna16_gemm.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h b/src/linked/torch/nvidia/ops/moe_wna16_marlin_gemm/vllm.h
@@ -8,8 +8,8 @@
 #include <optional>
 #include <utility>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/moe_wna16_marlin_gemm.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h b/src/linked/torch/nvidia/ops/topk_sigmoid/vllm.h
@@ -3,8 +3,8 @@
 
 #include <optional>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/topk_sigmoid.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/linked/torch/nvidia/ops/topk_softmax/vllm.h b/src/linked/torch/nvidia/ops/topk_softmax/vllm.h
@@ -3,8 +3,8 @@
 
 #include <optional>
 
-#include "linked/torch/nvidia/c10.h"
 #include "linked/torch/ops/topk_softmax.h"
+#include "torch/nvidia/c10.h"
 
 namespace infini::ops::linked::torch::nvidia {
 

diff --git a/src/torch/c10.h b/src/torch/c10.h
@@ -0,0 +1,13 @@
+#ifndef INFINI_OPS_TORCH_C10_H_
+#define INFINI_OPS_TORCH_C10_H_
+
+#include "device.h"
+
+namespace infini::ops {
+
+template <Device::Type kDev>
+struct C10;
+
+}  // namespace infini::ops
+
+#endif  // INFINI_OPS_TORCH_C10_H_
diff --git a/src/linked/torch/cambricon/c10.h → src/torch/cambricon/c10.h b/src/linked/torch/cambricon/c10.h → src/torch/cambricon/c10.h
@@ -1,13 +1,13 @@
-#ifndef INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_
-#define INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_
+#ifndef INFINI_OPS_TORCH_CAMBRICON_C10_H_
+#define INFINI_OPS_TORCH_CAMBRICON_C10_H_
 
 #include <c10/core/StreamGuard.h>
 #include <cnrt.h>
 #include <framework/core/MLUStream.h>
 
-#include "linked/torch/c10.h"
+#include "torch/c10.h"
 
-namespace infini::ops::linked::torch {
+namespace infini::ops {
 
 template <>
 struct C10<Device::Type::kCambricon> {
@@ -22,6 +22,6 @@ struct C10<Device::Type::kCambricon> {
   }
 };
 
-}  // namespace infini::ops::linked::torch
+}  // namespace infini::ops
 
-#endif  // INFINI_OPS_LINKED_TORCH_CAMBRICON_C10_H_
+#endif  // INFINI_OPS_TORCH_CAMBRICON_C10_H_