Linhu/kvcr 45650 land

python/sglang/srt/mem_cache/storage/kvcr/README-linker.md

@@ -4,7 +4,7 @@ The KVCR linker transfers registered GPU KV pages directly to and from KVCR's
 local DRAM or a hinted remote KVCR peer. It uses `deposit`, `query`, and `deliver`,
 with at most four layer deliveries outstanding. Its owner thread continuously
 polls KVCR, including while idle, so peer requests and transfer lifecycle events
-keep progressing. Framework source-pin requests are declined; peers can retrieve
+keep progressing. Framework source pinning is disabled; peers can retrieve
 completed deposits in KVCR-owned DRAM.
 
 For SWA/state layouts, ranks confirm a shared prefix boundary by fetching its
@@ -16,8 +16,8 @@ direct delivery, and layouts without trailing state skip this confirmation.
 ## Setup
 
 Install KVCR with its NIXL/UCX dependencies in the serving environment. Use the
-public API from KVCR revision `8a91e20` or a compatible revision that includes
-the local-copy timeout lifecycle fix.
+public API from KVCR revision `0fd754d8` (PR #133) or a compatible revision with
+optional framework-pinning callbacks and the local-copy timeout lifecycle fix.
 
 DRAM additions and removals carry `ownership="kvcr"` in the existing KV event
 stream. Engine GPU events and cache clears retain the default ownership. Dynamo

python/sglang/srt/mem_cache/storage/kvcr/kvcr_direct_linker.py

@@ -167,33 +167,44 @@ def prepare(self, indices_by_pool, owners_by_pool):
         return prepared
 
     def gather_layer(self, layer_id, prepared):
+        # One all-gather per layer: pack every buffer's owned rows back to back.
+        parts = []
         for name, (padded, counts, width) in prepared.items():
             layer = self.pools[name].layer_mapping.get(layer_id)
             if layer is None:
                 continue
             for buffers in self.buffers[name]:
                 for index in (layer,) if isinstance(layer, int) else layer:
-                    self._gather(buffers[index], padded, counts, width)
+                    parts.append((buffers[index], padded, counts, width))
+        if parts:
+            self._gather(parts)
 
-    def _gather(self, buf, padded, counts, width):
-        row_bytes = buf.shape[1]
-        need = width * row_bytes
+    def _gather(self, parts):
+        sizes = [width * buf.shape[1] for buf, _, _, width in parts]
+        need = sum(sizes)
         if self.staging.numel() < need:
             self.staging = torch.empty(need, dtype=torch.uint8, device=self.device)
         if self.gathered.numel() < need * self.size:
             self.gathered = torch.empty(
                 need * self.size, dtype=torch.uint8, device=self.device
             )
-        mine = self.staging[:need].view(width, row_bytes)
-        torch.index_select(buf, 0, padded[self.rank], out=mine)
+        offset = 0
+        for (buf, padded, _, width), size in zip(parts, sizes):
+            mine = self.staging[offset : offset + size].view(width, buf.shape[1])
+            torch.index_select(buf, 0, padded[self.rank], out=mine)
+            offset += size
         out = self.gathered[: need * self.size]
         torch.distributed.all_gather_into_tensor(
             out, self.staging[:need], group=self.group
         )
-        out = out.view(self.size, width, row_bytes)
-        for r in range(self.size):
-            if r != self.rank and counts[r]:
-                buf.index_copy_(0, padded[r, : counts[r]], out[r, : counts[r]])
+        out = out.view(self.size, need)
+        offset = 0
+        for (buf, padded, counts, width), size in zip(parts, sizes):
+            rows = out[:, offset : offset + size].view(self.size, width, buf.shape[1])
+            for r in range(self.size):
+                if r != self.rank and counts[r]:
+                    buf.index_copy_(0, padded[r, : counts[r]], rows[r, : counts[r]])
+            offset += size
 
 
 def _offset_source_ports(hints, offset):
@@ -420,10 +431,15 @@ def _run(self):
                 return
             self._runtime = error.kvcr_runtime
         while not self._closed:
-            active = self._fatal is None and (self._batches or self._offloads)
+            # Spin only while transfers are in flight; a batch that only awaits its
+            # forward (gather/drain) would otherwise contend for the GIL with it.
+            active = self._fatal is None and (
+                self._offloads
+                or any(not b.exhausted or b.active for b in self._batches)
+            )
             try:
                 function, args, future = self._commands.get(
-                    timeout=0 if active else 0.001
+                    timeout=0 if active else (0.0005 if self._batches else 0.001)
                 )
             except Empty:
                 if active:

python/sglang/srt/mem_cache/storage/kvcr/runtime.py

@@ -150,8 +150,6 @@ def __init__(self, layout, options: dict, *, rank: int):
         self.events = deque()
         self.removals = {}
         self._buffers = []
-        self._pin_results = {}
-        self._next_pin = 0
         self.client = None
         self.process_local = False
 
@@ -180,10 +178,6 @@ def inventory(event):
             )
             self.control_endpoint = control.endpoint
             bindings = KVCRBindings(
-                request_pin=self._decline_pin,
-                poll_pin_results=self._poll_pin_results,
-                release_pin=lambda _: True,
-                cancel_pin_request=lambda request: self._pin_results.pop(request, None),
                 framework_control=control,
                 key_adapter=self,
                 inventory_sink=inventory,
@@ -249,27 +243,24 @@ def inventory(event):
                 self._close_buffers()
             raise
 
-    def _decline_pin(self, keys):
-        self._next_pin += 1
-        self._pin_results[self._next_pin] = None
-        return self._next_pin
-
-    def _poll_pin_results(self):
-        results = list(self._pin_results.items())
-        self._pin_results.clear()
-        return results
+    def encode_to_framework(self, key: bytes) -> bytes:
+        return key
 
-    def encode(self, key: object) -> bytes:
-        if not isinstance(key, bytes):
+    def decode_from_framework(self, framework_key: object) -> bytes:
+        if not isinstance(framework_key, bytes):
             raise TypeError("KVCR keys must be encoded layout keys")
-        return key
+        return framework_key
 
-    def decode(self, key: bytes) -> int:
+    def encode_to_router(self, key: bytes) -> int:
         decoded = self._layout.decode_key(key)
         if decoded is None:
             raise ValueError("Unknown KVCR layout key")
         return int(decoded[1][:16], 16)
 
+    def decode_from_router(self, router_key: int | bytes) -> bytes:
+        # A router hash names a page across all pools, so it has no single key.
+        raise NotImplementedError("no single KVCR key per router hash")
+
     def log_stats(self, *, force=False):
         if not self._telemetry_enabled:
             return

test/registered/unit/mem_cache/test_kvcr_direct_linker.py

@@ -1323,14 +1323,9 @@ def test_refused_or_uncertain_guard_claim_never_falls_back(self):
                     with self.assertRaisesRegex(RuntimeError, "startup"):
                         runtime.close()
 
-    def test_declines_pins_and_cancels_removals_per_physical_key(self):
+    def test_cancels_removals_per_physical_key(self):
         runtime = self.make_runtime()
         bindings = runtime.client.bindings
-        request = bindings.request_pin([b"small:1234"])
-        cancelled = bindings.request_pin([b"large:1234"])
-        bindings.cancel_pin_request(cancelled)
-        self.assertEqual(list(bindings.poll_pin_results()), [(request, None)])
-        self.assertEqual(list(bindings.poll_pin_results()), [])
         for keys, removed in (
             ((b"small:1234", b"large:1234"), True),
             ((b"small:1234",), False),