support layerwise

Author: Clebrate

corelib/recsys_kvcache_manager/recsys_kvcache_manager/flex_kvcache_manager.py

@@ -39,6 +39,10 @@
 )
 from .kvcache_metadata import KVCacheMetadata
 from .kvcache_utils import KVIndexMeta, KVLookupResult
+from .flexkv_layerwise import (
+    DEFAULT_LAYERWISE_EVENTFD_SOCKET,
+    FlexKVLayerwiseEventfdSender,
+)
 
 
 @dataclass
@@ -56,6 +60,12 @@ class _FlexKVOnloadHandle:
     task_ids: List[int]
     uids: torch.Tensor
     slot_mappings: List[torch.Tensor]
+    layer_eventfds: Optional[List[int]] = None
+
+    def wait_layer(self, layer_idx: int) -> None:
+        if self.layer_eventfds is None:
+            return
+        os.read(self.layer_eventfds[layer_idx], 8)
 
 
 @dataclass
@@ -79,7 +89,7 @@ def __post_init__(self) -> None:
             [
                 self.num_layer,
                 self.num_block,
-                self._kv_dim,
+                self.kv_dim,
                 self.tokens_per_block,
                 self.num_head,
                 self.head_size,
@@ -113,6 +123,10 @@ def __init__(
         enable_mps: bool = False,
         hostkv_wait_timeout_ms: int = 0,
         host_kvstorage_fail_policy: str = "fail_open",
+        config_path: Optional[str] = None,
+        enable_layerwise: Optional[bool] = None,
+        layerwise_eventfd_socket: Optional[str] = None,
+        layerwise_counter_id: int = 0,
     ) -> None:
         self.mode = mode
         self.server_addr = server_addr
@@ -128,6 +142,17 @@ def __init__(
         self.hostkv_wait_timeout_ms = int(hostkv_wait_timeout_ms)
         self.host_kvstorage_fail_policy = host_kvstorage_fail_policy
         self.enable_mps = bool(enable_mps)
+        self.as_batch = bool(as_batch)
+        self.config_path = config_path or ""
+        self.enable_layerwise = enable_layerwise
+        self.layerwise_eventfd_socket = (
+            layerwise_eventfd_socket
+            or os.environ.get(
+                "FLEXKV_LAYERWISE_EVENTFD_SOCKET",
+                DEFAULT_LAYERWISE_EVENTFD_SOCKET,
+            )
+        )
+        self.layerwise_counter_id = int(layerwise_counter_id)
         self.backend_name = "flexkv"
 
         self._gpu_cache_table_list: Optional[List[torch.Tensor]] = None
@@ -139,6 +164,7 @@ def __init__(
         self._adapter = FlexKVClientAdapter(mode, server_addr, server_port)
         self._client = None
         self._ready = False
+        self._layerwise_eventfd_sender: Optional[FlexKVLayerwiseEventfdSender] = None
 
     def register_gpu_cache_tables(self, cache_table_list: List[torch.Tensor]) -> None:
         assert (
@@ -185,7 +211,7 @@ def _init_client(self) -> None:
         if self._client is not None:
             return
         try:
-            from flexkv.common.config import CacheConfig, ModelConfig
+            from flexkv.common.config import CacheConfig, GLOBAL_CONFIG_FROM_ENV, ModelConfig
             from flexkv.kvmanager import KVManager
         except Exception as e:
             raise RuntimeError(f"FlexKV SDK import failed: {e}") from e
@@ -208,6 +234,56 @@ def _init_client(self) -> None:
         if self.num_tmp_cpu_blocks > 0:
             cache_cfg_kwargs["num_tmp_cpu_blocks"] = self.num_tmp_cpu_blocks
         cache_cfg = CacheConfig(**cache_cfg_kwargs)
+        if self.config_path:
+            try:
+                from flexkv.common.config import (
+                    load_user_config_from_file,
+                    update_default_config_from_user_config,
+                )
+            except Exception as e:
+                raise RuntimeError(f"FlexKV config import failed: {e}") from e
+
+            user_cfg = load_user_config_from_file(self.config_path)
+            try:
+                from flexkv.common.config import RankInfo
+
+                rank_info_or_model_cfg = RankInfo(model_config=model_cfg)
+            except ImportError:
+                rank_info_or_model_cfg = model_cfg
+            update_default_config_from_user_config(
+                rank_info_or_model_cfg, cache_cfg, user_cfg
+            )
+        if self.enable_layerwise is None:
+            enable_layerwise_env = os.environ.get("RECSYS_FLEXKV_ENABLE_LAYERWISE", "0")
+            self.enable_layerwise = enable_layerwise_env.strip().lower() in {
+                "1",
+                "true",
+                "yes",
+                "on",
+            }
+        elif isinstance(self.enable_layerwise, str):
+            self.enable_layerwise = self.enable_layerwise.strip().lower() in {
+                "1",
+                "true",
+                "yes",
+                "on",
+            }
+        else:
+            self.enable_layerwise = bool(self.enable_layerwise)
+        GLOBAL_CONFIG_FROM_ENV.enable_layerwise_transfer = bool(self.enable_layerwise)
+        # FlexKV transfer workers are spawned in child processes, so mirror the
+        # in-process config override into env before KVManager starts them.
+        os.environ["FLEXKV_ENABLE_LAYERWISE_TRANSFER"] = (
+            "1" if self.enable_layerwise else "0"
+        )
+        if self.enable_layerwise:
+            os.environ["FLEXKV_LAYERWISE_EVENTFD_SOCKET"] = self.layerwise_eventfd_socket
+            if self._layerwise_eventfd_sender is None:
+                self._layerwise_eventfd_sender = FlexKVLayerwiseEventfdSender(
+                    num_layers=self.num_layers,
+                    socket_path=self.layerwise_eventfd_socket,
+                )
+            self._layerwise_eventfd_sender.start()
         self._client = KVManager(
             model_config=model_cfg,
             cache_config=cache_cfg,
@@ -373,12 +449,18 @@ def onboard_kvcache_launch(
             task_ids=onboard_task_ids,
             uids=torch.tensor(onboard_uids, dtype=torch.int64),
             slot_mappings=onboard_slot_mappings,
+            layer_eventfds=self._layerwise_eventfd_sender.layer_eventfds(
+                self.layerwise_counter_id
+            )
+            if self.enable_layerwise and self._layerwise_eventfd_sender is not None
+            else None,
         )
         onload_task_handle = HostKVTaskHandle(
             backend="flexkv",
             user_ids=onload_handle.uids,
             handle=onload_handle,
             status=HostKVTaskStatus.LAUNCHED,
+            is_layerwise=bool(self.enable_layerwise),
             metadata={
                 "onboard_start_indices": torch.tensor(
                     onboard_start_indices, dtype=torch.int32
@@ -387,7 +469,20 @@ def onboard_kvcache_launch(
             },
         )
 
-        self._client.launch(onload_handle.task_ids, onload_handle.slot_mappings)
+        use_batch = self.as_batch and len(onboard_task_ids) > 1
+        launch_kwargs = {"as_batch": use_batch}
+        if self.enable_layerwise:
+            launch_kwargs.update(
+                {
+                    "layerwise_transfer": True,
+                    "counter_id": self.layerwise_counter_id,
+                }
+            )
+        self._client.launch(
+            onload_handle.task_ids,
+            onload_handle.slot_mappings,
+            **launch_kwargs,
+        )
         return onload_task_handle
 
     def onboard_kvcache_wait(self, task_handle: HostKVTaskHandle) -> HostKVWaitResult:

corelib/recsys_kvcache_manager/recsys_kvcache_manager/flexkv_layerwise.py

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import socket
+import struct
+import threading
+import time
+from array import array
+from ctypes import CDLL, get_errno
+from typing import List, Optional
+
+
+DEFAULT_LAYERWISE_EVENTFD_SOCKET = "/tmp/flexkv_layerwise_eventfd.sock"
+
+
+class FlexKVLayerwiseEventfdSender:
+    """Creates mock layerwise eventfds and sends them to FlexKV's worker."""
+
+    def __init__(
+        self,
+        num_layers: int,
+        socket_path: str,
+        num_counters: int = 3,
+        timeout_s: float = 180.0,
+    ) -> None:
+        self.num_layers = int(num_layers)
+        self.socket_path = socket_path
+        self.num_counters = int(num_counters)
+        self.timeout_s = float(timeout_s)
+        self._eventfds: Optional[List[List[int]]] = None
+        self._thread: Optional[threading.Thread] = None
+
+    @staticmethod
+    def _create_eventfd() -> int:
+        if hasattr(os, "eventfd"):
+            return os.eventfd(0, 0)
+        libc = CDLL(None, use_errno=True)
+        fd = libc.eventfd(0, 0)
+        if fd < 0:
+            raise OSError(get_errno(), "eventfd creation failed")
+        return int(fd)
+
+    def create_eventfds(self) -> List[List[int]]:
+        if self._eventfds is None:
+            # FlexKV layerwise worker expects counter sets for triple buffering.
+            self._eventfds = [
+                [self._create_eventfd() for _ in range(self.num_layers)]
+                for _ in range(self.num_counters)
+            ]
+        return self._eventfds
+
+    def layer_eventfds(self, counter_id: int) -> List[int]:
+        eventfds = self.create_eventfds()
+        if counter_id < 0 or counter_id >= len(eventfds):
+            raise ValueError(
+                f"Invalid layerwise counter_id={counter_id}, "
+                f"expected [0, {len(eventfds)})"
+            )
+        return eventfds[counter_id]
+
+    def start(self) -> None:
+        if self._thread is not None:
+            return
+        self.create_eventfds()
+        self._thread = threading.Thread(
+            target=self._send_eventfds,
+            name="flexkv-layerwise-eventfd-sender",
+            daemon=True,
+        )
+        self._thread.start()
+
+    def _send_eventfds(self) -> None:
+        eventfds = self.create_eventfds()
+        deadline = time.time() + self.timeout_s
+        last_error: Optional[Exception] = None
+        while time.time() < deadline:
+            sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            try:
+                sock.connect(self.socket_path)
+                metadata = struct.pack(
+                    "iiii",
+                    0,
+                    1,
+                    self.num_layers,
+                    len(eventfds),
+                )
+                sock.sendall(metadata)
+                for counter_id, fds in enumerate(eventfds):
+                    fd_array = array("i", fds)
+                    sock.sendmsg(
+                        [struct.pack("i", counter_id)],
+                        [
+                            (
+                                socket.SOL_SOCKET,
+                                socket.SCM_RIGHTS,
+                                fd_array.tobytes(),
+                            )
+                        ],
+                    )
+                ack = sock.recv(1)
+                if ack != b"\x01":
+                    raise RuntimeError(
+                        f"FlexKV layerwise eventfd receiver returned ack={ack!r}"
+                    )
+                return
+            except (FileNotFoundError, ConnectionRefusedError, socket.timeout) as e:
+                last_error = e
+                time.sleep(0.05)
+            finally:
+                sock.close()
+        raise RuntimeError(
+            "Timed out sending mock layerwise eventfds to FlexKV "
+            f"socket {self.socket_path}: {last_error}"
+        )

corelib/recsys_kvcache_manager/recsys_kvcache_manager/kvcache_manager.py

@@ -327,6 +327,32 @@ def _build_host_kvstorage_manager_from_config(
                 }
             else:
                 flexkv_enable_mps = bool(flexkv_enable_mps_raw)
+            flexkv_as_batch_raw = extra.get("flexkv_as_batch", 1)
+            if isinstance(flexkv_as_batch_raw, str):
+                flexkv_as_batch = flexkv_as_batch_raw.strip().lower() in {
+                    "1",
+                    "true",
+                    "yes",
+                    "on",
+                }
+            else:
+                flexkv_as_batch = bool(flexkv_as_batch_raw)
+            flexkv_enable_layerwise = extra.get("flexkv_enable_layerwise", None)
+            if isinstance(flexkv_enable_layerwise, str):
+                flexkv_enable_layerwise = flexkv_enable_layerwise.strip().lower() in {
+                    "1",
+                    "true",
+                    "yes",
+                    "on",
+                }
+            elif flexkv_enable_layerwise is not None:
+                flexkv_enable_layerwise = bool(flexkv_enable_layerwise)
+            flexkv_layerwise_eventfd_socket = extra.get(
+                "flexkv_layerwise_eventfd_socket", None
+            )
+            flexkv_layerwise_counter_id = int(
+                extra.get("flexkv_layerwise_counter_id", 0)
+            )
 
             return FlexKVStorageManager(
                 mode=flexkv_mode,
@@ -343,6 +369,10 @@ def _build_host_kvstorage_manager_from_config(
                 enable_mps=flexkv_enable_mps,
                 host_kvstorage_fail_policy=flexkv_host_kvstorage_fail_policy,
                 hostkv_wait_timeout_ms=int(kvcache_config.offload_timeout_ms),
+                config_path=flexkv_config_path,
+                enable_layerwise=flexkv_enable_layerwise,
+                layerwise_eventfd_socket=flexkv_layerwise_eventfd_socket,
+                layerwise_counter_id=flexkv_layerwise_counter_id,
             )
         else:
             raise NotImplementedError(