From 8d788497c086a4d090879124f926a09bea58671b Mon Sep 17 00:00:00 2001
From: gongchensu <zhuyue_134@qq.com>
Date: Tue, 14 Jul 2026 09:08:21 +0000
Subject: [PATCH] feat(metax): add graph capture and replay support

---
 docs/backends.md                 |  2 +-
 src/native/cuda/metax/runtime_.h | 42 +++++++++++++++++++++++++++++++-
 tests/CMakeLists.txt             |  2 ++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/docs/backends.md b/docs/backends.md
index 26ad1ff..8ecb8f8 100644
--- a/docs/backends.md
+++ b/docs/backends.md
@@ -33,7 +33,7 @@ Current test expectations are:
 | CPU | No | Yes | No | Yes | No | Yes | No |
 | NVIDIA | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | Iluvatar | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
-| MetaX | Yes | Yes | Yes | Yes | Yes | Yes | No |
+| MetaX | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | Moore | Yes | Yes | No | Yes | Yes | Yes | No |
 | Hygon | Yes | Yes | Yes | Yes | Yes | Yes | No |
 | Cambricon | Yes | No | No | No | No | No | No |
diff --git a/src/native/cuda/metax/runtime_.h b/src/native/cuda/metax/runtime_.h
index 462f5f1..63d4d24 100644
--- a/src/native/cuda/metax/runtime_.h
+++ b/src/native/cuda/metax/runtime_.h
@@ -12,13 +12,20 @@ namespace infini::rt::runtime {
 
 template <>
 struct Runtime<Device::Type::kMetax>
-    : CudaRuntime<Runtime<Device::Type::kMetax>> {
+    : GraphRuntime<Runtime<Device::Type::kMetax>,
+                   CudaRuntime<Runtime<Device::Type::kMetax>>> {
   using Error = mcError_t;
 
   using Stream = mcStream_t;
 
+  using Graph = mcGraph_t;
+
+  using GraphExec = mcGraphExec_t;
+
   using Event = mcEvent_t;
 
+  using StreamCaptureMode = mcStreamCaptureMode;
+
   static constexpr Device::Type kDeviceType = Device::Type::kMetax;
 
   static constexpr Error kSuccess = mcSuccess;
@@ -124,6 +131,39 @@ struct Runtime<Device::Type::kMetax>
   static constexpr auto EventElapsedTime = [](auto&&... args) {
     return mcEventElapsedTime(std::forward<decltype(args)>(args)...);
   };
+
+  static constexpr auto kStreamCaptureModeGlobal =
+      mcStreamCaptureModeGlobal;
+
+  static constexpr auto kStreamCaptureModeThreadLocal =
+      mcStreamCaptureModeThreadLocal;
+
+  static constexpr auto kStreamCaptureModeRelaxed =
+      mcStreamCaptureModeRelaxed;
+
+  static constexpr auto StreamBeginCapture = [](auto&&... args) {
+    return mcStreamBeginCapture(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto StreamEndCapture = [](auto&&... args) {
+    return mcStreamEndCapture(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto GraphDestroy = [](auto&&... args) {
+    return mcGraphDestroy(std::forward<decltype(args)>(args)...);
+  };
+
+  static Error GraphInstantiate(GraphExec* graph_exec, Graph graph) {
+    return mcGraphInstantiate(graph_exec, graph, nullptr, nullptr, 0);
+  }
+
+  static constexpr auto GraphExecDestroy = [](auto&&... args) {
+    return mcGraphExecDestroy(std::forward<decltype(args)>(args)...);
+  };
+
+  static constexpr auto GraphLaunch = [](auto&&... args) {
+    return mcGraphLaunch(std::forward<decltype(args)>(args)...);
+  };
 };
 
 static_assert(Runtime<Device::Type::kMetax>::Validate());
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 1dd2048..42cdf6b 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -85,6 +85,8 @@ if(WITH_METAX)
         METAX infini::rt::Device::Type::kMetax
         infini/rt/metax/runtime_.h
         1 1 1 1 1 1 1 1)
+    add_infini_rt_backend_graph_test(
+        METAX infini::rt::Device::Type::kMetax 1)
 endif()
 
 if(WITH_MOORE)