feat(kv): add Fluxon KV integration for SGLang

.gitignore

@@ -60,3 +60,6 @@ bench_suite.lock
 deployment/local/
 setup_and_pack/pack_fluxonkv_pylib_env.yaml
 fluxon_rs/moka/
+fluxon_rs/*.btr
+fluxon_rs/moka.bak_non_git_*/
+build_*.rc

fluxon_py/__init__.py

@@ -66,6 +66,9 @@
     "KvFuture",
     "MemHolder",
     "FluxonMemHolder",
+    "GpuBufferRegistration",
+    "GpuDestination",
+    "GpuGetStartHandle",
     # Backend management
     "KvClientType",
     "new_store",
@@ -148,6 +151,9 @@
 
 _LAZY_PYO3 = {
     "FluxonMemHolder": ("kvclient.fluxon", "FluxonMemHolder"),
+    "GpuBufferRegistration": ("kvclient.fluxon", "GpuBufferRegistration"),
+    "GpuDestination": ("kvclient.fluxon", "GpuDestination"),
+    "GpuGetStartHandle": ("kvclient.fluxon", "GpuGetStartHandle"),
 }
 
 

fluxon_py/config.py

@@ -66,6 +66,7 @@ def _yaml_template():
   protocol_type:                       # Protocol type (('tcp'|'rdma'))
   rdma_device_names:                   # Explicit RDMA devices for protocol config (['{str}'](optional))
 pprof_duration_seconds:                # Dump pprof flamegraph after N seconds (int(optional))
+replica_writeback_hot_capacity_ratio: 0.75 # Owner-local hot working-set ratio in the open interval zero to one (float(optional))
 contribute_to_cluster_pool_size:       # Capacity contributed to cluster pool (dict(optional))
   dram: 1677721600                     # - DRAM contribution (int(multiple of 16777216))
   vram:                                # - VRAM contribution per GPU (dict(dynamic_key))
@@ -82,11 +83,18 @@ def _yaml_template():
   prefer_local_placement: false        # Prefer placing new KV writes on the requester-local owner when possible (bool(optional))
   short_circuit_put_payload_path: false # Keep large put_start allocation but skip payload memcpy + transfer (bool(optional))
   skip_put_end_commit: false           # Return success after payload transfer without put_done commit; inflight_put TTL cleanup only (bool(optional))
+  ssd_read_source_policy: legacy_remote_first # legacy_remote_first|local_ssd_only_first (str(optional))
+  owner_local_reserve_soft_wait_timeout_ms: # Local-reserve polling interval, >0 (int(optional))
+  owner_local_reserve_hard_timeout_ms: # Local-reserve end-to-end claim timeout, > soft wait (int(optional))
+  owner_local_reserve_expected_capacity: # Owner-only local-reserve prewarm target (dict(optional))
+    value_len:                         # Canonical value payload bytes, >0 and <=512 MiB (int)
+    payload_capacity_bytes:            # Expected payload capacity to keep resident, >0 (int)
   transport_mode:                      # transfer_only|transfer_with_rpc (str(optional))
   tcp_thread_reactor_shard_count:      # tcp_thread reactor shard count, 1..16 (int(optional))
   tcp_thread_bulk_lane_count:          # tcp_thread bulk lane count, 1..8 (int(optional))
   tcp_thread_control_lane_count:       # tcp_thread control lane count, 1..8 (int(optional))
   user_rpc_sync_handler_thread_count:  # Owner-dedicated sync user-RPC worker thread count, >0 (int(optional))
+  replica_task_max_inflight:           # Deprecated compatibility field, still validated as 1..64; direct remote-Put singleflight does not use a global actor queue (int(optional))
   require_transfer_rpc_fast_path_ready_timeout_seconds: # Require owner-owner transfer-rpc fast path before owner ready/shared.json publication (int(optional))
   rdma_device_names:                   # Explicit RDMA devices for benchmark/test fast-path fanout (['{str}'](optional))
   enable_side_transfer: false          # Enable TCP side-transfer fast-path (bool(optional))
@@ -110,6 +118,9 @@ def _yaml_template():
   cluster_name:                       # Cluster name (str)
   share_mem_path:                     # Shared bundle path for mmap.file/shared.json/peer metadata (str)
   large_file_paths:                   # Owner-mode ordered large-file roots (['{str}'](optional))
+  large_limit_size:                   # Optional per-root SSD capacity in bytes (list(optional))
+  ssd_write_rate_limit_bytes_per_sec: # Optional non-queueing SSD write rate (int(optional))
+  ssd_write_burst_bytes:              # Paired immediate SSD write burst (int(optional))
   p2p_listen_port:                    # P2P QUIC listen port override (int(optional))
   redis_compat:                       # Enable Redis protocol shim (dict(optional))
     listen_addr:                      # TCP listen addr, e.g. "127.0.0.1:16379" (str)
@@ -135,11 +146,16 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]:
         "prefer_local_placement",
         "short_circuit_put_payload_path",
         "skip_put_end_commit",
+        "ssd_read_source_policy",
+        "owner_local_reserve_soft_wait_timeout_ms",
+        "owner_local_reserve_hard_timeout_ms",
+        "owner_local_reserve_expected_capacity",
         "transport_mode",
         "tcp_thread_reactor_shard_count",
         "tcp_thread_bulk_lane_count",
         "tcp_thread_control_lane_count",
         "user_rpc_sync_handler_thread_count",
+        "replica_task_max_inflight",
         "require_transfer_rpc_fast_path_ready_timeout_seconds",
         "rdma_device_names",
         "enable_side_transfer",
@@ -172,6 +188,21 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]:
                 raise ValueError(f"{ctx}.{key} must be a bool")
             out[key] = value
 
+    ssd_read_source_policy = raw.get("ssd_read_source_policy")
+    if ssd_read_source_policy is not None:
+        if not isinstance(ssd_read_source_policy, str):
+            raise ValueError(f"{ctx}.ssd_read_source_policy must be a string")
+        allowed_ssd_read_source_policies = {
+            "legacy_remote_first",
+            "local_ssd_only_first",
+        }
+        if ssd_read_source_policy not in allowed_ssd_read_source_policies:
+            raise ValueError(
+                f"{ctx}.ssd_read_source_policy must be one of "
+                f"{sorted(allowed_ssd_read_source_policies)}, got {ssd_read_source_policy!r}"
+            )
+        out["ssd_read_source_policy"] = ssd_read_source_policy
+
     transport_mode = raw.get("transport_mode")
     transport_mode_was_explicit = transport_mode is not None
     side_transfer_role_raw = raw.get("side_transfer_role")
@@ -234,6 +265,7 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]:
         ("tcp_thread_reactor_shard_count", 1, 16),
         ("tcp_thread_bulk_lane_count", 1, 8),
         ("tcp_thread_control_lane_count", 1, 8),
+        ("replica_task_max_inflight", 1, 64),
     ):
         value = raw.get(key)
         if value is None:
@@ -254,6 +286,50 @@ def _normalize_test_spec_config(raw: Any, ctx: str) -> Dict[str, Any]:
             raise ValueError(f"{ctx}.user_rpc_sync_handler_thread_count must be > 0")
         out["user_rpc_sync_handler_thread_count"] = user_rpc_sync_handler_thread_count
 
+    reserve_timeouts: Dict[str, int] = {}
+    for key in (
+        "owner_local_reserve_soft_wait_timeout_ms",
+        "owner_local_reserve_hard_timeout_ms",
+    ):
+        value = raw.get(key)
+        if value is None:
+            continue
+        if isinstance(value, bool) or not isinstance(value, int):
+            raise ValueError(f"{ctx}.{key} must be an int")
+        if value <= 0:
+            raise ValueError(f"{ctx}.{key} must be > 0")
+        reserve_timeouts[key] = value
+        out[key] = value
+    soft_timeout_ms = reserve_timeouts.get("owner_local_reserve_soft_wait_timeout_ms", 10)
+    hard_timeout_ms = reserve_timeouts.get("owner_local_reserve_hard_timeout_ms", 10_000)
+    if hard_timeout_ms <= soft_timeout_ms:
+        raise ValueError(
+            f"{ctx}.owner_local_reserve_hard_timeout_ms must be greater than "
+            f"{ctx}.owner_local_reserve_soft_wait_timeout_ms"
+        )
+
+    expected_capacity = raw.get("owner_local_reserve_expected_capacity")
+    if expected_capacity is not None:
+        expected_ctx = f"{ctx}.owner_local_reserve_expected_capacity"
+        if not isinstance(expected_capacity, dict):
+            raise ValueError(f"{expected_ctx} must be a mapping")
+        unknown_expected = sorted(
+            set(expected_capacity.keys()) - {"value_len", "payload_capacity_bytes"}
+        )
+        if unknown_expected:
+            raise ValueError(f"{expected_ctx} contains unknown keys: {unknown_expected}")
+        normalized_expected: Dict[str, int] = {}
+        for key in ("value_len", "payload_capacity_bytes"):
+            value = expected_capacity.get(key)
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise ValueError(f"{expected_ctx}.{key} must be an int")
+            if value <= 0:
+                raise ValueError(f"{expected_ctx}.{key} must be > 0")
+            normalized_expected[key] = value
+        if normalized_expected["value_len"] > 512 * 1024 * 1024:
+            raise ValueError(f"{expected_ctx}.value_len must be <= 536870912")
+        out["owner_local_reserve_expected_capacity"] = normalized_expected
+
     side_transfer_worker_count = raw.get("side_transfer_worker_count")
     if side_transfer_worker_count is not None:
         if isinstance(side_transfer_worker_count, bool) or not isinstance(side_transfer_worker_count, int):
@@ -348,6 +424,23 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None:
         raise ValueError("fluxonkv_spec must be a mapping")
 
     is_zero_contribution = _is_zero_contribution_fluxonkv_config(cfg)
+    test_spec_config = cfg.get("test_spec_config") or {}
+    expected_capacity = test_spec_config.get("owner_local_reserve_expected_capacity")
+    hot_capacity_ratio = cfg.get("replica_writeback_hot_capacity_ratio")
+
+    if hot_capacity_ratio is not None:
+        if isinstance(hot_capacity_ratio, bool) or not isinstance(
+            hot_capacity_ratio, (int, float)
+        ):
+            raise ValueError(
+                "replica_writeback_hot_capacity_ratio must be a number in (0, 1)"
+            )
+        hot_capacity_ratio = float(hot_capacity_ratio)
+        if not 0.0 < hot_capacity_ratio < 1.0:
+            raise ValueError(
+                "replica_writeback_hot_capacity_ratio must be finite and in (0, 1)"
+            )
+        cfg["replica_writeback_hot_capacity_ratio"] = hot_capacity_ratio
 
     share_mem_path = spec.get("share_mem_path")
     if not isinstance(share_mem_path, str) or not share_mem_path.strip():
@@ -360,11 +453,22 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None:
         raise ValueError("fluxonkv_spec.transfer_engine has been removed from Fluxon KV config")
 
     if is_zero_contribution:
+        if hot_capacity_ratio is not None:
+            raise ValueError(
+                "replica_writeback_hot_capacity_ratio is only valid on owner configs"
+            )
+        if expected_capacity is not None:
+            raise ValueError(
+                "test_spec_config.owner_local_reserve_expected_capacity is only valid on owner configs"
+            )
         forbidden_spec_keys = [
             "etcd_addresses",
             "redis_compat",
             "sub_cluster",
             "large_file_paths",
+            "large_limit_size",
+            "ssd_write_rate_limit_bytes_per_sec",
+            "ssd_write_burst_bytes",
         ]
         for key in forbidden_spec_keys:
             if key in spec:
@@ -379,6 +483,23 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None:
     if int(contrib["dram"]) == 0:
         raise ValueError("owner mode requires non-zero contribute_to_cluster_pool_size.dram")
 
+    if expected_capacity is not None:
+        value_len = int(expected_capacity["value_len"])
+        payload_capacity_bytes = int(expected_capacity["payload_capacity_bytes"])
+        slot_size = max(value_len, 4096)
+        slot_size = (slot_size + 4095) & ~4095
+        slots_per_grant = (512 * 1024 * 1024) // slot_size
+        value_count = (payload_capacity_bytes + value_len - 1) // value_len
+        expected_grants = (value_count + slots_per_grant - 1) // slots_per_grant
+        physical_reserve_bytes = expected_grants * 512 * 1024 * 1024
+        owner_dram_bytes = int(contrib["dram"])
+        if physical_reserve_bytes > owner_dram_bytes:
+            raise ValueError(
+                "test_spec_config.owner_local_reserve_expected_capacity requires "
+                f"{physical_reserve_bytes} physical bytes across {expected_grants} grants, "
+                f"exceeding owner dram contribution {owner_dram_bytes}"
+            )
+
     if "etcd_addresses" not in spec:
         raise ValueError("fluxonkv_spec.etcd_addresses is required for owner mode")
     etcd_addresses = spec.get("etcd_addresses")
@@ -404,6 +525,26 @@ def _validate_fluxonkv_contract(cfg: Dict[str, Any]) -> None:
                 f"fluxonkv_spec.large_file_paths[{idx}] must be a non-empty string in owner mode"
             )
 
+    write_rate = spec.get("ssd_write_rate_limit_bytes_per_sec")
+    write_burst = spec.get("ssd_write_burst_bytes")
+    if (write_rate is None) != (write_burst is None):
+        raise ValueError(
+            "fluxonkv_spec.ssd_write_rate_limit_bytes_per_sec and "
+            "ssd_write_burst_bytes must be configured together"
+        )
+    if write_rate is not None:
+        if (
+            isinstance(write_rate, bool)
+            or not isinstance(write_rate, int)
+            or write_rate <= 0
+            or isinstance(write_burst, bool)
+            or not isinstance(write_burst, int)
+            or write_burst <= 0
+        ):
+            raise ValueError("SSD write rate and burst must both be positive integers")
+        if "large_limit_size" not in spec:
+            raise ValueError("SSD write admission requires fluxonkv_spec.large_limit_size")
+
 
 class FluxonKvClientConfig():
     """Configuration class for KV Cache stores that reads from YAML config files."""
@@ -434,6 +575,10 @@ def __init__(self, config_dict: Dict[str, Any]):
             raise ValueError(
                 "exactly one of [mooncake_spec, fluxonkv_spec] is required (and the chosen spec must not be null)"
             )
+        if "replica_writeback_hot_capacity_ratio" in plain and not has_fluxon:
+            raise ValueError(
+                "replica_writeback_hot_capacity_ratio requires fluxonkv_spec"
+            )
 
         pprof_duration_seconds = plain.get("pprof_duration_seconds")
         if pprof_duration_seconds is None:
@@ -750,8 +895,8 @@ def parse_type_recursive(type_str: str) -> Optional[Tuple[str, Dict[str, Any]]]:
                         return parsed_type_name, merged_params
                     return type_name, {"constraint": constraint}
 
-        # 6) Primitive types: str, int, bool, None
-        if type_str in ["str", "int", "bool", "None"]:
+        # 6) Primitive types: str, int, float, bool, list, None
+        if type_str in ["str", "int", "float", "bool", "list", "None"]:
             return type_str, {}
 
         debug_print("type_str ", type_str, "not matched to any type")
@@ -871,13 +1016,29 @@ def raise_validation_error(msg: str):
                             raise_validation_error(f"Value must be multiple of {multiple}, got {value}")
                         else:
                             return None
-
+
+    elif type_name == "float":
+        if isinstance(value, bool) or not isinstance(value, (int, float)):
+            if raise_err:
+                raise_validation_error(
+                    f"Expected float-compatible number, got {type(value).__name__}"
+                )
+            else:
+                return None
+
     elif type_name == "bool":
         if not isinstance(value, bool):
             if raise_err:
                 raise_validation_error(f"Expected bool, got {type(value).__name__}")
             else:
                 return None
+
+    elif type_name == "list":
+        if not isinstance(value, list):
+            if raise_err:
+                raise_validation_error(f"Expected list, got {type(value).__name__}")
+            else:
+                return None
 
     elif type_name == "dict":
         if not isinstance(value, dict):