perf(unified-memory): batch lazy-compaction mapping lookup (#34066)

test/registered/unit/mem_cache/test_multi_ended_allocator.py

@@ -79,6 +79,24 @@ class _FakeKVCache:
         self.buf[dst_loc] = self.buf[src_loc].clone()
 
 
+class _RejectScalarIndexTensor:
+    """Tensor proxy that rejects one-row-at-a-time mapping lookups."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+    def __getattr__(self, name):
+        return getattr(self.tensor, name)
+
+    def __getitem__(self, index):
+        if isinstance(index, int):
+            raise AssertionError("physical_to_virtual was read one row at a time")
+        return self.tensor[index]
+
+    def __setitem__(self, index, value):
+        self.tensor[index] = value
+
+
 class TestUnifiedKVPoolViews(unittest.TestCase):
     def test_min_slot_index_and_disjoint_bytes(self):
         full = _make_mha_spec("full", "up", layer_num=4)
@@ -2049,6 +2067,27 @@ class TestLazyCompaction(unittest.TestCase):
                 continue  # freed
             self.assertEqual(int(fa.physical_to_virtual[p].item()), v)
 
+    def test_lazy_flush_gathers_survivor_mappings_as_one_batch(self):
+        """Compaction must not synchronize once per relocated survivor."""
+        _pool, fa, kv = self._make_full(lazy=True)
+        values = fa.alloc(12)
+        self._stamp_kv(kv, fa, values)
+        fa.free(values[1:5].clone())
+
+        physical_to_virtual = fa.physical_to_virtual
+        fa.physical_to_virtual = _RejectScalarIndexTensor(physical_to_virtual)
+        try:
+            self.assertEqual(fa._flush(urgent=True), 4)
+        finally:
+            fa.physical_to_virtual = physical_to_virtual
+
+        for virtual in values.tolist():
+            physical = int(fa.virtual_to_physical[virtual].item())
+            if physical == -1:
+                continue
+            self.assertEqual(int(fa.physical_to_virtual[physical].item()), virtual)
+            self.assertEqual(int(kv.buf[physical].item()), virtual)
+
     def _replay_sequence(self, ops, lazy: bool):
         """Run a given alloc/free op trace under eager OR lazy mode and
         return the final (live virtual set, alloc-time KV stamps)."""