diff --git a/python/sglang/srt/mem_cache/storage/kvcr/README-linker.md b/python/sglang/srt/mem_cache/storage/kvcr/README-linker.md index 3d2ae0dafec1..e91872d94842 100644 --- a/python/sglang/srt/mem_cache/storage/kvcr/README-linker.md +++ b/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 diff --git a/python/sglang/srt/mem_cache/storage/kvcr/kvcr_direct_linker.py b/python/sglang/srt/mem_cache/storage/kvcr/kvcr_direct_linker.py index ce0f09c0c8ca..146211b30361 100644 --- a/python/sglang/srt/mem_cache/storage/kvcr/kvcr_direct_linker.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/storage/kvcr/runtime.py b/python/sglang/srt/mem_cache/storage/kvcr/runtime.py index 80b457d9795a..504b91c52804 100644 --- a/python/sglang/srt/mem_cache/storage/kvcr/runtime.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_kvcr_direct_linker.py b/test/registered/unit/mem_cache/test_kvcr_direct_linker.py index df551cbd5aa6..6df990354fa2 100644 --- a/test/registered/unit/mem_cache/test_kvcr_direct_linker.py +++ b/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),