Debug [skip ci]

Author: UNIDY2002

mooncake-ep/src/mooncake_backend.cpp

@@ -1077,11 +1077,26 @@ void MooncakeBackend::advanceRecvOp(RecvOpStateData& opData) {
 
         case RecvOpState::READY_TO_ALLOC: {
             // Check if the slot is already allocated by another active operation
-            // If so, find the next available slot
-            while (allocatedRecvSlots_.find(opData.baseSlot) != allocatedRecvSlots_.end()) {
+            // If so, find the next available slot that doesn't cause wrap-around
+            int attempts = 0;
+            while (attempts < static_cast<int>(kP2PNumSlots)) {
+                // Check if current slot is free and doesn't cause wrap-around
+                if (allocatedRecvSlots_.find(opData.baseSlot) == allocatedRecvSlots_.end() &&
+                    opData.baseSlot + opData.numSlotsNeeded <= static_cast<int>(kP2PNumSlots)) {
+                    // Found a valid slot
+                    break;
+                }
+                // Try next slot, wrapping around if needed
                 opData.baseSlot = (opData.baseSlot + 1) % static_cast<int>(kP2PNumSlots);
+                attempts++;
             }
 
+            TORCH_CHECK(attempts < static_cast<int>(kP2PNumSlots),
+                        "P2P recv: failed to find available slot after ", attempts, " attempts");
+            TORCH_CHECK(opData.baseSlot + opData.numSlotsNeeded <= static_cast<int>(kP2PNumSlots),
+                        "P2P recv: slot range would exceed buffer: baseSlot=", opData.baseSlot,
+                        ", numSlots=", opData.numSlotsNeeded);
+            
             // Mark slot as allocated
             allocatedRecvSlots_.insert(opData.baseSlot);