feat(gms): add vLLM and SGLang host-tier adapters

Signed-off-by: mkhadkevich <mkhadkevich@nvidia.com>

Author: hutm

lib/gms_kv_ring/engines/sglang/__init__.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""SGLang adapter for the gms_kv_ring handle.
+
+`install_for_sglang(...)` monkey-patches SGLang's RadixCache so that:
+
+  • on eviction, evicted blocks are pushed to the evict ring
+  • on cache hit, restored block_id pairs are pushed to the restore
+    ring and the compute stream waits on the per-slot counter
+
+The monkey-patch is intentionally small. Engines pass us their
+existing block-table machinery; we only intercept eviction and
+restore call-sites.
+
+NOTE: this file imports SGLang lazily (only inside `install_for_sglang`)
+so unit tests for the handle don't require SGLang to be importable.
+"""
+
+from gms_kv_ring.engines.sglang.install_kv_ring import install_for_sglang
+
+__all__ = ["install_for_sglang"]

lib/gms_kv_ring/engines/sglang/gds_connector.py

@@ -0,0 +1,75 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""GDS-direct connector glue for SGLang's RadixCache / HiCache.
+
+The connector logic is engine-agnostic — see
+`gms_kv_ring.engines.gds_block_connector.BlockGdsConnector`. This
+module re-exports it under an SGLang-flavored name and documents
+the SGLang-specific integration recipe.
+
+SGLang specifics:
+  - RadixCache nodes carry token-index slots into the KV pool.
+    What vLLM calls "block_id," SGLang calls a per-token slot
+    index. The "block" we offload to storage is typically a fixed
+    page (e.g., 64 tokens) that aligns with SGLang's `page_size`.
+  - For block-granular GDS spill/restore, the caller assembles
+    page-aligned slot groups and treats each as a "block id."
+
+Integration recipe:
+
+    from gms_kv_ring.engines.sglang.install_kv_ring import install_for_sglang
+    from gms_kv_ring.engines.sglang.gds_connector import SGLangGdsConnector
+
+    handle = install_for_sglang(
+        engine_id=..., daemon_socket=..., layers=...,
+    )
+
+    def block_layout(page_id):
+        # SGLang specifics: layers laid out separately, each layer
+        # has shape (num_pages, page_size_bytes). One "block" here
+        # = one page across all attention layers.
+        return [(L, page_id * page_size_bytes, page_size_bytes)
+                for L in range(n_layers)]
+
+    gds = SGLangGdsConnector(handle, block_layout)
+
+    # In RadixCache.evict(num_tokens):
+    #   - Pick victim pages (existing LRU/freq logic in SGLang).
+    #   - If gds.is_available(), spill them to durable storage
+    #     so they survive engine restart and can be promoted on
+    #     re-access. Otherwise use the standard host_tier path.
+    def evict(self, num_tokens):
+        victim_pages = self._select_victims(num_tokens)
+        if gds.is_available():
+            gds.evict_blocks_to_storage(victim_pages)
+        else:
+            # standard host_tier path via handle.record_evict per
+            # victim, same as the existing on_evict hook.
+            ...
+
+    # On match_prefix cache hit, before serving from HBM:
+    #   - Identify the hit pages.
+    #   - If they're currently in storage (because they were
+    #     previously evicted), promote them back to HBM in-place
+    #     so the attention path can read from the pool as usual.
+    def restore_prefix(self, hit_pages):
+        if gds.is_available():
+            ok = gds.restore_blocks_from_storage(hit_pages)
+            # ok[page_id]=False → engine must fall to cold compute
+            # for that page (the slot is "missing" from the cache).
+            return ok
+        else:
+            # standard host_tier promote + restore_ring path
+            ...
+"""
+
+from __future__ import annotations
+
+from gms_kv_ring.engines.gds_block_connector import BlockGdsConnector, EvictResult
+
+# Engine-flavored alias. Same class; SGLang-specific recipe lives
+# in this module's docstring above.
+SGLangGdsConnector = BlockGdsConnector
+
+__all__ = ["SGLangGdsConnector", "EvictResult"]

lib/gms_kv_ring/engines/sglang/install_kv_ring.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Wire a GMSKvRing handle into SGLang's RadixCache.
+
+Public entry point: `install_for_sglang(engine_id, daemon_socket,
+layer_descs, *, on_evict=None, on_hit=None)`.
+
+The function returns a `GMSKvRing` handle. The caller (engine
+bootstrap) is responsible for closing it on engine shutdown.
+
+The on_evict / on_hit hooks the caller can pass in are optional —
+if omitted, we install a default monkey-patch on
+`sglang.srt.mem_cache.radix_cache.RadixCache.evict` and on the
+`match_prefix` path that publishes the obvious record sets.
+
+Why callback-injection instead of always monkey-patching: SGLang's
+internals shift between releases (we've seen #199 hangs from
+cross-stream sync). Giving the engine the option to call
+`handle.record_evict(...)` from its own well-tested hook is safer
+than auto-patching across SGLang versions.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Callable, Optional
+
+from gms_kv_ring.engines.handle import GMSKvRing
+
+logger = logging.getLogger(__name__)
+
+
+def install_for_sglang(
+    *,
+    engine_id: str,
+    daemon_socket: str,
+    layers: list[dict],
+    on_evict: Optional[Callable[[GMSKvRing], None]] = None,
+    on_hit: Optional[Callable[[GMSKvRing], None]] = None,
+    evict_ring_capacity: int = 4096,
+    restore_ring_capacity: int = 4096,
+    num_counters: int = 512,
+) -> GMSKvRing:
+    """Build the handle and (optionally) wire SGLang hooks to it.
+
+    Parameters
+    ----------
+    engine_id : str
+        Stable per-engine identifier. Use the same value across
+        restarts so the daemon can reuse host-tier slots.
+    daemon_socket : str
+        Unix-socket path the daemon is listening on.
+    layers : list[dict]
+        Per-layer descriptors. Each entry: {layer_idx, va, size, stride}.
+        `va` is the engine-process virtual address of the layer base
+        (i.e. the engine has already CUDA-mapped the pool). The
+        daemon's pool attachment will rely on these.
+    on_evict, on_hit : callable | None
+        Optional callbacks invoked exactly once with the handle, so
+        the caller can wire its own engine hooks. If None, this is a
+        plain-handle install — useful for tests and for callers that
+        want to wire their hooks explicitly.
+    """
+    handle = GMSKvRing(
+        engine_id=engine_id,
+        daemon_socket=daemon_socket,
+        layers=layers,
+        evict_ring_capacity=evict_ring_capacity,
+        restore_ring_capacity=restore_ring_capacity,
+        num_counters=num_counters,
+    )
+    if on_evict is not None:
+        on_evict(handle)
+    if on_hit is not None:
+        on_hit(handle)
+    logger.info("gms_kv_ring installed for SGLang engine %r", engine_id)
+    return handle

lib/gms_kv_ring/engines/vllm/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""vLLM adapter for the gms_kv_ring handle.
+
+Same shape as the SGLang adapter: callers pass layer descriptors,
+get back a `GMSKvRing` handle, and wire `record_evict` / `record_restore`
+into the engine's eviction (KVCacheConnector.evict_kv_blocks) and
+cache-hit (KVCacheConnector.start_load_kv) call-sites.
+"""
+
+from gms_kv_ring.engines.vllm.install_kv_ring import install_for_vllm
+
+__all__ = ["install_for_vllm"]

lib/gms_kv_ring/engines/vllm/gds_connector.py

@@ -0,0 +1,58 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""GDS-direct connector glue for vLLM's KVCacheConnectorV1.
+
+The connector logic is engine-agnostic — see
+`gms_kv_ring.engines.gds_block_connector.BlockGdsConnector`. This
+module re-exports it under a vLLM-flavored name and documents the
+specific integration recipe for vLLM's KVCacheConnectorV1.
+
+Integration recipe (inside a vLLM KVCacheConnectorV1 implementation):
+
+    from gms_kv_ring.engines.vllm.install_kv_ring import install_for_vllm
+    from gms_kv_ring.engines.vllm.gds_connector import VllmGdsConnector
+
+    handle = install_for_vllm(
+        engine_id=..., daemon_socket=..., layers=...,
+    )
+
+    def block_layout(block_id):
+        # vLLM-version-specific. For typical paged-attention with
+        # contiguous per-layer regions:
+        #   offset = block_id * block_size_bytes
+        #   size   = block_size_bytes
+        return [(L, block_id * block_size_bytes, block_size_bytes)
+                for L in range(n_layers)]
+
+    gds = VllmGdsConnector(handle, block_layout)
+
+    # vLLM's evict_kv_blocks hook:
+    def evict_kv_blocks(self, block_ids, ipc_event_handle=b""):
+        if gds.is_available():
+            gds.evict_blocks_to_storage(block_ids)
+        else:
+            # standard host_tier path
+            for b in block_ids:
+                handle.record_evict(b, block_layout(b), ipc_event_handle)
+
+    # vLLM's start_load_kv hook (on cache hit):
+    def start_load_kv(self, request):
+        block_ids = request.hit_blocks  # vLLM-version-specific
+        if gds.is_available():
+            ok = gds.restore_blocks_from_storage(block_ids)
+            # Caller decides cold-fallback for ok[b]==False entries.
+        else:
+            # standard host_tier promote + restore_ring path
+            ...
+"""
+
+from __future__ import annotations
+
+from gms_kv_ring.engines.gds_block_connector import BlockGdsConnector, EvictResult
+
+# Engine-flavored alias. Same class; vLLM-specific recipe lives in
+# this module's docstring above.
+VllmGdsConnector = BlockGdsConnector
+
+__all__ = ["VllmGdsConnector", "EvictResult"]

lib/gms_kv_ring/engines/vllm/install_kv_ring.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Wire a GMSKvRing handle into vLLM.
+
+vLLM's KVCacheConnectorV1 interface exposes the two hooks we need:
+
+  • `evict_kv_blocks(block_ids, ipc_event_handle)` — call
+    `handle.record_evict(...)` here.
+  • `start_load_kv(request)` — translate the request's hit prefix
+    into (src_block, dest_block) pairs and call
+    `handle.record_restore(...)` then enqueue
+    `handle.wait_restore(stream, slot, target)` on the compute stream.
+
+This module deliberately does NOT import vLLM at module load —
+unit tests run without vLLM installed. The caller imports vLLM in
+its connector and uses our returned handle from inside its hook
+methods.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Callable, Optional
+
+from gms_kv_ring.engines.handle import GMSKvRing
+
+logger = logging.getLogger(__name__)
+
+
+def install_for_vllm(
+    *,
+    engine_id: str,
+    daemon_socket: str,
+    layers: list[dict],
+    on_evict: Optional[Callable[[GMSKvRing], None]] = None,
+    on_hit: Optional[Callable[[GMSKvRing], None]] = None,
+    evict_ring_capacity: int = 4096,
+    restore_ring_capacity: int = 4096,
+    num_counters: int = 512,
+) -> GMSKvRing:
+    """Build the handle and (optionally) invoke wiring callbacks.
+
+    See `engines/sglang/install_kv_ring.py` for parameter semantics —
+    they are identical."""
+    handle = GMSKvRing(
+        engine_id=engine_id,
+        daemon_socket=daemon_socket,
+        layers=layers,
+        evict_ring_capacity=evict_ring_capacity,
+        restore_ring_capacity=restore_ring_capacity,
+        num_counters=num_counters,
+    )
+    if on_evict is not None:
+        on_evict(handle)
+    if on_hit is not None:
+        on_hit(handle)
+    logger.info("gms_kv_ring installed for vLLM engine %r", engine_id)
+    return handle