fix(pcie): retire capture channel aliases

sparkinfer/comm/pcie/pcie_dcp_a2a.cu

@@ -78,6 +78,10 @@ struct RankStaging {
   void *ptrs[kMaxRanks];
 };
 
+struct DoubleStaging {
+  RankStaging slots[2];
+};
+
 template <typename T> struct __align__(16) Pack {
   T values[8];
 };
@@ -97,6 +101,21 @@ static DINLINE FlagType load_flag(FlagType *address) {
   return value;
 }
 
+template <int world_size>
+DINLINE void select_staging(RankStaging &staging,
+                            const DoubleStaging &staging_options,
+                            Signal *self) {
+  if (threadIdx.x == 0) {
+    const int slot =
+        int(load_flag(&self->self_counter[blockIdx.x][0]) & 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>
@@ -149,7 +168,7 @@ 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,
@@ -169,6 +188,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);
@@ -286,7 +308,7 @@ __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) {
   constexpr int kPackElems = 8;
@@ -299,6 +321,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;
@@ -339,11 +364,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 +379,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,9 +419,8 @@ 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;
+    // Both staging slabs are kernel arguments. Each block selects from its
+    // pre-increment execution counter, so capture/replay cannot bake a slot.
     const int heads_per_rank = total_heads / world_size_;
     const int rows = batch * heads_per_rank;
     const int warps_per_block = threads / 32;
@@ -405,8 +429,8 @@ public:
 
 #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)
     switch (world_size_) {
@@ -455,18 +479,17 @@ 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 selected from device execution state in the kernel,
+    // including for back-to-back eager launches.
     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));
 
 #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)
     switch (world_size_) {
     case 2:
       LAUNCH(2);

sparkinfer/comm/pcie/pcie_dcp_a2a.py

@@ -37,15 +37,15 @@ def _is_supported_bhd_layout(tensor: torch.Tensor) -> bool:
         return False
     batch, heads, head_dim = (int(value) for value in tensor.shape)
     stride_batch, stride_head, _ = (int(value) for value in tensor.stride())
-    packed_token_major = (
-        stride_batch == heads * head_dim and stride_head == head_dim
-    )
+    packed_token_major = stride_batch == heads * head_dim and stride_head == head_dim
     capacity_strided_head_major = (
         stride_batch == head_dim
         and stride_head >= batch * head_dim
         and stride_head % 8 == 0
     )
     return packed_token_major or capacity_strided_head_major
+
+
 @dataclass(frozen=True)
 class _StagingLayout:
     signal_bytes: int
@@ -386,9 +386,7 @@ def _validate(
                 f"output shape must be {expected_out}, got {tuple(out.shape)}"
             )
         if not _is_supported_bhd_layout(partial_output):
-            raise ValueError(
-                "partial_output must be packed token-major or head-major"
-            )
+            raise ValueError("partial_output must be packed token-major or head-major")
         if not partial_lse.is_contiguous():
             raise ValueError("partial_lse must be contiguous")
         if not _is_supported_bhd_layout(out):
@@ -766,6 +764,7 @@ def all_gather_heads(
     @contextmanager
     def capture(self, stream: object = None):
         """Bind nested CUDA captures to the enclosing stream's channel."""
+        previous_channels: Optional[dict[int, PCIeDCPA2A]] = None
         if self.single_channel:
             channel = self.for_stream(stream)
         else:
@@ -774,6 +773,7 @@ def capture(self, stream: object = None):
                     "PCIe DCP A2A capture context must be entered before CUDA "
                     "graph capture starts"
                 )
+            previous_channels = dict(self._channels)
             stream_key = _current_stream_key(self.device, stream)
             key = 0 if stream_key is None else int(stream_key)
             # Keep a graph-owned channel even when CUDA recycles the enclosing
@@ -787,6 +787,19 @@ def capture(self, stream: object = None):
             popped = self._capture_channel_stack.pop()
             if popped is not channel:
                 raise RuntimeError("PCIe DCP A2A capture channel stack corrupted")
+            if previous_channels is not None:
+                # Captured graph nodes retain the channel through
+                # ``_all_channels``. Restore eager mappings and discard every
+                # nested capture-stream alias so a recycled stream key cannot
+                # hand this graph-owned channel to a later unwrapped capture.
+                for key, mapped in tuple(self._channels.items()):
+                    if mapped is not channel:
+                        continue
+                    previous = previous_channels.get(key)
+                    if previous is None:
+                        del self._channels[key]
+                    else:
+                        self._channels[key] = previous
 
     def close(self) -> None:
         if self._closed: