fix(pcie): harden replay and CUDA IPC lifecycle

pyproject.toml

@@ -39,7 +39,7 @@ where = ["."]
 include = ["sparkinfer*"]
 
 [tool.setuptools.package-data]
-"sparkinfer.comm.pcie" = ["*.cu"]
+"sparkinfer.comm.pcie" = ["*.cu", "*.h"]
 
 [tool.ruff.lint]
 select = ["E", "F", "B", "SIM"]

sparkinfer/comm/pcie/ipc_handle_registry.h

@@ -0,0 +1,63 @@
+#pragma once
+
+#include <map>
+#include <utility>
+
+namespace sparkinfer::pcie {
+
+// Owns imported IPC handles whose close routine reports an error code instead
+// of throwing. Successful closes clear their entries so repeated cleanup is
+// idempotent; failed entries remain available to an explicit retry.
+template <typename Key, typename Handle, typename Error, Error Success>
+class IpcHandleRegistry {
+ public:
+  using CloseFn = Error (*)(Handle) noexcept;
+  using Map = std::map<Key, Handle>;
+  using iterator = typename Map::iterator;
+
+  explicit IpcHandleRegistry(CloseFn close) noexcept : close_(close) {}
+
+  IpcHandleRegistry(const IpcHandleRegistry&) = delete;
+  IpcHandleRegistry& operator=(const IpcHandleRegistry&) = delete;
+  IpcHandleRegistry(IpcHandleRegistry&&) = delete;
+  IpcHandleRegistry& operator=(IpcHandleRegistry&&) = delete;
+
+  ~IpcHandleRegistry() noexcept {
+    (void)close_all_noexcept();
+  }
+
+  iterator find(const Key& key) {
+    return handles_.find(key);
+  }
+
+  iterator end() {
+    return handles_.end();
+  }
+
+  std::pair<iterator, bool> emplace(const Key& key, Handle handle) {
+    return handles_.emplace(key, handle);
+  }
+
+  Error close_all_noexcept() noexcept {
+    Error first_error = Success;
+    for (auto& entry : handles_) {
+      Handle& handle = entry.second;
+      if (handle == Handle{}) {
+        continue;
+      }
+      const Error error = close_(handle);
+      if (error == Success) {
+        handle = Handle{};
+      } else if (first_error == Success) {
+        first_error = error;
+      }
+    }
+    return first_error;
+  }
+
+ private:
+  Map handles_;
+  CloseFn close_;
+};
+
+}  // namespace sparkinfer::pcie

sparkinfer/comm/pcie/pcie_dcp_a2a.cu

@@ -20,6 +20,7 @@
 #include <algorithm>
 #include <array>
 #include <cmath>
+#include <cstdint>
 #include <cstdlib>
 #include <sstream>
 #include <stdexcept>
@@ -66,6 +67,8 @@ static int dcp_block_limit_override() {
 }
 
 struct Signal {
+  alignas(128) FlagType staging_generation;
+  FlagType active_staging_slot;
   alignas(128) FlagType self_counter[kMaxBlocks][kMaxRanks];
   alignas(128) FlagType peer_counter[2][kMaxBlocks][kMaxRanks * kFlagStride];
 };
@@ -78,6 +81,10 @@ struct RankStaging {
   void *ptrs[kMaxRanks];
 };
 
+struct DoubleStaging {
+  RankStaging slots[2];
+};
+
 template <typename T> struct __align__(16) Pack {
   T values[8];
 };
@@ -97,6 +104,40 @@ static DINLINE FlagType load_flag(FlagType *address) {
   return value;
 }
 
+static DINLINE FlagType load_flag_gpu(FlagType *address) {
+  FlagType value;
+  asm volatile("ld.relaxed.gpu.global.u32 %0, [%1];"
+               : "=r"(value)
+               : "l"(address)
+               : "memory");
+  return value;
+}
+
+__global__ void advance_staging_slot_kernel(Signal *self) {
+  if (threadIdx.x == 0) {
+    const FlagType generation = self->staging_generation;
+    self->active_staging_slot = generation & FlagType{1};
+    self->staging_generation = generation + FlagType{1};
+  }
+}
+
+template <int world_size>
+DINLINE void select_staging(RankStaging &staging,
+                            const DoubleStaging &staging_options,
+                            Signal *self) {
+  if (threadIdx.x == 0) {
+    // The one-CTA control node runs earlier on this stream.  Every worker CTA
+    // therefore observes one operation-wide slot regardless of worker grid
+    // size or rank launch skew.
+    const int slot = int(load_flag_gpu(&self->active_staging_slot) & FlagType{1});
+#pragma unroll
+    for (int peer = 0; peer < world_size; ++peer) {
+      staging.ptrs[peer] = staging_options.slots[slot].ptrs[peer];
+    }
+  }
+  __syncthreads();
+}
+
 // pre_sync orders in-kernel staging stores (from every warp of the block)
 // before the flag post; without staging the flags can go out immediately.
 template <int world_size, bool pre_sync>
@@ -117,6 +158,16 @@ DINLINE void start_barrier(const RankSignals &signals, Signal *self, int rank) {
   __syncthreads();
 }
 
+DINLINE void test_post_barrier_delay(int rank, int delayed_rank,
+                                     uint64_t delay_cycles) {
+  if (delay_cycles != 0 && rank == delayed_rank && threadIdx.x == 0) {
+    const uint64_t start = clock64();
+    while (clock64() - start < delay_cycles) {
+    }
+  }
+  __syncthreads();
+}
+
 DINLINE float to_float(half value) { return __half2float(value); }
 DINLINE float to_float(nv_bfloat16 value) { return __bfloat162float(value); }
 
@@ -149,14 +200,15 @@ template <typename T, int world_size>
 __global__ void __launch_bounds__(512, 1)
     dcp_lse_reduce_kernel(const T *__restrict__ local_output,
                           const float *__restrict__ local_lse,
-                          RankStaging staging, int64_t lse_offset,
+                          DoubleStaging staging_options, int64_t lse_offset,
                           RankSignals signals, Signal *self,
                           T *__restrict__ output, int rank, int batch,
                           int total_heads, int head_dim,
                           int64_t input_stride_batch,
                           int64_t input_stride_head,
                           int64_t output_stride_batch,
-                          int64_t output_stride_head, bool natural_log) {
+                          int64_t output_stride_head, bool natural_log,
+                          int delayed_rank, uint64_t delay_cycles) {
   constexpr int kPackElems = 8;
   const int heads_per_rank = total_heads / world_size;
   const int packs_per_head = head_dim / kPackElems;
@@ -169,6 +221,9 @@ __global__ void __launch_bounds__(512, 1)
   const auto *local_packs = reinterpret_cast<const Pack<T> *>(local_output);
   auto *output_packs = reinterpret_cast<Pack<T> *>(output);
 
+  __shared__ RankStaging staging;
+  select_staging<world_size>(staging, staging_options, self);
+
   auto *staging_out = reinterpret_cast<Pack<T> *>(staging.ptrs[rank]);
   auto *staging_lse = reinterpret_cast<float *>(
       reinterpret_cast<char *>(staging.ptrs[rank]) + lse_offset);
@@ -192,6 +247,7 @@ __global__ void __launch_bounds__(512, 1)
     }
   }
   start_barrier<world_size, true>(signals, self, rank);
+  test_post_barrier_delay(rank, delayed_rank, delay_cycles);
 
   // Rotated source pointers so every later access uses a compile-time
   // index; the self source reads the local tensors directly (still hot in
@@ -286,9 +342,10 @@ __global__ void __launch_bounds__(512, 1)
 template <typename T, int world_size>
 __global__ void __launch_bounds__(512, 1)
     all_gather_heads_kernel(const T *__restrict__ local_input,
-                            RankStaging staging, RankSignals signals,
+                            DoubleStaging staging_options, RankSignals signals,
                             Signal *self, T *__restrict__ output, int rank,
-                            int batch, int local_heads, int head_dim) {
+                            int batch, int local_heads, int head_dim,
+                            int delayed_rank, uint64_t delay_cycles) {
   constexpr int kPackElems = 8;
   const int packs_per_head = head_dim / kPackElems;
   const int total_heads = local_heads * world_size;
@@ -299,6 +356,9 @@ __global__ void __launch_bounds__(512, 1)
   const int warp_stride = gridDim.x * warps_per_block;
   const auto *local_packs = reinterpret_cast<const Pack<T> *>(local_input);
 
+  __shared__ RankStaging staging;
+  select_staging<world_size>(staging, staging_options, self);
+
   auto *staging_out = reinterpret_cast<Pack<T> *>(staging.ptrs[rank]);
   for (int row = warp_first; row < rows; row += warp_stride) {
     const int batch_index = row / total_heads;
@@ -313,6 +373,7 @@ __global__ void __launch_bounds__(512, 1)
     }
   }
   start_barrier<world_size, true>(signals, self, rank);
+  test_post_barrier_delay(rank, delayed_rank, delay_cycles);
 
   auto *output_packs = reinterpret_cast<Pack<T> *>(output);
   for (int row = warp_first; row < rows; row += warp_stride) {
@@ -339,11 +400,11 @@ public:
   int world_size_;
   RankSignals signals_{};
   Signal *self_signal_;
-  RankStaging staging_[2]{};
+  DoubleStaging staging_{};
   int64_t output_capacity_elems_;
   int64_t lse_offset_;
   int64_t lse_capacity_;
-  int slot_ = 0;
+
 
   PCIeDCPA2A(Signal **signals,
              const std::vector<std::array<void *, 2>> &staging,
@@ -354,8 +415,8 @@ public:
         lse_capacity_(lse_capacity) {
     for (int peer = 0; peer < world_size_; ++peer) {
       signals_.signals[peer] = signals[peer];
-      staging_[0].ptrs[peer] = staging[peer][0];
-      staging_[1].ptrs[peer] = staging[peer][1];
+      staging_.slots[0].ptrs[peer] = staging[peer][0];
+      staging_.slots[1].ptrs[peer] = staging[peer][1];
     }
   }
 
@@ -394,21 +455,27 @@ public:
       throw std::runtime_error("threads must be a multiple of 32");
     }
 
-    // Staging happens inside the kernel (warp-per-row, before the start
-    // barrier); no host staging memcpys are issued.
-    const int slot = slot_++ % 2;
+    // Select one staging slot per execution in a graph-capturable device node.
+    // Worker CTA count may change large -> small -> large without leaving
+    // dormant per-block parity counters behind.
     const int heads_per_rank = total_heads / world_size_;
     const int rows = batch * heads_per_rank;
     const int warps_per_block = threads / 32;
     const int blocks = std::max(
         1, std::min(block_limit, (rows + warps_per_block - 1) / warps_per_block));
+    advance_staging_slot_kernel<<<1, 1, 0, stream>>>(self_signal_);
+    CHECK_CUDA_SUCCESS(cudaGetLastError());
+    const int delayed_rank =
+        env_int("SPARKINFER_PCIE_DCP_TEST_DELAY_RANK", -1);
+    const uint64_t delay_cycles = static_cast<uint64_t>(std::max(
+        0, env_int("SPARKINFER_PCIE_DCP_TEST_POST_BARRIER_DELAY_CYCLES", 0)));
 
 #define LAUNCH(world)                                                          \
   dcp_lse_reduce_kernel<T, world><<<blocks, threads, 0, stream>>>(             \
-      partial_output, partial_lse, staging_[slot], lse_offset_, signals_,      \
-      self_signal_, output, rank_, batch, total_heads, head_dim,               \
+      partial_output, partial_lse, staging_, lse_offset_, signals_,           \
+      self_signal_, output, rank_, batch, total_heads, head_dim,              \
       input_stride_batch, input_stride_head, output_stride_batch,              \
-      output_stride_head, natural_log)
+      output_stride_head, natural_log, delayed_rank, delay_cycles)
     switch (world_size_) {
     case 2:
       LAUNCH(2);
@@ -455,18 +522,23 @@ public:
       throw std::runtime_error("threads must be a multiple of 32");
     }
 
-    // Staging happens inside the kernel (warp-per-row, before the start
-    // barrier); no host staging memcpy is issued.
-    const int slot = slot_++ % 2;
+    // Slot ownership is published by a device control node, including for
+    // back-to-back eager launches and every CUDA graph replay.
     const int rows = batch * total_heads;
     const int warps_per_block = threads / 32;
     const int blocks = std::max(
         1, std::min(block_limit, (rows + warps_per_block - 1) / warps_per_block));
+    advance_staging_slot_kernel<<<1, 1, 0, stream>>>(self_signal_);
+    CHECK_CUDA_SUCCESS(cudaGetLastError());
+    const int delayed_rank =
+        env_int("SPARKINFER_PCIE_DCP_TEST_DELAY_RANK", -1);
+    const uint64_t delay_cycles = static_cast<uint64_t>(std::max(
+        0, env_int("SPARKINFER_PCIE_DCP_TEST_POST_BARRIER_DELAY_CYCLES", 0)));
 
 #define LAUNCH(world)                                                          \
   all_gather_heads_kernel<T, world><<<blocks, threads, 0, stream>>>(           \
-      local_input, staging_[slot], signals_, self_signal_, output, rank_,      \
-      batch, local_heads, head_dim)
+      local_input, staging_, signals_, self_signal_, output, rank_, batch,    \
+      local_heads, head_dim, delayed_rank, delay_cycles)
     switch (world_size_) {
     case 2:
       LAUNCH(2);