Support intranode DeepEP without IMEX fabric

csrc/BinaryGraphIO.cpp

@@ -12,6 +12,46 @@
 
 namespace foundry {
 
+namespace {
+
+bool validate_binary_layout(const uint8_t* data, size_t file_size,
+                            const binary_format::FileHeader& header) {
+  namespace bf = binary_format;
+
+  if (header.num_sections != bf::SECTION_COUNT) {
+    return false;
+  }
+
+  const uint64_t directory_size =
+      static_cast<uint64_t>(header.num_sections) * sizeof(bf::SectionEntry);
+  if (directory_size > file_size - sizeof(bf::FileHeader)) {
+    return false;
+  }
+
+  const auto* sections = reinterpret_cast<const bf::SectionEntry*>(data + sizeof(bf::FileHeader));
+  for (uint32_t i = 0; i < header.num_sections; ++i) {
+    const auto& section = sections[i];
+    if (section.section_type != i || section.offset > file_size ||
+        section.size > file_size - section.offset) {
+      return false;
+    }
+  }
+
+  const uint64_t legacy_dependency_size =
+      static_cast<uint64_t>(header.num_dependencies) * sizeof(bf::LegacyBinDependency);
+  const uint64_t dependency_size =
+      static_cast<uint64_t>(header.num_dependencies) * sizeof(bf::BinDependency);
+  const uint64_t stored_dependency_size = sections[bf::SECTION_DEPENDENCY_TABLE].size;
+  const bool has_dependencies = (header.flags & bf::FLAG_HAS_DEPENDENCIES) != 0;
+  if (has_dependencies != (header.num_dependencies != 0)) {
+    return false;
+  }
+  return stored_dependency_size == legacy_dependency_size ||
+         stored_dependency_size == dependency_size;
+}
+
+}  // namespace
+
 void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::object& root) {
   namespace json = boost::json;
   namespace bf = binary_format;
@@ -270,7 +310,19 @@ void CUDAGraph::save_binary(const std::string& bin_path, const boost::json::obje
   deps.reserve(deps_array.size());
   for (const auto& dep_val : deps_array) {
     const json::object& d = dep_val.as_object();
-    deps.push_back({d.at("from").to_number<uint32_t>(), d.at("to").to_number<uint32_t>()});
+    bf::BinDependency dep = {};
+    dep.from_id = d.at("from").to_number<uint32_t>();
+    dep.to_id = d.at("to").to_number<uint32_t>();
+    if (d.contains("from_port")) {
+      dep.from_port = static_cast<uint8_t>(d.at("from_port").to_number<unsigned int>());
+    }
+    if (d.contains("to_port")) {
+      dep.to_port = static_cast<uint8_t>(d.at("to_port").to_number<unsigned int>());
+    }
+    if (d.contains("type")) {
+      dep.type = static_cast<uint8_t>(d.at("type").to_number<unsigned int>());
+    }
+    deps.push_back(dep);
   }
 
   // ---- Generators ----
@@ -397,7 +449,7 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) {
   if (file_size < sizeof(bf::FileHeader))
     return json::value{};
   const bf::FileHeader& header = *reinterpret_cast<const bf::FileHeader*>(data);
-  if (!bf::validate_header(header))
+  if (!bf::validate_header(header) || !validate_binary_layout(data, file_size, header))
     return json::value{};
 
   // Read section directory
@@ -653,13 +705,26 @@ boost::json::value read_and_parse_binary_graph(const std::string& bin_path) {
   // Dependencies
   if (header.flags & bf::FLAG_HAS_DEPENDENCIES) {
     auto [dep_data, dep_size] = section(bf::SECTION_DEPENDENCY_TABLE);
-    const bf::BinDependency* deps = reinterpret_cast<const bf::BinDependency*>(dep_data);
+    const size_t dependency_entry_size = dep_size / header.num_dependencies;
+    const bool has_edge_data = dependency_entry_size == sizeof(bf::BinDependency);
     json::array deps_array;
     deps_array.reserve(header.num_dependencies);
     for (uint32_t i = 0; i < header.num_dependencies; i++) {
+      const uint8_t* record = dep_data + i * dependency_entry_size;
+      const auto* legacy = reinterpret_cast<const bf::LegacyBinDependency*>(record);
       json::object d;
-      d["from"] = deps[i].from_id;
-      d["to"] = deps[i].to_id;
+      d["from"] = legacy->from_id;
+      d["to"] = legacy->to_id;
+      if (has_edge_data) {
+        const auto* dep = reinterpret_cast<const bf::BinDependency*>(record);
+        d["from_port"] = static_cast<unsigned int>(dep->from_port);
+        d["to_port"] = static_cast<unsigned int>(dep->to_port);
+        d["type"] = static_cast<unsigned int>(dep->type);
+      } else {
+        d["from_port"] = 0;
+        d["to_port"] = 0;
+        d["type"] = 0;
+      }
       deps_array.push_back(d);
     }
     root["dependencies"] = std::move(deps_array);
@@ -725,7 +790,8 @@ BinaryGraphFile read_binary_graph_file(const std::string& bin_path) {
   in.close();
 
   memcpy(&result.header, result.data.data(), sizeof(bf::FileHeader));
-  if (!bf::validate_header(result.header)) {
+  if (!bf::validate_header(result.header) ||
+      !validate_binary_layout(result.data.data(), file_size, result.header)) {
     result.data.clear();
     return result;
   }

csrc/CUDAGraph.cpp

@@ -16,6 +16,8 @@
 #include <cstddef>
 #include <cstring>
 #include <iostream>
+#include <map>
+#include <tuple>
 #include <unordered_map>
 #include <cuda.h>
 #include <fstream>
@@ -127,6 +129,47 @@ MempoolId_t torch_graph_pool_handle(bool is_user_created = true) {
 
 }  // namespace
 
+CUresult add_graph_dependencies(CUgraph graph, const std::vector<CUgraphNode>& from_nodes,
+                                const std::vector<CUgraphNode>& to_nodes,
+                                const std::vector<CUgraphEdgeData>& edge_data) {
+  TORCH_INTERNAL_ASSERT(from_nodes.size() == to_nodes.size());
+  TORCH_INTERNAL_ASSERT(from_nodes.size() == edge_data.size());
+
+  struct DependencyBatch {
+    std::vector<CUgraphNode> from_nodes;
+    std::vector<CUgraphNode> to_nodes;
+    std::vector<CUgraphEdgeData> edge_data;
+  };
+  using EdgeDataKey = std::tuple<unsigned char, unsigned char, unsigned char>;
+  std::map<EdgeDataKey, DependencyBatch> batches;
+
+  // CUDA 13.0 silently normalizes non-default entries when one call mixes edge annotations.
+  // Group identical annotations so each driver call remains lossless without adding edges singly.
+  for (size_t index = 0; index < edge_data.size(); ++index) {
+    const auto& data = edge_data[index];
+    auto& batch = batches[{data.from_port, data.to_port, data.type}];
+    batch.from_nodes.push_back(from_nodes[index]);
+    batch.to_nodes.push_back(to_nodes[index]);
+    batch.edge_data.push_back(data);
+  }
+
+  for (auto& [key, batch] : batches) {
+    (void)key;
+#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000)
+    CUresult result = cuGraphAddDependencies(graph, batch.from_nodes.data(), batch.to_nodes.data(),
+                                             batch.edge_data.data(), batch.edge_data.size());
+#else
+    CUresult result =
+        cuGraphAddDependencies_v2(graph, batch.from_nodes.data(), batch.to_nodes.data(),
+                                  batch.edge_data.data(), batch.edge_data.size());
+#endif
+    if (result != CUDA_SUCCESS) {
+      return result;
+    }
+  }
+  return CUDA_SUCCESS;
+}
+
 static bool _cuda_graphs_debug = false;
 
 static CUDAGeneratorStateRegistry global_generator_state_registry;
@@ -934,6 +977,7 @@ void CUDAGraph::analyze_captured_graph() {
         GraphDependency dep;
         dep.from_index = from_it->second;
         dep.to_index = to_it->second;
+        dep.edge_data = edges[i];
         graph_dependencies.push_back(dep);
       }
     }
@@ -1335,7 +1379,8 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t
     nodes_array.push_back(node_obj);
   }
 
-  // Compute topology key = node types + cluster dim values per kernel node.
+  // Compute topology key = node types + cluster dim values per kernel node +
+  // dependency endpoints and edge data.
   // cuGraphExecUpdate does NOT propagate kernel node attribute changes
   // (set via cuGraphKernelNodeSetAttribute) to the CUgraphExec — only
   // CUDA_KERNEL_NODE_PARAMS fields are updated. This means cluster dim
@@ -1381,6 +1426,22 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t
         }
       }
     }
+
+    std::vector<GraphDependency> sorted_dependencies = graph_dependencies;
+    std::sort(sorted_dependencies.begin(), sorted_dependencies.end(),
+              [](const GraphDependency& lhs, const GraphDependency& rhs) {
+                return std::make_tuple(lhs.from_index, lhs.to_index, lhs.edge_data.from_port,
+                                       lhs.edge_data.to_port, lhs.edge_data.type) <
+                       std::make_tuple(rhs.from_index, rhs.to_index, rhs.edge_data.from_port,
+                                       rhs.edge_data.to_port, rhs.edge_data.type);
+              });
+    topology_key += "|D";
+    for (const auto& dep : sorted_dependencies) {
+      topology_key += ";" + std::to_string(dep.from_index) + ">" + std::to_string(dep.to_index) +
+                      ":" + std::to_string(static_cast<int>(dep.edge_data.from_port)) + ":" +
+                      std::to_string(static_cast<int>(dep.edge_data.to_port)) + ":" +
+                      std::to_string(static_cast<int>(dep.edge_data.type));
+    }
     root["topology_key"] = topology_key;
   }
 
@@ -1447,6 +1508,9 @@ void CUDAGraph::save(const std::string& json_path, const OutputTensors& output_t
     json::object dep_obj;
     dep_obj["from"] = dep.from_index;
     dep_obj["to"] = dep.to_index;
+    dep_obj["from_port"] = static_cast<int>(dep.edge_data.from_port);
+    dep_obj["to_port"] = static_cast<int>(dep.edge_data.to_port);
+    dep_obj["type"] = static_cast<int>(dep.edge_data.type);
     deps_array.push_back(dep_obj);
   }
   root["dependencies"] = deps_array;
@@ -2031,25 +2095,34 @@ GraphLoadResult CUDAGraph::load(const std::string& json_path, MempoolId_t pool)
   if (!deps_array.empty()) {
     std::vector<CUgraphNode> from_nodes;
     std::vector<CUgraphNode> to_nodes;
+    std::vector<CUgraphEdgeData> edge_data;
     from_nodes.reserve(deps_array.size());
     to_nodes.reserve(deps_array.size());
+    edge_data.reserve(deps_array.size());
 
     for (const auto& dep_val : deps_array) {
       const json::object& dep_obj = dep_val.as_object();
       int from_id = dep_obj.at("from").to_number<int>();
       int to_id = dep_obj.at("to").to_number<int>();
+      CUgraphEdgeData data{};
+      if (dep_obj.contains("from_port")) {
+        data.from_port = static_cast<decltype(data.from_port)>(
+            dep_obj.at("from_port").to_number<unsigned int>());
+      }
+      if (dep_obj.contains("to_port")) {
+        data.to_port =
+            static_cast<decltype(data.to_port)>(dep_obj.at("to_port").to_number<unsigned int>());
+      }
+      if (dep_obj.contains("type")) {
+        data.type = static_cast<decltype(data.type)>(dep_obj.at("type").to_number<unsigned int>());
+      }
 
       from_nodes.push_back(id_to_node[from_id]);
       to_nodes.push_back(id_to_node[to_id]);
+      edge_data.push_back(data);
     }
 
-#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000)
-    CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(),
-                                                 nullptr, deps_array.size());
-#else
-    CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(),
-                                                    nullptr, deps_array.size());
-#endif
+    CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data);
     if (dep_result != CUDA_SUCCESS) {
       fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n",
               dep_result);

csrc/CUDAGraphParallel.cpp

@@ -776,22 +776,31 @@ GraphLoadResult CUDAGraph::build_graph_from_parsed(ParsedGraphData&& parsed, CUc
   if (!deps_array.empty()) {
     std::vector<CUgraphNode> from_nodes;
     std::vector<CUgraphNode> to_nodes;
+    std::vector<CUgraphEdgeData> edge_data;
     from_nodes.reserve(deps_array.size());
     to_nodes.reserve(deps_array.size());
+    edge_data.reserve(deps_array.size());
 
     for (const auto& dep_val : deps_array) {
       const json::object& dep_obj = dep_val.as_object();
       from_nodes.push_back(id_to_node[dep_obj.at("from").to_number<int>()]);
       to_nodes.push_back(id_to_node[dep_obj.at("to").to_number<int>()]);
+      CUgraphEdgeData data{};
+      if (dep_obj.contains("from_port")) {
+        data.from_port = static_cast<decltype(data.from_port)>(
+            dep_obj.at("from_port").to_number<unsigned int>());
+      }
+      if (dep_obj.contains("to_port")) {
+        data.to_port =
+            static_cast<decltype(data.to_port)>(dep_obj.at("to_port").to_number<unsigned int>());
+      }
+      if (dep_obj.contains("type")) {
+        data.type = static_cast<decltype(data.type)>(dep_obj.at("type").to_number<unsigned int>());
+      }
+      edge_data.push_back(data);
     }
 
-#if (defined(CUDA_VERSION) && CUDA_VERSION >= 13000)
-    CUresult dep_result = cuGraphAddDependencies(cuGraph, from_nodes.data(), to_nodes.data(),
-                                                 nullptr, deps_array.size());
-#else
-    CUresult dep_result = cuGraphAddDependencies_v2(cuGraph, from_nodes.data(), to_nodes.data(),
-                                                    nullptr, deps_array.size());
-#endif
+    CUresult dep_result = add_graph_dependencies(cuGraph, from_nodes, to_nodes, edge_data);
     if (dep_result != CUDA_SUCCESS) {
       fprintf(stderr, "[foundry LOAD ERROR] cuGraphAddDependencies FAILED with error %d\n",
               dep_result);