diff --git a/python/sglang/srt/arg_groups/model_path_hook.py b/python/sglang/srt/arg_groups/model_path_hook.py index 0093fac31515..b54508ce21f8 100644 --- a/python/sglang/srt/arg_groups/model_path_hook.py +++ b/python/sglang/srt/arg_groups/model_path_hook.py @@ -275,16 +275,8 @@ def handle_load_format(server_args: Any): "that selects IPC loading automatically." ) - # Speculative decoding loads an extra draft model whose weights the - # daemon does not export, so refuse the combination up front instead of - # failing deep inside draft-worker load (draft-model daemon TBD). - if cfg.weight_cache_mode != "off" and cfg.speculative_algorithm is not None: - raise ValueError( - "--weight-cache-mode is not supported together with speculative " - "decoding (--speculative-algorithm): the weight cache daemon does " - "not export the draft model's weights. Disable one of them " - "(--weight-cache-mode off) for this configuration." - ) + # The target uses IPC; draft workers keep their own ordinary disk loader. + # Role isolation is enforced in build_load_config, not on shared ServerArgs. def validate_transfer_engine(server_args: Any): diff --git a/python/sglang/srt/layers/engram.py b/python/sglang/srt/layers/engram.py index 68c129933623..76ca42bf0955 100644 --- a/python/sglang/srt/layers/engram.py +++ b/python/sglang/srt/layers/engram.py @@ -267,6 +267,52 @@ def init_history(self, num_req_slots: int, device) -> None: def from_config( cls, config, layout: EngramLayout, *, image_token_id: Optional[int] = None ) -> EngramHasher: + if ( + torch.empty(0).device.type == "meta" + and get_model().weight_cache_mode != "off" + ): + # IPC supplies the already-built hash buffers and pad_id. Rebuilding + # the HF tokenizer map here would dominate a zero-copy restart. + result = cls.__new__(cls) + nn.Module.__init__(result) + result.max_ngram_size = layout.max_ngram_size + result.pad_id = 0 # replaced from the post-load manifest + result.image_token_id = image_token_id + result.history = None # request state always belongs to this engine + result.pad_row = 0 + result.register_buffer( + "token_map", + torch.empty(config.vocab_size, dtype=torch.int64), + persistent=False, + ) + result.register_buffer( + "multipliers", + torch.empty( + len(layout.layer_ids), layout.max_ngram_size, dtype=torch.int64 + ), + persistent=False, + ) + result.register_buffer( + "primes", + torch.empty( + len(layout.layer_ids), + layout.max_ngram_size - 1, + layout.n_heads, + dtype=torch.int64, + ), + persistent=False, + ) + result.register_buffer( + "offsets", + torch.empty( + len(layout.layer_ids), + (layout.max_ngram_size - 1) * layout.n_heads, + dtype=torch.int64, + ), + persistent=False, + ) + return result + # The compressed token map is built with the HF normalizers, so the HF # tokenizer backend is used here whatever the serving backend is. tokenizer = get_tokenizer( @@ -585,10 +631,24 @@ def _hip_runtime() -> ctypes.CDLL: def _registered_device_ptr(host_ptr: int) -> int: """ Address kernels must use for registered host memory. - UVA makes it the host address on CUDA; HIP may map it elsewhere. + Resolve the address for this process/context; do not reuse an exporter pointer. """ if not _is_hip: - return host_ptr + path = find_loaded_library("libcudart") + if path is None: + raise RuntimeError("libcudart is not loaded in the current process") + lib = ctypes.CDLL(path) + lib.cudaHostGetDevicePointer.restype = ctypes.c_int + lib.cudaHostGetDevicePointer.argtypes = [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_void_p, + ctypes.c_uint, + ] + device_ptr = ctypes.c_void_p() + err = lib.cudaHostGetDevicePointer(ctypes.byref(device_ptr), host_ptr, 0) + if err != 0 or not device_ptr.value: + raise RuntimeError(f"cudaHostGetDevicePointer failed: {err}") + return device_ptr.value device_ptr = ctypes.c_void_p() err = _hip_runtime().hipHostGetDevicePointer(ctypes.byref(device_ptr), host_ptr, 0) if err != 0 or not device_ptr.value: @@ -599,8 +659,8 @@ def _registered_device_ptr(host_ptr: int) -> int: class _HostTable: """Host-memory backing for one engram table ('shared' or 'per_rank' layout). - Lives for the whole process: the mapping, the memfd and the cudaHostRegister - pin are never released because the table is read by every forward. + Ordinary registered tables keep their process lifetime. Imported and VMM + mappings belong to the model and must outlive all tensor/kernel users. """ def __init__(self, layout: str, nbytes: int, name: str, group): @@ -613,8 +673,23 @@ def __init__(self, layout: str, nbytes: int, name: str, group): self.nbytes = nbytes self.group = group self.dirty = False - if layout == "shared": - self.fd = self._open_shared_fd(nbytes, name) + from sglang.srt.runtime_context import get_server_args, publish_role + from sglang.srt.weight_cache.model_state import host_table_backend + + self.backend = host_table_backend(get_server_args().model_loader_extra_config) + if self.backend == "cuda_vmm": + self._init_vmm() + return + + if layout == "shared" or publish_role() == "weight_cache_daemon": + if layout == "shared": + self.fd = self._open_shared_fd(nbytes, name) + else: + # An anonymous MAP_PRIVATE shard cannot survive worker restart. + # A daemon-owned memfd keeps the same per-rank row ownership + # while allowing clients to map the physical pages via an FD. + self.fd = os.memfd_create(name, 0) + os.ftruncate(self.fd, nbytes) self.mm = mmap.mmap( self.fd, nbytes, @@ -635,7 +710,8 @@ def __init__(self, layout: str, nbytes: int, name: str, group): if layout == "per_rank": # Cached checkpoint pages, left by a previous server or by the loader, # make the 512 MiB huge-page faults fall back, so empty them first. - _drop_page_cache_once("before pre-faulting the per-rank shard") + if self.fd is None: + _drop_page_cache_once("before pre-faulting the per-rank shard") np.frombuffer(self.mm, dtype=np.uint8)[:: mmap.PAGESIZE] = 0 if layout == "shared": # Every rank holds the fd before rank 0 continues; the /proc path only @@ -646,6 +722,127 @@ def __init__(self, layout: str, nbytes: int, name: str, group): raise RuntimeError(f"cudaHostRegister({nbytes} bytes) failed: {err}") self.device_ptr = _registered_device_ptr(self.bytes.data_ptr()) + def _init_vmm(self): + from multiprocessing.reduction import DupFd + + from sglang.srt.weight_cache.host_vmm import HostVmmAllocation + + if self.layout == "per_rank" or self.group.rank_in_group == 0: + self._vmm = HostVmmAllocation.create(self.nbytes) + if self.layout == "shared": + payload = None + if self.group.rank_in_group == 0: + payload = ( + self._vmm.allocation_nbytes, + self._vmm.numa_id, + [DupFd(self._vmm.fd) for _ in range(self.group.world_size - 1)], + ) + allocation_nbytes, numa_id, handles = self.group.broadcast_object( + payload, src=0 + ) + if self.group.rank_in_group != 0: + fd = handles[self.group.rank_in_group - 1].detach() + self._vmm = HostVmmAllocation.from_fd( + fd, self.nbytes, allocation_nbytes, numa_id + ) + self.group.barrier() + self.fd = self._vmm.fd + self.bytes = self._vmm.bytes + self.device_ptr = self._vmm.device_ptr + + @classmethod + def from_vmm_fd(cls, layout, nbytes, fd, allocation_nbytes, numa_id): + from sglang.srt.weight_cache.host_vmm import HostVmmAllocation + + if layout not in ("shared", "per_rank"): + os.close(fd) + raise ValueError("Invalid host VMM layout") + result = cls.__new__(cls) + result.layout, result.nbytes = layout, nbytes + result.group, result.dirty, result.backend = None, False, "cuda_vmm" + result._vmm = HostVmmAllocation.from_fd( + fd, nbytes, allocation_nbytes, numa_id, gpu_view=True + ) + result.fd = result._vmm.fd + result.bytes = result._vmm.bytes + result.device_ptr = result._vmm.device_ptr + return result + + @classmethod + def from_fd(cls, layout: str, nbytes: int, fd: int): + """Adopt a daemon-owned table FD and register this process's mapping. + + No page pre-fault, checkpoint copy, or huge-page collapse is performed. + FD ownership transfers to this object, including on registration failure. + """ + if layout not in ("shared", "per_rank") or nbytes <= 0: + os.close(fd) + raise ValueError("Invalid imported Engram table layout/size") + result = cls.__new__(cls) + result.layout, result.nbytes, result.fd = layout, nbytes, fd + result.backend = "memfd" + result.group, result.dirty = None, False + result._ipc_registered = False + try: + if os.fstat(fd).st_size != nbytes: + raise ValueError("Engram memfd size does not match its descriptor") + result.mm = mmap.mmap( + fd, + nbytes, + flags=mmap.MAP_SHARED, + prot=mmap.PROT_READ | mmap.PROT_WRITE, + ) + result.bytes = torch.frombuffer(result.mm, dtype=torch.uint8) + tic = time.perf_counter() + err = torch.cuda.cudart().cudaHostRegister( + result.bytes.data_ptr(), nbytes, 0 + ) + if int(err) != 0: + raise RuntimeError( + f"cudaHostRegister({nbytes} imported bytes) failed: {err}" + ) + result._ipc_registered = True + result.device_ptr = _registered_device_ptr(result.bytes.data_ptr()) + logger.info( + "Engram IPC host registration: layout=%s, bytes=%d, elapsed=%.3fs", + layout, + nbytes, + time.perf_counter() - tic, + ) + return result + except Exception: + result.close_import() + raise + + def close_import(self): + """Release a client registration after all kernel/tensor users retire.""" + if getattr(self, "_vmm", None) is not None: + if hasattr(self, "bytes"): + del self.bytes + self._vmm.close() + self.fd = None + return + if getattr(self, "_ipc_registered", False): + torch.cuda.cudart().cudaHostUnregister(self.bytes.data_ptr()) + self._ipc_registered = False + if hasattr(self, "bytes"): + del self.bytes + try: + if hasattr(self, "mm"): + self.mm.close() + finally: + if self.fd is not None: + os.close(self.fd) + self.fd = None + + def __del__(self): + # Ordinary loader tables keep their existing process-lifetime policy. + if hasattr(self, "_ipc_registered") or hasattr(self, "_vmm"): + try: + self.close_import() + except Exception: + pass + def _open_shared_fd(self, nbytes: int, name: str) -> int: owner = None if self.group.rank_in_group == 0: @@ -691,8 +888,17 @@ def finish_load(self, label: str): self.dirty = False if self.layout == "shared": self.group.barrier() + if getattr(self, "backend", "memfd") == "cuda_vmm": + logger.info( + "engram host table %s: layout=%s, CUDA HOST_NUMA VMM, logical_bytes=%d, numa=%d", + label, + self.layout, + self.nbytes, + self._vmm.numa_id, + ) + return mapped_kb, huge_kb = _huge_pages_backing(self.bytes.data_ptr()) - if self.layout == "per_rank" and huge_kb < mapped_kb * 0.98: + if self.layout == "per_rank" and self.fd is None and huge_kb < mapped_kb * 0.98: # The loader's own reads refilled the page cache; empty it again so the # collapse can find contiguous memory. drop_checkpoint_page_cache() @@ -705,7 +911,7 @@ def finish_load(self, label: str): f"({pct:.0f}%)" ) if huge_kb == 0: - knob = "shmem_enabled" if self.layout == "shared" else "enabled" + knob = "shmem_enabled" if self.fd is not None else "enabled" logger.warning( "%s. No huge pages: expect ~10x slower lookups (one TLB miss per row); " "transparent_hugepage/%s is '%s'", @@ -756,6 +962,17 @@ def _init_host_table(self, num_embeddings: int, dim: int, layer_id: int): n = num_embeddings if layout == "shared" else self.rows w_bytes = n * dim s_bytes = n * (dim // FP8_BLOCK_SIZE) + if torch.empty(0).is_meta: + # IPC restores the host mapping after the skeleton is constructed. + # Allocating/registering a second huge host table here defeats it. + self.weight = nn.Parameter( + torch.empty(n, dim, dtype=torch.float8_e4m3fn), requires_grad=False + ) + self.scale = nn.Parameter( + torch.empty(n, dim // FP8_BLOCK_SIZE, dtype=torch.float8_e8m0fnu), + requires_grad=False, + ) + return self.host_table = _HostTable( layout, max(1, w_bytes + s_bytes), # mmap requires storage even for an empty shard. diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 5aa907ecbc9b..4c4d0b15a9e5 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1242,6 +1242,7 @@ def load_model(self): draft_model_idx=self.draft_model_idx, weight_cache_mode=get_model().weight_cache_mode, weight_cache_socket=get_model().weight_cache_socket, + is_draft_worker=self.is_draft_worker, load_group=self.load_group, ) diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index 1210e458dbaa..c8ccd6ca93a3 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -208,6 +208,7 @@ def build_load_config( weight_cache_mode: str, weight_cache_socket: str | None, load_group: LoadGroup = _DEFAULT_LOAD_GROUP, + is_draft_worker: bool = False, ) -> LoadConfig: from sglang.srt.configs.modelopt_config import ModelOptConfig @@ -235,8 +236,8 @@ def build_load_config( modelopt_config=modelopt_config, rl_quant_profile=get_model().rl_quant_profile, draft_model_idx=draft_model_idx, - weight_cache_mode=weight_cache_mode, - weight_cache_socket=weight_cache_socket, + weight_cache_mode="off" if is_draft_worker else weight_cache_mode, + weight_cache_socket=None if is_draft_worker else weight_cache_socket, load_group=load_group, ) @@ -252,7 +253,7 @@ def maybe_enable_ipc_weight_cache( so a second call (e.g. a weight reload) can't overwrite the captured fallback format. """ - if get_model().weight_cache_mode == "off": + if load_config.weight_cache_mode == "off": return if load_config.load_format != LoadFormat.IPC_CACHE: @@ -310,7 +311,7 @@ def load_model_with_memory_saver( # In zero-copy IPC mode, the weights are shared with the daemon via # CUDA IPC and must not be offloaded/reloaded by the memory saver. - is_ipc_zero_copy = get_model().weight_cache_mode != "off" + is_ipc_zero_copy = load_config.weight_cache_mode != "off" if is_ipc_zero_copy and enable_cpu_backup: logger.warning( "[ModelRunner] Disabling weights CPU backup in zero-copy IPC mode — " diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index a916a0c32f64..2477b58ee910 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -365,7 +365,19 @@ def _validate_default_loader_extra_config( # options, so newer versions can add parameters without a SGLang update. return - allowed_keys = {"enable_multithread_load", "num_threads"} + allowed_keys = { + "enable_multithread_load", + "num_threads", + "dsv41_weight_cache", + "dsv41_host_table_backend", + } + from sglang.srt.weight_cache.model_state import host_table_backend + + host_table_backend(extra_config) + if "dsv41_weight_cache" in extra_config and not isinstance( + extra_config["dsv41_weight_cache"], bool + ): + raise ValueError("dsv41_weight_cache must be a boolean") if load_format == LoadFormat.FASTSAFETENSORS: allowed_keys.add("enable_gds") if "enable_gds" in extra_config and not isinstance( diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 5b82f536a5eb..999dcb146313 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2806,9 +2806,12 @@ def bind_next(self, next_layer: Optional[nn.Module]) -> None: """Save the next layer for mhc. Must be called before `_init_boundaries`""" if not self.hc_pre_from_prev_sublayer: return - self._next_layer = None + # A non-owning seam reference must not register the next layer as a + # child: that duplicates every downstream weight in state_dict/IPC. + self._modules.pop("_next_layer", None) + object.__setattr__(self, "_next_layer", None) if isinstance(next_layer, DeepseekV4DecoderLayer) and next_layer.engram is None: - self._next_layer = next_layer + object.__setattr__(self, "_next_layer", next_layer) def _init_hyper_connections(self) -> None: """(Re)bind the two hyper-connections to the current mixing weights.""" diff --git a/python/sglang/srt/weight_cache/README.md b/python/sglang/srt/weight_cache/README.md new file mode 100644 index 000000000000..246de6379e44 --- /dev/null +++ b/python/sglang/srt/weight_cache/README.md @@ -0,0 +1,192 @@ +# Experimental DeepSeek V4.1 recovery + +This extension extends the existing Weight Cache Daemon feature introduced in +[#27139](https://github.com/sgl-project/sglang/pull/27139) and described in +[Fast Engine Recovery: Sub-Second Engine Restart for SGLang via Weight Cache Daemon](https://www.lmsys.org/blog/2026-08-21-sglang-fast-recovery) +to DeepSeek V4.1, with Engram in either GPU or host memory. + +The article addresses the repeated checkpoint I/O and weight processing paid on +every engine restart, which prolong service outages during worker failures and +rolling updates. It retains processed weights in a persistent daemon and lets +replacement engines reuse them through CUDA IPC. This extension builds on that feature +and addresses the gaps that prevent the same recovery path from working for V4.1. + +**Why this extension is needed:** + +- **V4.1 needs more than tensor import.** Its FP8 dense / FP4 expert path has + post-load layouts, quantization metadata and cached tensor references that a + replacement engine must restore consistently. Related failures are documented + in [#39794](https://github.com/sgl-project/sglang/pull/39794). This extension restores + the processed state, rebinds mHC and TopK references, and checks backend/layout + compatibility instead of merely relaxing the IPC admission check or repeating + transforms on shared weights. +- **Host Engram needs its own resident recovery path.** GPU tensor IPC cannot + share its host backing. This extension preserves daemon-owned host tables and imports + them in replacement workers; optional HOST_NUMA VMM reuses allocation handles + to avoid repeated `cudaHostRegister` on unchanged tables. +- **Faster checkpoint reads still pay checkpoint I/O on restart.** The + InstantTensor and positioned-read PRs below accelerate loading from storage. + This extension complements them by reusing already processed GPU weights and host + Engram after a worker restart, while the daemon survives. Their cold-load + optimizations can still benefit initial daemon startup, but are not included + in this patch. + +The result is a V4.1-specific implementation of the article's resident-weight +recovery path, plus an experimental prewarmed standby for its active-standby +scenario. Complete weight reload and service recovery results are reported below; +full long-context acceptance and other remaining gates are still pending. + +### Related work + +- [#27139 — Support fast engine recovery through weight cache](https://github.com/sgl-project/sglang/pull/27139): the original feature this extension extends. +- [#38306 — Support post-load lifecycle and MXFP4 IPC](https://github.com/sgl-project/sglang/pull/38306): related model-support work describing K3/DSV3 validation; this extension adds V4.1-specific state restoration and host Engram recovery. +- [#39794 — Reject routed MXFP4 experts for IPC](https://github.com/sgl-project/sglang/pull/39794): documents why accepting a base FP8 configuration is insufficient for routed-expert post-load layouts. +- [#40453 — Add InstantTensor loader support](https://github.com/sgl-project/sglang/pull/40453), and [#42919 — Load V4.1 with the Engram host table](https://github.com/sgl-project/sglang/pull/42919): accelerate checkpoint loading while keeping host-resident tables off the GPU stream. +- [#42920 — Parallel positioned reads for Engram host tables](https://github.com/sgl-project/sglang/pull/42920): accelerates storage-to-host loading with parallel `preadv`; this extension addresses resident recovery and mapping reuse instead. + +## Host Engram configuration + +Use Linux/NVIDIA CUDA with the repository's supported dependencies. The optional +`cuda_vmm` backend additionally requires `cuda.bindings.driver`, device support +for HOST_NUMA VMM and POSIX FD export/import. Unsupported devices fail explicitly. + +Set the same environment, checkpoint, source, backend choices, and model overrides +in the daemon and every client. For example, for a local TP4/EP1 deployment: + +```bash +export MODEL_PATH=/path/to/DeepSeek-V4.1-Flash +export SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE=1 +export SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT=per_rank + +common=( + --model-path "$MODEL_PATH" --trust-remote-code + --tp 4 --attn-dp-size 1 --ep-size 1 + --moe-a2a-backend none --moe-runner-backend flashinfer_mxfp4 + --attention-backend dsv4 + --json-model-override-args '{"vision_n_layers":0}' + --model-loader-extra-config '{"dsv41_weight_cache":true,"dsv41_host_table_backend":"cuda_vmm"}' + --context-length 4096 --max-total-tokens 8192 --max-running-requests 4 + --chunked-prefill-size 4096 --max-prefill-tokens 4096 + --mem-fraction-static 0.65 --random-seed 42 + --disable-cuda-graph --disable-prefill-cuda-graph --disable-overlap-schedule +) +# Terminal 1: load and retain the processed model. +python -m sglang.srt.weight_cache.daemon "${common[@]}" +# Terminal 2: import resident weights; wait for the daemon to be ready first. +python -m sglang.launch_server "${common[@]}" \ + --weight-cache-mode client --host 127.0.0.1 --port 22941 --nccl-port 22942 +``` + +These commands describe the tested configuration; capacity and backend support +must be checked for your hardware. Omit the host-table environment settings for +GPU-resident Engram. Select `"dsv41_host_table_backend":"memfd"` for the comparison +backend. The opt-in `dsv41_weight_cache` flag is required for V4.1 IPC recovery. + +VMM allocates NUMA host storage once in the daemon, exports its allocation FD, +and imports it in each worker. Clients expose a CUDA tensor alias to physical +host memory without copying the table or calling `cudaHostRegister` again. Each +process still pays for its own GPU access mapping. CPU checkpoint loading and +initial allocation remain cold-start costs. + +The post-load manifest restores tensor aliases, cached tensor references, +quantization flags, mHC seams and TopK correction bias references. The contract +and layout fingerprint reject incompatible backend/configuration/source choices. +Host descriptors validate logical size, padded allocation size, NUMA placement, +and tensor bounds. Host-table objects own imported mappings, CUDA handles and +FDs; retire tensor/kernel users before closing these objects. + +When upgrading this implementation, restart the weight daemon before importing +its cache with upgraded clients. The compatibility fingerprint includes the +daemon, IPC loader, model-load orchestration and weight-processing source files; +a client rejects a cache produced by different recovery code. This intentionally +requires a fresh daemon load after source changes. + +## Prewarmed standby workers + +`python -m sglang.srt.weight_cache.standby --config config.json` manages two +independent client process groups. A worker becomes ready only after HTTP health +and a generation warmup. It promotes an already ready worker after active-worker +failure and replenishes the failed worker in the background. Manage the weight +daemon separately. + +A configuration has two `workers`, each with a distinct `name`, loopback HTTP +`url`, and a `command` argument list that launches its client server. Optional +worker `env` entries extend the inherited environment. Use distinct HTTP/NCCL +ports and budget both workers' private KV pools. Other settings include `host`, +`port`, `log_dir`, `warmup`, `health_path`, `health_interval`, `health_timeout`, +`poll_interval`, `startup_timeout`, `warmup_timeout`, and `restart_backoff`. +Inspect `/recovery/status` for readiness and promotion events. + +Only stateless generation/OpenAI inference routes and model metadata are exposed. +Requests containing `session_params` are rejected. Failed requests are never +replayed; truncated streams are closed. HTTP health does not prove GPU forward +progress. Two workers sharing the same GPU do not isolate device failures, and +only one worker remains available while the standby is rebuilding. + +## Validation and limitations + +The experiment used four NVIDIA L20D GPUs, TP4/EP1/attention-DP1, +DeepSeek-V4.1-Flash native FP8 dense/FP4 expert weights, context 4096 and a maximum +of four requests. CUDA Graph and overlap were disabled. + +### Engram in GPU memory + +Disk reload compared with resident CUDA IPC recovery: + +| Metric | Before | After | Speedup | +| --- | --- | --- | --- | +| Complete main-model weight reload | **50.56 s** | **1.02 s** | **49.6×** | +| Complete service recovery: process launch → healthy service | **95.09 s** | **36.53 s** | **2.60×** | + +Values are medians: two final matched disk runs and three recovery runs. +Compilation/autotuning caches were warm. The weight reload covers the entire +main model, including Engram and post-load state; service recovery also includes +communication, KV and runtime initialization. + +### Engram in shared host memory + +| Host layout and comparison | Metric | Before | After | Speedup | +| --- | --- | --- | --- | --- | +| Shared table: disk reload → resident IPC + memfd | Complete main-model weight reload | **85.45 s** | **14.44 s** | **5.9×** | +| Per-rank shards: existing memfd recovery → HOST_NUMA VMM | Complete main-model weight reload | **4.77 s** | **3.36 s** | **1.42×** | +| Per-rank shards: existing memfd recovery → HOST_NUMA VMM | Launch-helper return → polled service readiness | **45.54 s** | **40.53 s** | **1.12×** | + +Host recovery values are medians of three runs; the shared-table disk baseline +has one sample. Shared tables and per-rank shards are separate experiments: +85.45 s is not a matched baseline for 3.36 s. Exact process-launch to healthy +service time was not measured for the host VMM comparison; its readiness timer +starts after the launch helper returns. + +Both placements require an already loaded, surviving daemon for optimized +recovery; initial daemon startup is excluded. + +Across 30 prompt comparisons, all text and input/output token IDs matched. +24/30 also matched logprobs exactly. The 1401-token case differed even in a memfd +same-service repeat; strict long-context numerical acceptance remains unresolved. +8/51/203 GB allocation probes matched an independent CPU gather oracle. Four-rank +shared-table gather and CUDA Graph replay passed at component scope only; full +model shared layout and full model CUDA Graphs were not validated. + +Earlier memfd standby fault injection measured 205–422 ms from launcher SIGKILL +to the first response token. This is takeover by an already warmed worker, not +VMM worker reconstruction. VMM plus standby has not had a combined full-model +fault-injection run. + +The experiment used driver 580.105.08, CUDA driver API 13.0, Torch 2.13.0+cu130, +FlashInfer 0.6.18 and sglang-kernel 0.4.7. The frozen source expected kernel 0.4.9; +its version check was explicitly bypassed. This dependency mismatch is an open +validation gap, not a recommended deployment configuration. + +Full-capacity throughput, speculative/draft execution, other topologies, +multi-node operation, daemon loss and in-flight state recovery are unvalidated. +This extension does not implement the parallel checkpoint-read optimization in +SGLang #42920. Keep this feature experimental until the remaining gates pass. + +The CPU contract tests can be run individually: + +```bash +for test in test_host_vmm_weight_cache test_host_engram_weight_cache \ + test_dsv41_weight_cache_state test_weight_cache_protocol test_warm_standby; do + python "test/registered/unit/model_loader/${test}.py" +done +``` diff --git a/python/sglang/srt/weight_cache/daemon.py b/python/sglang/srt/weight_cache/daemon.py index 864e7b104e7e..b00781e39b3f 100644 --- a/python/sglang/srt/weight_cache/daemon.py +++ b/python/sglang/srt/weight_cache/daemon.py @@ -291,17 +291,17 @@ def load(self): from sglang.srt.layers.moe import initialize_moe_config initialize_moe_config() + # Scheduler normally initializes these before creating ModelRunner. + # The daemon has no Scheduler, but must select identical weight + # transforms (notably V4.1 block-FP8 served by MXFP8 dense GEMMs). + _initialize_gemm_configs() # Initialize distributed backend for model loading # (must be done after server_args and model_config are available) # Build model config first, then init distributed - model_config = ModelConfig( - model_path=self.model_path, - trust_remote_code=self.trust_remote_code, - revision=self.revision, - dtype=self.dtype, - quantization=self.quantization, - ) + # Resolve the same model overrides as the engine (e.g. V4.1's vision + # layer count). A hand-maintained subset silently builds another model. + model_config = ModelConfig.from_server_args(server_args) # Build cache config fingerprint before loading the model. # Loading may mutate hf_config.quantization_config (e.g. via @@ -318,6 +318,9 @@ def load(self): # Refuse unsupported quant methods before creating distributed groups # or touching model weights. + from .model_state import layout_fingerprint, model_state_contract + + contract = model_state_contract(model_config, self.model_loader_extra_config) check_ipc_quant_support(quant_method, quant_config, where="daemon") # The initialized groups are the authority for rank identity. This @@ -351,6 +354,10 @@ def load(self): quant_config_hash=hash_quant_config(quant_config), dtype=str(model_config.dtype), revision=self.revision or "", + model_state_contract=contract, + weight_layout_hash=( + layout_fingerprint(model_config, server_args) if contract else "" + ), **compute_env_stamp(), ) @@ -447,13 +454,32 @@ def _export_state(self): # Also export non-persistent buffers (not in state_dict but needed # for inference, e.g. rotary embedding cos_sin_cache) non_persistent_count = 0 - for name, buf in self.model.named_buffers(): + # Shared non-persistent buffers need every module path too: replacing + # one meta buffer does not update its aliases in other modules. + for name, buf in self.model.named_buffers(remove_duplicate=False): if name not in state_dict_names: state_tensors[name] = (buf.data, False) non_persistent_count += 1 + manifest = None + if self.config.model_state_contract: + from .model_state import export_model_state + + manifest = export_model_state(self.model, state_tensors) self.transport_backend = choose_daemon_transport_backend(state_tensors) + if manifest is not None: + from .model_state import host_table_exports + + host_entries, host_fds = host_table_exports(self.model, state_tensors) + if host_entries: + self.transport_backend.configure_host_table_exports( + host_entries, host_fds + ) self.state_entries = self.transport_backend.prepare_export(state_tensors) + if manifest is not None: + from .model_state import STATE_ENTRY + + self.state_entries[STATE_ENTRY] = {"model_state": manifest} # Log approximate serialized metadata size (not payload-backed bytes). # Only the handle blob carries real weight, so measure it directly: @@ -886,10 +912,25 @@ def launch_weight_cache_daemons( ) -if __name__ == "__main__": - worker_parser = argparse.ArgumentParser(add_help=False) +def _initialize_gemm_configs(): + from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config + from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config + from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config + + initialize_fp8_gemm_config() + initialize_fp4_gemm_config() + initialize_bf16_gemm_config() + + +def _parse_worker_cli_args(argv=None): + # --tp is a server alias, not an abbreviation of the worker's --tp-rank. + worker_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False) WeightCacheDaemonArgs.add_cli_args(worker_parser) - worker_ns, server_argv = worker_parser.parse_known_args() + return worker_parser.parse_known_args(argv) + + +if __name__ == "__main__": + worker_ns, server_argv = _parse_worker_cli_args() from sglang.srt.server_args import prepare_server_args diff --git a/python/sglang/srt/weight_cache/host_vmm.py b/python/sglang/srt/weight_cache/host_vmm.py new file mode 100644 index 000000000000..93ff1c30cabe --- /dev/null +++ b/python/sglang/srt/weight_cache/host_vmm.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Process-local mappings of exportable CUDA HOST_NUMA allocations. + +The daemon allocates once. Clients import the CUDA allocation FD and establish +CPU/GPU access without cudaHostRegister. Tensor users must retire before close. +""" + +from __future__ import annotations + +import ctypes +import logging +import os +import time + +import torch + +logger = logging.getLogger(__name__) + + +def _checked(result): + if int(result[0]) != 0: + raise RuntimeError(f"CUDA host VMM operation failed: {result}") + return result[1] if len(result) == 2 else result[1:] + + +def _cuda_driver(): + try: + from cuda.bindings import driver + except ImportError as error: + raise RuntimeError("CUDA host VMM requires cuda.bindings.driver") from error + # Establish Torch's current primary context; do not retain a second context. + torch.empty(0, device="cuda") + return driver, torch.cuda.current_device() + + +class _CudaArrayView: + def __init__(self, address, nbytes): + self.__cuda_array_interface__ = { + "shape": (nbytes,), + "typestr": "|u1", + "data": (int(address), False), + "version": 3, + } + + +class HostVmmAllocation: + def __init__(self, driver, device): + self.driver, self.device = driver, device + self.fd = None + self.handle = None + self.address = None + self.mapped = False + self.bytes = None + self.closed = False + self.gpu_view = False + self.mapping_timings = {} + + @staticmethod + def _props(driver, numa_id): + prop = driver.CUmemAllocationProp() + prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.requestedHandleTypes = ( + driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ) + prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA + prop.location.id = numa_id + return prop + + @staticmethod + def _require_support(driver, device): + for name in ( + "CU_DEVICE_ATTRIBUTE_HOST_NUMA_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED", + "CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR_SUPPORTED", + ): + attribute = getattr(driver.CUdevice_attribute, name, None) + if ( + attribute is None + or _checked(driver.cuDeviceGetAttribute(attribute, device)) != 1 + ): + raise RuntimeError(f"CUDA host VMM requires {name} on device {device}") + + @classmethod + def create(cls, nbytes): + if type(nbytes) is not int or nbytes <= 0: + raise ValueError("Invalid host VMM logical size") + driver, device = _cuda_driver() + result = cls(driver, device) + try: + cls._require_support(driver, device) + numa_id = _checked( + driver.cuDeviceGetAttribute( + driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID, device + ) + ) + result.numa_id = max(0, numa_id) + result.nbytes = nbytes + prop = cls._props(driver, result.numa_id) + granularity = _checked( + driver.cuMemGetAllocationGranularity( + prop, + driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ) + ) + result.allocation_nbytes = ( + (nbytes + granularity - 1) // granularity * granularity + ) + tic = time.perf_counter() + result.handle = _checked( + driver.cuMemCreate(result.allocation_nbytes, prop, 0) + ) + result.fd = int( + _checked( + driver.cuMemExportToShareableHandle( + result.handle, prop.requestedHandleTypes, 0 + ) + ) + ) + os.set_inheritable(result.fd, False) + result._map() + logger.info( + "Engram VMM create: bytes=%d, allocation_bytes=%d, numa=%d, elapsed=%.3fs", + nbytes, + result.allocation_nbytes, + result.numa_id, + time.perf_counter() - tic, + ) + return result + except Exception: + result.close() + raise + + @classmethod + def from_fd(cls, fd, nbytes, allocation_nbytes, numa_id, *, gpu_view=False): + """Adopt fd, including on failure; verify the imported CUDA allocation.""" + result = None + try: + if ( + type(nbytes) is not int + or nbytes <= 0 + or type(allocation_nbytes) is not int + or allocation_nbytes < nbytes + or type(numa_id) is not int + or numa_id < 0 + ): + raise ValueError("Invalid host VMM allocation descriptor") + driver, device = _cuda_driver() + result = cls(driver, device) + result.fd = fd + result.gpu_view = gpu_view + result.nbytes, result.allocation_nbytes, result.numa_id = ( + nbytes, + allocation_nbytes, + numa_id, + ) + cls._require_support(driver, device) + tic = time.perf_counter() + result.handle = _checked( + driver.cuMemImportFromShareableHandle( + fd, + driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, + ) + ) + prop = _checked( + driver.cuMemGetAllocationPropertiesFromHandle(result.handle) + ) + if ( + prop.type != driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + or prop.location.type + != driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA + or prop.location.id != numa_id + ): + raise ValueError( + "Imported allocation is not the declared HOST_NUMA backing" + ) + granularity = _checked( + driver.cuMemGetAllocationGranularity( + prop, + driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_MINIMUM, + ) + ) + if ( + allocation_nbytes + != (nbytes + granularity - 1) // granularity * granularity + ): + raise ValueError("Host VMM allocation size/granularity mismatch") + result._map() + logger.info( + "Engram VMM import: bytes=%d, allocation_bytes=%d, numa=%d, elapsed=%.3fs", + nbytes, + allocation_nbytes, + numa_id, + time.perf_counter() - tic, + ) + return result + except Exception: + if result is not None: + result.close() + else: + os.close(fd) + raise + + def _map(self): + driver = self.driver + tic = time.perf_counter() + self.address = _checked( + driver.cuMemAddressReserve(self.allocation_nbytes, 0, 0, 0) + ) + self.mapping_timings["reserve"] = time.perf_counter() - tic + tic = time.perf_counter() + _checked( + driver.cuMemMap(self.address, self.allocation_nbytes, 0, self.handle, 0) + ) + self.mapping_timings["map"] = time.perf_counter() - tic + self.mapped = True + host = driver.CUmemAccessDesc() + host.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_HOST_NUMA + host.location.id = self.numa_id + host.flags = driver.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + gpu = driver.CUmemAccessDesc() + gpu.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + gpu.location.id = self.device + gpu.flags = host.flags + descriptors = [gpu] if self.gpu_view else [host, gpu] + tic = time.perf_counter() + _checked( + driver.cuMemSetAccess( + self.address, self.allocation_nbytes, descriptors, len(descriptors) + ) + ) + self.mapping_timings["access"] = time.perf_counter() - tic + if self.gpu_view: + # Inference only needs GPU access. A CUDA tensor view represents + # HOST_NUMA backing without eagerly installing CPU PTEs for it. + self.bytes = torch.as_tensor( + _CudaArrayView(self.address, self.nbytes), device=f"cuda:{self.device}" + ) + if self.bytes.data_ptr() != int(self.address): + raise RuntimeError("CUDA host VMM view unexpectedly copied storage") + else: + self.bytes = torch.frombuffer( + (ctypes.c_uint8 * self.nbytes).from_address(int(self.address)), + dtype=torch.uint8, + ) + logger.info( + "Engram VMM mapping: gpu_view=%s, bytes=%d, steps=%s", + self.gpu_view, + self.nbytes, + self.mapping_timings, + ) + + @property + def device_ptr(self): + return int(self.address) + + def close(self): + if self.closed: + return + self.closed = True + self.bytes = None + # Attempt all cleanup steps, even if a context has already shut down. + failures = [] + if self.mapped: + try: + _checked(self.driver.cuMemUnmap(self.address, self.allocation_nbytes)) + except Exception as error: + failures.append(error) + self.mapped = False + if self.address is not None: + try: + _checked( + self.driver.cuMemAddressFree(self.address, self.allocation_nbytes) + ) + except Exception as error: + failures.append(error) + self.address = None + if self.handle is not None: + try: + _checked(self.driver.cuMemRelease(self.handle)) + except Exception as error: + failures.append(error) + self.handle = None + if self.fd is not None: + os.close(self.fd) + self.fd = None + if failures: + logger.warning("CUDA host VMM cleanup failures: %s", failures) + + def __del__(self): + try: + self.close() + except Exception: + pass diff --git a/python/sglang/srt/weight_cache/ipc_loader.py b/python/sglang/srt/weight_cache/ipc_loader.py index 6bf6a9305b4b..bc7547fbc849 100644 --- a/python/sglang/srt/weight_cache/ipc_loader.py +++ b/python/sglang/srt/weight_cache/ipc_loader.py @@ -102,6 +102,9 @@ def load_model( # (client mode) or serving wrong-numerics IPC weights. Checked here so # it applies regardless of whether the daemon is reachable. quant_method, engine_quant_config = self._resolve_engine_quant(model_config) + from .model_state import model_state_contract + + model_state_contract(model_config, self.load_config.model_loader_extra_config) check_ipc_quant_support(quant_method, engine_quant_config, where="client") # Try to fetch state from daemon @@ -293,6 +296,9 @@ def _set_module_tensor(model, name, tensor, is_param=True): # requires_grad=False: the IPC memory is shared/read-only and SGLang # is inference-only, so autograd must never write into it. new_param = nn.Parameter(tensor, requires_grad=False) + old_param = obj._parameters.get(leaf_name) + if old_param is not None: + new_param.__dict__.update(old_param.__dict__) setattr(obj, leaf_name, new_param) else: # register_buffer raises KeyError if the name already exists as a @@ -341,11 +347,35 @@ def _load_zero_copy_mode( # remove_duplicate=False mirrors the daemon's export (which keys tied # weights under every name) so a tied parameter is recognized under all # of its names here too. + from .model_state import ( + STATE_ENTRY, + TENSOR_PREFIX, + model_state_contract, + validate_model_state, + ) + + state_entry = entries.get(STATE_ENTRY) + manifest = state_entry["model_state"] if state_entry else None + contract = model_state_contract( + model_config, self.load_config.model_loader_extra_config + ) + if bool(manifest) != bool(contract): + raise ValueError("IPC cache is missing the required model post-load state") + if manifest is not None: + validate_model_state(model, manifest) + expected = set(manifest["schema"]) | set(manifest["extras"]) | {STATE_ENTRY} + if set(entries) != expected: + raise ValueError( + "V4.1 IPC exported tensor set does not match its manifest" + ) + existing_params = { name: param for name, param in model.named_parameters(remove_duplicate=False) } - existing_buffers = {name: buf for name, buf in model.named_buffers()} + existing_buffers = { + name: buf for name, buf in model.named_buffers(remove_duplicate=False) + } existing_names = set(existing_params) | set(existing_buffers) imported_refs = [] @@ -357,9 +387,34 @@ def _load_zero_copy_mode( # Iterate over ALL daemon entries (not just model params/buffers). # This ensures post-quantization parameters (weight_scale, etc.) # that were created by process_weights_after_loading are also mapped. + extra_tensors = {} for name, entry in entries.items(): + if name == STATE_ENTRY: + continue imported_tensor = self._transport_backend.import_tensor(entry) + if name.startswith(TENSOR_PREFIX): + if manifest is None: + raise ValueError("Unexpected model-state tensor in IPC cache") + shape, dtype = manifest["extras"][name] + if ( + list(imported_tensor.shape) != shape + or str(imported_tensor.dtype) != dtype + ): + raise ValueError(f"V4.1 IPC derived tensor schema mismatch: {name}") + extra_tensors[name] = imported_tensor + imported_refs.append(imported_tensor) + continue is_param = entry.get("is_param", True) + if manifest is not None: + shape, dtype, kind = manifest["schema"][name] + if ( + list(imported_tensor.shape) != shape + or str(imported_tensor.dtype) != dtype + or is_param != kind + ): + raise ValueError( + f"V4.1 IPC processed tensor schema mismatch: {name}" + ) if name in existing_names: # Existing parameter/buffer — validate shape/dtype @@ -367,7 +422,7 @@ def _load_zero_copy_mode( ref_param = existing_params[name] else: ref_param = existing_buffers[name] - if ( + if manifest is None and ( imported_tensor.shape != ref_param.shape or imported_tensor.dtype != ref_param.dtype ): @@ -397,6 +452,9 @@ def _load_zero_copy_mode( f"not merely uninitialized weights:\n" + "\n".join(mismatched) ) + if manifest is not None and any("host_table" in e for e in entries.values()): + self._transport_backend.bind_host_tables(model, entries) + # After mapping every daemon entry, any tensor still on the meta device # is one the daemon did NOT provide. Filling it with torch.empty() would # hand the model uninitialized GPU memory — silently producing wrong @@ -432,6 +490,11 @@ def _load_zero_copy_mode( f"{'...' if len(still_on_meta_buffers) > 10 else ''}" ) + if manifest is not None: + from .model_state import restore_model_state + + restore_model_state(model, manifest, extra_tensors) + map_elapsed = time.perf_counter() - map_tic # Stash IPC refs on the model to prevent GC (which would unmap the memory) @@ -517,6 +580,13 @@ def _fetch_from_cache(self, model_config, device_config) -> Optional[dict]: quant_method, quant_config = self._resolve_engine_quant(model_config) + from sglang.srt.runtime_context import get_server_args + + from .model_state import layout_fingerprint, model_state_contract + + contract = model_state_contract( + model_config, self.load_config.model_loader_extra_config + ) engine_config = CacheConfig( model_path=model_config.model_path, model_arch=( @@ -542,6 +612,12 @@ def _fetch_from_cache(self, model_config, device_config) -> Optional[dict]: quant_config_hash=hash_quant_config(quant_config), dtype=str(model_config.dtype), revision=model_config.revision or "", + model_state_contract=contract, + weight_layout_hash=( + layout_fingerprint(model_config, get_server_args()) + if contract + else "" + ), **compute_env_stamp(), ) diff --git a/python/sglang/srt/weight_cache/model_state.py b/python/sglang/srt/weight_cache/model_state.py new file mode 100644 index 000000000000..e67eb7219bdb --- /dev/null +++ b/python/sglang/srt/weight_cache/model_state.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental post-load state transport for DeepSeek V4.1. + +Tensor payloads use the existing IPC transport. The manifest contains only +references and primitive metadata; no tensors, callables, or CUDA resources. +Host Engram uses daemon-owned memfds or CUDA HOST_NUMA VMM allocations. Mutable request state is not cached. +""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os + +import torch + +STATE_ENTRY = "__sglang_model_state__" +TENSOR_PREFIX = "__sglang_model_tensor__/" +CONTRACT = "deepseek_v41_post_load_v3" +# Post-load flags consumed by the model, rather than purely constructor flags. +MODULE_FLAGS = ( + "ape_converted", + "_mxfp4_backend", + "_mega_moe_weights_built", + "block_fp8_mxfp8_ready", + "intermediate_size_per_partition", + "hidden_pad", + "intermediate_pad", + "pad_id", + "image_token_id", + "k_scale_float", + "v_scale_float", + "use_padded_loading", +) +PARAM_FLAGS = ("format_ue8m0",) + + +def is_dsv41(model_config) -> bool: + cfg = getattr(model_config, "hf_text_config", model_config.hf_config) + return getattr(cfg, "model_type", None) in ("deepseek_v41", "deepseek_v4.1") + + +def model_state_contract(model_config, extra_config) -> str: + if not is_dsv41(model_config): + return "" + if isinstance(extra_config, str): + extra_config = json.loads(extra_config) + if not isinstance(extra_config or {}, dict): + raise ValueError("model-loader-extra-config must be a JSON object") + if (extra_config or {}).get("dsv41_weight_cache") is not True: + raise ValueError( + "DeepSeek V4.1 IPC recovery is experimental; explicitly set " + "model-loader-extra-config must contain dsv41_weight_cache=true." + ) + return CONTRACT + + +def host_table_backend(extra_config): + if isinstance(extra_config, str): + extra_config = json.loads(extra_config) + backend = (extra_config or {}).get("dsv41_host_table_backend", "memfd") + if backend not in ("memfd", "cuda_vmm"): + raise ValueError("dsv41_host_table_backend must be memfd or cuda_vmm") + return backend + + +def layout_fingerprint(model_config, server_args) -> str: + """Stamp choices that change packed weights or derived model state.""" + fields = ( + "moe_runner_backend", + "fp8_gemm_runner_backend", + "fp4_gemm_runner_backend", + "bf16_gemm_backend", + "quantization", + "enable_shared_experts_fusion", + "enable_deterministic_inference", + "flashinfer_mxfp4_moe_precision", + "enable_symm_mem", + ) + data = { + "contract": CONTRACT, + "host_backend": host_table_backend( + getattr(server_args, "model_loader_extra_config", None) + ), + "config": model_config.hf_config.to_dict(), + "flags": {k: getattr(server_args, k, None) for k in fields}, + "env": { + k: v + for k, v in sorted(os.environ.items()) + if k.startswith( + ( + "SGLANG_OPT_", + "SGLANG_DSV4", + "SGLANG_ENABLE_DSV41_ENGRAM", + "SGLANG_DEEPGEMM_", + "SGLANG_MXFP", + "SGLANG_USE_AITER", + ) + ) + }, + } + libraries = {} + for package in ( + "sglang-kernel", + "flashinfer-python", + "deep-gemm", + "triton", + "tvm-ffi", + "tilelang", + ): + try: + libraries[package] = importlib.metadata.version(package) + except importlib.metadata.PackageNotFoundError: + libraries[package] = "" + data["libraries"] = libraries + # Two checkouts may have the same library version but different weight + # transforms. Include source content, not paths, in the identity. + from pathlib import Path + + root = Path(__file__).resolve().parents[1] + paths = ( + "weight_cache/model_state.py", + "weight_cache/daemon.py", + "weight_cache/ipc_loader.py", + "model_executor/model_runner_components/load_model_utils.py", + "models/deepseek_v4.py", + "layers/engram.py", + "weight_cache/transport.py", + "weight_cache/host_vmm.py", + "layers/attention/dsv4/compressor.py", + "layers/quantization/fp8.py", + "layers/quantization/fp8_utils.py", + "layers/quantization/fp4_utils.py", + "layers/quantization/unquant.py", + "layers/quantization/mxfp4_flashinfer_trtllm_moe.py", + "layers/moe/mega_moe.py", + ) + data["source"] = { + p: hashlib.sha256((root / p).read_bytes()).hexdigest() for p in paths + } + return hashlib.sha256(json.dumps(data, sort_keys=True).encode()).hexdigest() + + +def _primitive(value): + if value is None or type(value) in (bool, int, float, str): + return True + return isinstance(value, (list, tuple)) and all(_primitive(v) for v in value) + + +def _tensor_tree(value): + if isinstance(value, torch.Tensor): + return True + return ( + isinstance(value, (tuple, list)) + and bool(value) + and all(_tensor_tree(v) or v is None for v in value) + ) + + +def _class_name(obj): + cls = type(obj) + return f"{cls.__module__}.{cls.__qualname__}" + + +def export_model_state(model, state_tensors): + """Capture derived tensor references and metadata after weight transforms. + + This runs before any request/KV initialization. Module registries and + quantization runner objects remain constructed in the client; they are + never pickled. Derived tensor trees keep their tuple/list structure. + """ + manifest = { + "contract": CONTRACT, + "schema": {}, + "objects": {}, + "attrs": [], + "params": {}, + "extras": {}, + } + for name, (tensor, is_param) in state_tensors.items(): + manifest["schema"][name] = [list(tensor.shape), str(tensor.dtype), is_param] + tensor_ids = {} + + def encode(value): + if isinstance(value, torch.Tensor): + key = tensor_ids.get(id(value)) + if key is None: + key = f"{TENSOR_PREFIX}{len(tensor_ids)}" + tensor_ids[id(value)] = key + state_tensors[key] = (value.detach(), False) + manifest["extras"][key] = [list(value.shape), str(value.dtype)] + return {"tensor": key} + return { + "tuple" if isinstance(value, tuple) else "list": [ + None if v is None else encode(v) for v in value + ] + } + + def capture(path, obj, quant=False): + manifest["objects"][path] = _class_name(obj) + for attr, value in vars(obj).items(): + if _tensor_tree(value): + manifest["attrs"].append([path, attr, encode(value)]) + elif attr in MODULE_FLAGS or (quant and _primitive(value)): + if not _primitive(value): + raise ValueError(f"Non-primitive post-load flag: {path}.{attr}") + manifest["attrs"].append([path, attr, {"value": value}]) + + for path, module in model.named_modules(): + capture(path, module) + method = getattr(module, "quant_method", None) + if method is not None: + qpath = f"{path}.quant_method" if path else "quant_method" + capture(qpath, method, quant=True) + wrapped = getattr(method, "_fp8", None) + if wrapped is not None: + capture(qpath + "._fp8", wrapped, quant=True) + for name, param in model.named_parameters(remove_duplicate=False): + flags = {k: getattr(param, k) for k in PARAM_FLAGS if hasattr(param, k)} + if flags: + manifest["params"][name] = flags + return manifest + + +def host_table_exports(model, state_tensors): + """Describe registered Engram views without serializing CPU tensor payloads.""" + entries, fds = {}, [] + for path, module in model.named_modules(): + table = getattr(module, "host_table", None) + if table is None: + continue + if table.fd is None: + raise ValueError( + "IPC host Engram requires daemon-owned shareable host storage" + ) + descriptor = { + "index": len(fds), + "module": path, + "layout": table.layout, + "nbytes": table.nbytes, + } + if getattr(table, "backend", "memfd") == "cuda_vmm": + descriptor.update( + backend="cuda_vmm", + allocation_nbytes=table._vmm.allocation_nbytes, + numa_id=table._vmm.numa_id, + ) + fds.append(table.fd) + for attr in ("weight", "scale"): + name = f"{path}.{attr}" + tensor, is_param = state_tensors[name] + if tensor.device.type != "cpu" or not tensor.is_contiguous(): + raise ValueError(f"Invalid host Engram tensor: {name}") + entries[name] = { + "host_table": descriptor, + "host_offset": tensor.data_ptr() - table.bytes.data_ptr(), + "shape": list(tensor.shape), + "dtype": str(tensor.dtype).replace("torch.", ""), + "is_param": is_param, + } + return entries, fds + + +def _resolve(model, path): + obj = model + for part in path.split(".") if path else (): + obj = getattr(obj, part) + return obj + + +def validate_model_state(model, manifest): + if manifest.get("contract") != CONTRACT: + raise ValueError("Missing or incompatible V4.1 post-load state contract") + for path, cls in manifest["objects"].items(): + if _class_name(_resolve(model, path)) != cls: + raise ValueError(f"V4.1 IPC object class mismatch at {path}") + + +def restore_model_state(model, manifest, tensors): + validate_model_state(model, manifest) + + def decode(value): + if "tensor" in value: + return tensors[value["tensor"]] + if "value" in value: + return value["value"] + if "tuple" in value: + return tuple(None if v is None else decode(v) for v in value["tuple"]) + if "list" in value: + return [None if v is None else decode(v) for v in value["list"]] + raise ValueError("Invalid V4.1 tensor descriptor") + + for path, attr, value in manifest["attrs"]: + setattr(_resolve(model, path), attr, decode(value)) + for name, flags in manifest["params"].items(): + for flag, value in flags.items(): + if flag not in PARAM_FLAGS: + raise ValueError(f"Unknown V4.1 parameter flag: {flag}") + setattr(_resolve(model, name), flag, value) + # Bind only process-local references. Never call post_load_weights(), which + # would transform shared scales/APE a second time and corrupt all clients. + layers = getattr(getattr(model, "model", None), "layers", ()) + for i, layer in enumerate(layers): + mlp = getattr(layer, "mlp", None) + topk_config = getattr(getattr(mlp, "topk", None), "topk_config", None) + gate = getattr(mlp, "gate", None) + if topk_config is not None and gate is not None: + # TopKConfig is a dataclass, not a registered module. Its cached + # parameter reference still points to the constructor's meta bias. + topk_config.correction_bias = gate.e_score_correction_bias + if hasattr(layer, "_init_hyper_connections"): + layer._init_hyper_connections() + if hasattr(layer, "bind_next"): + layer.bind_next(layers[i + 1] if i + 1 < len(layers) else None) diff --git a/python/sglang/srt/weight_cache/protocol.py b/python/sglang/srt/weight_cache/protocol.py index 46cfe588d82a..977374933197 100644 --- a/python/sglang/srt/weight_cache/protocol.py +++ b/python/sglang/srt/weight_cache/protocol.py @@ -55,6 +55,9 @@ class CacheConfig(msgspec.Struct): device_capability: str # local compute capability, e.g. "8.0" ("" if N/A) torch_version: str # torch.__version__ of the process that built the weights + model_state_contract: str = "" + weight_layout_hash: str = "" + def matches(self, other: "CacheConfig") -> bool: """Check if two configs are compatible for weight sharing.""" return self == other diff --git a/python/sglang/srt/weight_cache/standby.py b/python/sglang/srt/weight_cache/standby.py new file mode 100644 index 000000000000..b2af9713bc8d --- /dev/null +++ b/python/sglang/srt/weight_cache/standby.py @@ -0,0 +1,434 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental warm-standby gateway for resident-weight clients. + +Run with ``python -m sglang.srt.weight_cache.standby --config config.json``. +The weight daemon is managed separately. Commands must launch independent client +process groups with distinct HTTP/NCCL ports and bounded private KV pools. +Only stateless inference is routed. Failed requests are never replayed. +""" + +import argparse +import asyncio +import contextlib +import json +import logging +import os +import signal +import time +from dataclasses import dataclass, field +from pathlib import Path +from urllib.parse import urlsplit + +import aiohttp +from aiohttp import web + +logger = logging.getLogger(__name__) +HOP_HEADERS = { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "host", + "content-length", +} +ROUTES = { + "/generate", + "/v1/chat/completions", + "/v1/completions", + "/v1/models", + "/get_model_info", +} + + +def forward_headers(headers): + excluded = HOP_HEADERS | { + value.strip().lower() for value in headers.get("Connection", "").split(",") + } + return {key: value for key, value in headers.items() if key.lower() not in excluded} + + +@dataclass +class Worker: + name: str + url: str + command: list + env: dict = field(default_factory=dict) + state: str = "starting" + process: object = None + generation: int = 0 + ready_at: float = 0.0 + last_error: str = "" + next_probe: float = 0.0 + next_start: float = 0.0 + log_file: object = None + + +class StandbyGateway: + def __init__(self, config): + specs = config["workers"] + if len(specs) != 2: + raise ValueError("Exactly two workers are required") + if len({s["name"] for s in specs}) != 2 or len({s["url"] for s in specs}) != 2: + raise ValueError("Worker names and URLs must be distinct") + for spec in specs: + parsed = urlsplit(spec["url"]) + if ( + parsed.scheme != "http" + or parsed.hostname not in {"127.0.0.1", "localhost", "::1"} + or parsed.path not in {"", "/"} + ): + raise ValueError("Workers must use loopback HTTP URLs") + if not spec.get("command") or not all( + isinstance(x, str) for x in spec["command"] + ): + raise ValueError("Each worker needs a non-empty command argument list") + self.workers = [ + Worker(s["name"], s["url"].rstrip("/"), s["command"], s.get("env", {})) + for s in specs + ] + self.config = config + self.active = None + self.epoch = 0 + self.events = [] + self.session = None + self.tasks = [] + self.log_dir = Path(config.get("log_dir", "standby-logs")) + self.warmup = config.get( + "warmup", + { + "text": "The capital of France is", + "sampling_params": {"temperature": 0, "max_new_tokens": 4}, + }, + ) + if self.warmup.get("stream") or "session_params" in self.warmup: + raise ValueError("Warmup must be non-streaming and stateless") + for key in ( + "poll_interval", + "health_interval", + "health_timeout", + "restart_backoff", + "startup_timeout", + "warmup_timeout", + ): + if float(config.get(key, 1)) <= 0: + raise ValueError(f"{key} must be positive") + + def event(self, kind, **fields): + record = { + "event": kind, + "monotonic": time.monotonic(), + "time": time.time(), + **fields, + } + self.events.append(record) + self.events = self.events[-100:] + logger.info(json.dumps(record)) + + def choose(self): + for worker in self.workers: + if ( + worker.state == "ready" + and worker.process is not None + and worker.process.returncode is not None + ): + worker.state = "failed" + worker.last_error = f"launcher exited {worker.process.returncode}" + self.event( + "failed", + worker=worker.name, + generation=worker.generation, + reason=worker.last_error, + ) + worker.next_start = time.monotonic() + self.config.get( + "restart_backoff", 2 + ) + if self.active is not None and self.active.state == "ready": + return self.active + old = self.active.name if self.active else None + self.active = next((w for w in self.workers if w.state == "ready"), None) + if self.active is not None: + self.epoch += 1 + self.event( + "promoted", previous=old, worker=self.active.name, epoch=self.epoch + ) + return self.active + + def fail(self, worker, error): + if worker.state == "failed": + return + worker.state = "failed" + worker.last_error = str(error) + self.event( + "failed", + worker=worker.name, + generation=worker.generation, + reason=str(error), + ) + self.choose() + + async def probe(self, worker): + timeout = aiohttp.ClientTimeout(total=self.config.get("health_timeout", 2.0)) + try: + async with self.session.get( + worker.url + self.config.get("health_path", "/health"), timeout=timeout + ) as response: + await response.read() + return response.status == 200 + except (aiohttp.ClientError, asyncio.TimeoutError): + return False + + async def stop_worker(self, worker): + process = worker.process + if process is not None: + # Kill descendants even if the launcher has already exited. + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) + try: + await asyncio.wait_for(process.wait(), timeout=5) + except asyncio.TimeoutError: + pass + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGKILL) + await process.wait() + worker.process = None + if worker.log_file is not None: + worker.log_file.close() + worker.log_file = None + + async def prepare(self, worker): + worker.state = "starting" + worker.generation += 1 + worker.log_file = open( + self.log_dir / f"{worker.name}.{worker.generation}.log", "ab", buffering=0 + ) + try: + worker.process = await asyncio.create_subprocess_exec( + *worker.command, + env={**os.environ, **worker.env}, + stdout=worker.log_file, + stderr=asyncio.subprocess.STDOUT, + start_new_session=True, + ) + self.event( + "started", + worker=worker.name, + generation=worker.generation, + pid=worker.process.pid, + ) + deadline = time.monotonic() + self.config.get("startup_timeout", 600) + while time.monotonic() < deadline: + if worker.process.returncode is not None: + raise RuntimeError(f"launcher exited {worker.process.returncode}") + if await self.probe(worker): + async with self.session.post( + worker.url + "/generate", + json=self.warmup, + timeout=aiohttp.ClientTimeout( + total=self.config.get("warmup_timeout", 180) + ), + ) as response: + body = await response.read() + if response.status != 200: + raise RuntimeError( + f"warmup HTTP {response.status}: {body[:200]!r}" + ) + result = json.loads(body) + if ( + not isinstance(result, dict) + or "text" not in result + or "error" in result + ): + raise RuntimeError( + "warmup did not return a generation response" + ) + if worker.process.returncode is not None: + raise RuntimeError("launcher exited during warmup") + worker.ready_at = time.monotonic() + worker.next_probe = worker.ready_at + self.config.get( + "health_interval", 1 + ) + worker.state = "ready" + self.event( + "ready", worker=worker.name, generation=worker.generation + ) + self.choose() + return + await asyncio.sleep(0.5) + raise TimeoutError("worker startup timeout") + except ( + OSError, + RuntimeError, + ValueError, + aiohttp.ClientError, + asyncio.TimeoutError, + ) as error: + self.fail(worker, error) + await self.stop_worker(worker) + worker.next_start = time.monotonic() + self.config.get("restart_backoff", 2) + + async def maintain(self, worker, predecessor=None): + # Sequential initial preparation avoids simultaneous model initialization. + if predecessor is not None: + while predecessor.state != "ready": + await asyncio.sleep(0.1) + try: + while True: + if worker.state in {"starting", "failed"}: + await self.stop_worker(worker) + await asyncio.sleep(max(0, worker.next_start - time.monotonic())) + await self.prepare(worker) + elif worker.process.returncode is not None: + self.fail(worker, f"launcher exited {worker.process.returncode}") + worker.next_start = time.monotonic() + self.config.get( + "restart_backoff", 2 + ) + elif time.monotonic() >= worker.next_probe: + generation = worker.generation + healthy = await self.probe(worker) + if ( + generation == worker.generation + and worker.state == "ready" + and not healthy + ): + self.fail(worker, "health probe failed") + worker.next_start = time.monotonic() + self.config.get( + "restart_backoff", 2 + ) + worker.next_probe = time.monotonic() + self.config.get( + "health_interval", 1 + ) + await asyncio.sleep(self.config.get("poll_interval", 0.05)) + finally: + await self.stop_worker(worker) + + async def start(self, app): + self.log_dir.mkdir(parents=True, exist_ok=True) + self.session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout(total=None, sock_connect=1), + auto_decompress=False, + ) + self.tasks = [ + asyncio.create_task(self.maintain(self.workers[0])), + asyncio.create_task(self.maintain(self.workers[1], self.workers[0])), + ] + + async def close(self, app): + for task in self.tasks: + task.cancel() + await asyncio.gather(*self.tasks, return_exceptions=True) + for worker in self.workers: + await self.stop_worker(worker) + if self.session is not None: + await self.session.close() + + async def status(self, request): + self.choose() + return web.json_response( + { + "active": self.active.name if self.active else None, + "epoch": self.epoch, + "standby_ready": sum(w.state == "ready" for w in self.workers) == 2, + "workers": [ + { + "name": w.name, + "state": w.state, + "generation": w.generation, + "pid": w.process.pid if w.process else None, + "ready_at": w.ready_at, + "last_error": w.last_error, + } + for w in self.workers + ], + "events": self.events, + } + ) + + async def health(self, request): + return web.Response(status=200 if self.choose() else 503) + + async def forward(self, request): + if request.path not in ROUTES: + raise web.HTTPNotFound() + body = await request.read() + if request.method == "POST": + try: + payload = json.loads(body) + except (ValueError, UnicodeDecodeError): + raise web.HTTPBadRequest(text="Invalid JSON") + if not isinstance(payload, dict): + raise web.HTTPBadRequest(text="Expected a JSON object") + if "session_params" in payload: + raise web.HTTPConflict( + text="Session state is not recoverable across workers" + ) + worker = self.choose() + if worker is None: + raise web.HTTPServiceUnavailable(text="No warmed worker available") + generation = worker.generation + epoch = self.epoch + outgoing = None + try: + async with self.session.request( + request.method, + worker.url + request.path_qs, + data=body, + headers=forward_headers(request.headers), + allow_redirects=False, + ) as response: + headers = forward_headers(response.headers) + headers.update( + { + "X-SGLang-Recovery-Worker": worker.name, + "X-SGLang-Recovery-Epoch": str(epoch), + } + ) + outgoing = web.StreamResponse(status=response.status, headers=headers) + await outgoing.prepare(request) + async for chunk in response.content.iter_any(): + await outgoing.write(chunk) + await outgoing.write_eof() + return outgoing + except (aiohttp.ClientError, asyncio.TimeoutError) as error: + if worker.generation == generation: + self.fail(worker, error) + # No retry: the failed worker may already have accepted the request. + if outgoing is not None and outgoing.prepared: + if request.transport is not None: + request.transport.close() + return outgoing + raise web.HTTPServiceUnavailable( + text="Worker failed; request was not replayed" + ) + + def app(self): + app = web.Application( + client_max_size=self.config.get("client_max_size", 16 * 1024 * 1024) + ) + app.on_startup.append(self.start) + app.on_cleanup.append(self.close) + app.router.add_get("/recovery/status", self.status) + app.router.add_get("/health", self.health) + app.router.add_route("*", "/{path:.*}", self.forward) + return app + + +def main(): + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("--config", required=True) + args = parser.parse_args() + config = json.loads(Path(args.config).read_text()) + logging.basicConfig(level=logging.INFO, format="%(message)s") + gateway = StandbyGateway(config) + web.run_app( + gateway.app(), + host=config.get("host", "127.0.0.1"), + port=config.get("port", 22940), + ) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/weight_cache/transport.py b/python/sglang/srt/weight_cache/transport.py index f770fbb335fa..374f9e641e53 100644 --- a/python/sglang/srt/weight_cache/transport.py +++ b/python/sglang/srt/weight_cache/transport.py @@ -5,6 +5,7 @@ import array import logging +import math import os import socket import struct @@ -31,29 +32,38 @@ def _send_fd(sock: socket.socket, fd: int, index: int) -> None: sent = sock.sendmsg( [payload], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds.tobytes())] ) - if sent != len(payload): - raise RuntimeError(f"sendmsg sent {sent} bytes, expected {len(payload)}") + if sent <= 0: + raise RuntimeError("sendmsg failed to send the FD header") + if sent < len(payload): + sock.sendall(payload[sent:]) def _recv_fd(sock: socket.socket) -> Tuple[int, int]: fd_item_size = array.array("i").itemsize - data, ancdata, _, _ = sock.recvmsg( - _FD_INDEX_STRUCT.size, socket.CMSG_SPACE(fd_item_size) + data, ancdata, flags, _ = sock.recvmsg( + _FD_INDEX_STRUCT.size, + socket.CMSG_SPACE(fd_item_size), + getattr(socket, "MSG_CMSG_CLOEXEC", 0), ) - if len(data) != _FD_INDEX_STRUCT.size: - raise RuntimeError( - f"received truncated fd header: {len(data)} < {_FD_INDEX_STRUCT.size}" - ) - index = _FD_INDEX_STRUCT.unpack(data)[0] fds = array.array("i") for level, cmsg_type, cmsg_data in ancdata: if level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS: fds.frombytes(cmsg_data[: len(cmsg_data) - (len(cmsg_data) % fd_item_size)]) - if len(fds) != 1: + try: + if flags & socket.MSG_CTRUNC or len(fds) != 1: + raise RuntimeError(f"expected one complete fd, got {len(fds)}") + while len(data) < _FD_INDEX_STRUCT.size: + chunk = sock.recv(_FD_INDEX_STRUCT.size - len(data)) + if not chunk: + raise RuntimeError("received truncated fd header") + data += chunk + index = _FD_INDEX_STRUCT.unpack(data)[0] + os.set_inheritable(fds[0], False) + return int(index), int(fds[0]) + except Exception: for fd in fds: os.close(fd) - raise RuntimeError(f"expected one fd, got {len(fds)}") - return int(index), int(fds[0]) + raise class WeightCacheTransportBackend(ABC): @@ -91,11 +101,25 @@ def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor: class TorchIpcTransportBackend(WeightCacheTransportBackend): name = TORCH_IPC_BACKEND + def __init__(self): + self._host_entries = {} + self._host_export_fds = [] + self._host_received_fds = {} + self._host_descriptors = {} + self._host_tables = {} + + def configure_host_table_exports(self, entries, fds): + self._host_entries = entries + self._host_export_fds = list(fds) # borrowed; daemon model owns them + def prepare_export( self, state_tensors: Mapping[str, Tuple[torch.Tensor, bool]] ) -> Dict[str, Dict[str, Any]]: entries: Dict[str, Dict[str, Any]] = {} for name, (tensor, is_param) in state_tensors.items(): + if name in self._host_entries: + entries[name] = self._host_entries[name] + continue entries[name] = { "handle": MultiprocessingSerializer.serialize( tensor.data, output_str=True @@ -124,17 +148,137 @@ def send_fetch_state_response( "pid": pid, "transport_backend": self.name, "preloaded_weights_bytes": preloaded_weights_bytes, + "host_table_count": len(self._host_export_fds), }, ) + for index, fd in enumerate(self._host_export_fds): + _send_fd(conn, fd, index) def recv_fetch_state_response( self, sock: socket.socket, result: Dict[str, Any] ) -> Dict[str, Any]: + count = result.get("host_table_count", 0) + if type(count) is not int or count < 0: + raise ValueError("Invalid host table FD count") + descriptors = {} + for name, entry in result["entries"].items(): + desc = entry.get("host_table") + if desc is None: + continue + index = desc["index"] + if ( + type(index) is not int + or not 0 <= index < count + or desc["layout"] not in ("shared", "per_rank") + or type(desc["nbytes"]) is not int + or desc["nbytes"] <= 0 + or desc.get("backend", "memfd") not in ("memfd", "cuda_vmm") + or name not in (desc["module"] + ".weight", desc["module"] + ".scale") + ): + raise ValueError("Invalid host table descriptor") + if desc.get("backend") == "cuda_vmm" and ( + type(desc.get("allocation_nbytes")) is not int + or desc["allocation_nbytes"] < desc["nbytes"] + or type(desc.get("numa_id")) is not int + or desc["numa_id"] < 0 + ): + raise ValueError("Invalid CUDA host VMM descriptor") + if index in descriptors and descriptors[index] != desc: + raise ValueError("Conflicting host table descriptors") + descriptors[index] = desc + if set(descriptors) != set(range(count)): + raise ValueError("Host table descriptor/FD count mismatch") + received = {} + try: + for expected in range(count): + index, fd = _recv_fd(sock) + received[expected] = fd + if index != expected or ( + descriptors[expected].get("backend", "memfd") == "memfd" + and os.fstat(fd).st_size != descriptors[expected]["nbytes"] + ): + raise ValueError("Host table FD index/size mismatch") + self._host_received_fds = received + self._host_descriptors = descriptors + except Exception: + for fd in received.values(): + os.close(fd) + raise return result def import_tensor(self, entry: Dict[str, Any]) -> torch.Tensor: + if "host_table" in entry: + return self._import_host_tensor(entry) return MultiprocessingSerializer.deserialize(entry["handle"]) + def _import_host_tensor(self, entry): + from sglang.srt.layers.engram import _HostTable + + desc = entry["host_table"] + index = desc["index"] + if self._host_descriptors.get(index) != desc: + raise ValueError("Host table tensor descriptor mismatch") + shape, offset = entry["shape"], entry["host_offset"] + dtype = getattr(torch, entry["dtype"]) + itemsize = torch.empty((), dtype=dtype, device="cpu").element_size() + if ( + any(type(n) is not int or n < 0 for n in shape) + or type(offset) is not int + or offset < 0 + or offset % itemsize + ): + raise ValueError("Invalid host table tensor shape/offset") + size = math.prod(shape) * itemsize + if offset + size > desc["nbytes"]: + raise ValueError("Host table tensor exceeds its backing") + if index not in self._host_tables: + fd = self._host_received_fds.pop(index) + if desc.get("backend", "memfd") == "cuda_vmm": + self._host_tables[index] = _HostTable.from_vmm_fd( + desc["layout"], + desc["nbytes"], + fd, + desc["allocation_nbytes"], + desc["numa_id"], + ) + else: + self._host_tables[index] = _HostTable.from_fd( + desc["layout"], desc["nbytes"], fd + ) + return ( + self._host_tables[index] + .bytes.narrow(0, offset, size) + .view(dtype) + .view(shape) + ) + + def bind_host_tables(self, model, entries): + for index, table in self._host_tables.items(): + desc = self._host_descriptors[index] + module = model.get_submodule(desc["module"]) + if module.host_table is not None: + raise ValueError("Client allocated a host table before IPC recovery") + weight = entries[desc["module"] + ".weight"] + scale = entries[desc["module"] + ".scale"] + if ( + weight["host_offset"] != 0 + or scale["host_offset"] + != module.weight.numel() * module.weight.element_size() + ): + raise ValueError("Host Engram weight/scale layout mismatch") + module.host_table = table + module._host_table_ptrs = ( + table.device_ptr, + table.device_ptr + scale["host_offset"], + ) + + def __del__(self): + for fd in getattr(self, "_host_received_fds", {}).values(): + try: + os.close(fd) + except OSError: + pass + class VmmFdTransportBackend(WeightCacheTransportBackend): """Placeholder for the CUDA VMM + fd-passing transport. diff --git a/test/registered/unit/model_loader/test_dsv41_weight_cache_state.py b/test/registered/unit/model_loader/test_dsv41_weight_cache_state.py new file mode 100644 index 000000000000..c9b122c1fad4 --- /dev/null +++ b/test/registered/unit/model_loader/test_dsv41_weight_cache_state.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for processed-weight mapping; CUDA recovery is tested separately.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.weight_cache.ipc_loader import IpcModelLoader +from sglang.srt.weight_cache.model_state import ( + CONTRACT, + STATE_ENTRY, + TENSOR_PREFIX, + export_model_state, + model_state_contract, + restore_model_state, + validate_model_state, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class ToyLayer(nn.Module): + def __init__(self): + super().__init__() + self.weight = nn.Parameter(torch.empty(2, 4), requires_grad=False) + self.scale = nn.Parameter(torch.empty(2, 2), requires_grad=False) + self.cached_parts = None + self.ape_converted = False + self.quant_method = SimpleNamespace(use_mxfp8=False, weight_block_size=[32, 32]) + + +class ToyModel(nn.Module): + def __init__(self): + super().__init__() + self.layer = ToyLayer() + + +def config(): + return SimpleNamespace( + hf_config=SimpleNamespace(model_type="deepseek_v41"), dtype=torch.float32 + ) + + +def processed_state(): + model = ToyModel() + # Post-loading transformed shape/dtype differs from the constructor. + model.layer.weight = nn.Parameter( + torch.arange(8).reshape(2, 4).to(torch.int8), requires_grad=False + ) + model.layer.scale = nn.Parameter( + torch.arange(4).reshape(1, 4).float(), requires_grad=False + ) + model.layer.scale.format_ue8m0 = True + model.layer.cached_parts = (model.layer.weight[:, :2], model.layer.scale) + model.layer.ape_converted = True + model.layer.quant_method.use_mxfp8 = True + tensors = {k: (v, True) for k, v in model.state_dict().items()} + manifest = export_model_state(model, tensors) + entries = {k: {"tensor": v, "is_param": kind} for k, (v, kind) in tensors.items()} + entries[STATE_ENTRY] = {"model_state": manifest} + return model, tensors, manifest, entries + + +class TestDsv41WeightCacheState(CustomTestCase): + def test_daemon_initializes_dense_gemm_backends(self): + from sglang.srt.weight_cache.daemon import _initialize_gemm_configs + + with ( + patch( + "sglang.srt.layers.quantization.fp8_utils.initialize_fp8_gemm_config" + ) as fp8, + patch( + "sglang.srt.layers.quantization.fp4_utils.initialize_fp4_gemm_config" + ) as fp4, + patch( + "sglang.srt.layers.quantization.unquant.initialize_bf16_gemm_config" + ) as bf16, + ): + _initialize_gemm_configs() + fp8.assert_called_once_with() + fp4.assert_called_once_with() + bf16.assert_called_once_with() + + def test_router_config_rebinds_to_imported_bias(self): + model = ToyModel() + bias = nn.Parameter(torch.ones(4), requires_grad=False) + gate = nn.Module() + gate.e_score_correction_bias = bias + topk_config = SimpleNamespace(correction_bias=torch.empty(4, device="meta")) + layer = nn.Module() + layer.mlp = nn.Module() + layer.mlp.gate = gate + layer.mlp.topk = SimpleNamespace(topk_config=topk_config) + model.model = nn.Module() + model.model.layers = nn.ModuleList([layer]) + tensors = {k: (v, True) for k, v in model.state_dict().items()} + manifest = export_model_state(model, tensors) + restore_model_state(model, manifest, {}) + self.assertIs(topk_config.correction_bias, bias) + + def test_daemon_exports_shared_nonpersistent_buffer_aliases(self): + from sglang.srt.weight_cache.daemon import WeightCacheDaemon + + model = ToyModel() + model.other = ToyLayer() + shared = torch.ones(3) + model.layer.register_buffer("freqs_cis", shared, persistent=False) + model.other.register_buffer("freqs_cis", shared, persistent=False) + daemon = WeightCacheDaemon.__new__(WeightCacheDaemon) + daemon.model = model + daemon.gpu_id = 0 + daemon.state_entries = {} + daemon.config = SimpleNamespace(model_state_contract="") + backend = SimpleNamespace( + name="test", + prepare_export=lambda tensors: { + name: {"tensor": tensor, "is_param": kind} + for name, (tensor, kind) in tensors.items() + }, + ) + with patch( + "sglang.srt.weight_cache.daemon.choose_daemon_transport_backend", + return_value=backend, + ): + daemon._export_state() + self.assertEqual( + daemon.state_entries["layer.freqs_cis"]["tensor"].data_ptr(), + shared.data_ptr(), + ) + self.assertEqual( + daemon.state_entries["other.freqs_cis"]["tensor"].data_ptr(), + shared.data_ptr(), + ) + + def test_processed_tensors_map_without_reprocessing(self): + source, _, _, entries = processed_state() + loader = IpcModelLoader( + LoadConfig(model_loader_extra_config={"dsv41_weight_cache": True}) + ) + loader._transport_backend = SimpleNamespace(import_tensor=lambda e: e["tensor"]) + with patch( + "sglang.srt.weight_cache.ipc_loader._initialize_model", + side_effect=lambda *args: ToyModel(), + ): + model = loader._load_zero_copy_mode(config(), None, entries, None) + self.assertEqual(model.layer.weight.dtype, torch.int8) + self.assertEqual(model.layer.scale.shape, (1, 4)) + self.assertEqual(model.layer.weight.data_ptr(), source.layer.weight.data_ptr()) + self.assertEqual(model.layer.scale.data_ptr(), source.layer.scale.data_ptr()) + self.assertTrue(model.layer.scale.format_ue8m0) + self.assertTrue(model.layer.ape_converted) + self.assertTrue(model.layer.quant_method.use_mxfp8) + self.assertTrue( + torch.equal(model.layer.cached_parts[0], source.layer.cached_parts[0]) + ) + + def test_missing_contract_fails_before_mapping(self): + _, _, _, entries = processed_state() + entries.pop(STATE_ENTRY) + loader = IpcModelLoader( + LoadConfig(model_loader_extra_config={"dsv41_weight_cache": True}) + ) + with patch( + "sglang.srt.weight_cache.ipc_loader._initialize_model", + side_effect=lambda *args: ToyModel(), + ): + with self.assertRaisesRegex(ValueError, "missing"): + loader._load_zero_copy_mode(config(), None, entries, None) + + def test_bad_tensor_schema_is_rejected(self): + _, _, manifest, entries = processed_state() + manifest["schema"]["layer.scale"][0] = [99] + loader = IpcModelLoader( + LoadConfig(model_loader_extra_config={"dsv41_weight_cache": True}) + ) + loader._transport_backend = SimpleNamespace(import_tensor=lambda e: e["tensor"]) + with patch( + "sglang.srt.weight_cache.ipc_loader._initialize_model", + side_effect=lambda *args: ToyModel(), + ): + with self.assertRaisesRegex(ValueError, "schema mismatch"): + loader._load_zero_copy_mode(config(), None, entries, None) + + def test_class_drift_rejected(self): + _, _, manifest, _ = processed_state() + manifest["objects"]["layer"] = "wrong.Layer" + with self.assertRaisesRegex(ValueError, "class mismatch"): + validate_model_state(ToyModel(), manifest) + + def test_mhc_seam_does_not_duplicate_state_dict(self): + from sglang.srt.models.deepseek_v4 import DeepseekV4DecoderLayer + + layers = [] + for _ in range(2): + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + nn.Module.__init__(layer) + layer.hc_pre_from_prev_sublayer = True + layer.engram = None + layer.register_parameter("test_weight", nn.Parameter(torch.ones(1))) + layers.append(layer) + layers[0].bind_next(layers[1]) + self.assertIs(layers[0]._next_layer, layers[1]) + self.assertNotIn("_next_layer", layers[0]._modules) + self.assertEqual(list(layers[0].state_dict()), ["test_weight"]) + + def test_meta_hasher_skips_tokenizer_and_request_history(self): + from sglang.srt.layers.engram import EngramHasher, EngramLayout + from sglang.srt.runtime_context import publish, reset_context + from sglang.srt.server_args import ServerArgs + + reset_context() + self.addCleanup(reset_context) + publish( + ServerArgs(model_path="dummy", weight_cache_mode="client"), role="tokenizer" + ) + layout = EngramLayout( + max_ngram_size=3, + layer_ids=(1, 14), + num_embeddings=(20, 20), + primes=(((11,), (13,)), ((17,), (19,))), + n_heads=1, + head_dim=2, + ) + with ( + torch.device("meta"), + patch( + "sglang.srt.layers.engram.get_tokenizer", + side_effect=AssertionError("tokenizer rebuilt"), + ), + ): + hasher = EngramHasher.from_config(SimpleNamespace(vocab_size=20), layout) + self.assertTrue(hasher.token_map.is_meta) + self.assertIsNone(hasher.history) + self.assertEqual(hasher.max_ngram_size, 3) + + def test_weight_layout_fingerprint_changes_with_backend(self): + from sglang.srt.weight_cache.model_state import layout_fingerprint + + cfg = SimpleNamespace( + hf_config=SimpleNamespace(to_dict=lambda: {"model_type": "deepseek_v41"}) + ) + a = layout_fingerprint( + cfg, SimpleNamespace(fp8_gemm_runner_backend="deep_gemm") + ) + b = layout_fingerprint( + cfg, SimpleNamespace(fp8_gemm_runner_backend="flashinfer_cutedsl") + ) + self.assertNotEqual(a, b) + + def test_weight_layout_fingerprint_rejects_recovery_source_drift(self): + from pathlib import Path + + from sglang.srt.weight_cache.model_state import layout_fingerprint + + cfg = SimpleNamespace( + hf_config=SimpleNamespace(to_dict=lambda: {"model_type": "deepseek_v41"}) + ) + args = SimpleNamespace() + baseline = layout_fingerprint(cfg, args) + read_bytes = Path.read_bytes + for source in ( + "weight_cache/daemon.py", + "weight_cache/ipc_loader.py", + "model_executor/model_runner_components/load_model_utils.py", + "layers/quantization/fp4_utils.py", + "layers/quantization/unquant.py", + ): + with self.subTest(source=source): + + def changed_source(path): + data = read_bytes(path) + return ( + data + b"\n# different checkout\n" + if path.as_posix().endswith(source) + else data + ) + + with patch.object(Path, "read_bytes", changed_source): + self.assertNotEqual(baseline, layout_fingerprint(cfg, args)) + + def test_daemon_parser_does_not_consume_tp_alias(self): + from sglang.srt.weight_cache.daemon import _parse_worker_cli_args + + worker, remaining = _parse_worker_cli_args( + ["--tp", "4", "--model-path", "/model"] + ) + self.assertIsNone(worker.tp_rank) + self.assertIsNone(worker.gpu_id) + self.assertEqual(remaining, ["--tp", "4", "--model-path", "/model"]) + + def test_private_host_backing_cannot_be_exported(self): + from sglang.srt.weight_cache.model_state import host_table_exports + + model = ToyModel() + model.layer.host_table = SimpleNamespace(fd=None) + with self.assertRaisesRegex(ValueError, "shareable host storage"): + host_table_exports(model, {}) + + def test_daemon_json_opt_in_is_accepted(self): + self.assertEqual( + model_state_contract(config(), '{"dsv41_weight_cache":true}'), CONTRACT + ) + + def test_experiment_requires_explicit_opt_in(self): + with self.assertRaisesRegex(ValueError, "experimental"): + model_state_contract(config(), {}) + + def test_plain_derived_tensors_are_not_in_manifest(self): + _, tensors, manifest, _ = processed_state() + self.assertTrue(any(k.startswith(TENSOR_PREFIX) for k in tensors)) + # Protocol metadata contains no tensor payloads. + import json + + json.dumps(manifest) + + +class TestDraftCacheIsolation(CustomTestCase): + def test_only_target_selects_ipc(self): + from sglang.srt.configs.load_config import LoadFormat + from sglang.srt.model_executor.model_runner_components.load_model_utils import ( + build_load_config, + maybe_enable_ipc_weight_cache, + ) + from sglang.srt.runtime_context import publish, reset_context + from sglang.srt.server_args import ServerArgs + + reset_context() + self.addCleanup(reset_context) + args = ServerArgs(model_path="dummy", weight_cache_mode="client") + publish(args, role="tokenizer") + common = dict( + server_args=args, + tp_rank=0, + remote_instance_weight_transporter_engine=None, + remote_instance_weight_transporter_session_id=None, + draft_model_idx=None, + weight_cache_mode="client", + weight_cache_socket="/tmp/target.sock", + ) + target = build_load_config(**common) + draft = build_load_config(**common, is_draft_worker=True) + maybe_enable_ipc_weight_cache(load_config=target) + maybe_enable_ipc_weight_cache(load_config=draft) + self.assertEqual(target.load_format, LoadFormat.IPC_CACHE) + self.assertNotEqual(draft.load_format, LoadFormat.IPC_CACHE) + self.assertEqual(draft.weight_cache_mode, "off") + self.assertIsNone(draft.weight_cache_socket) + self.assertEqual(target.weight_cache_socket, "/tmp/target.sock") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_loader/test_host_engram_weight_cache.py b/test/registered/unit/model_loader/test_host_engram_weight_cache.py new file mode 100644 index 000000000000..86bc33d40951 --- /dev/null +++ b/test/registered/unit/model_loader/test_host_engram_weight_cache.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host Engram FD ownership, mapping, and loader recovery contracts.""" + +import array +import gc +import mmap +import os +import socket +import struct +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from torch import nn + +from sglang.srt.configs.load_config import LoadConfig +from sglang.srt.weight_cache.ipc_loader import IpcModelLoader +from sglang.srt.weight_cache.model_state import ( + STATE_ENTRY, + export_model_state, + host_table_exports, +) +from sglang.srt.weight_cache.protocol import recv_msg +from sglang.srt.weight_cache.transport import ( + TorchIpcTransportBackend, + _recv_fd, + _send_fd, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class HostLayer(nn.Module): + def __init__(self): + super().__init__() + self.host_table = None + self.weight = nn.Parameter( + torch.empty(2, 32, dtype=torch.float8_e4m3fn), requires_grad=False + ) + self.scale = nn.Parameter( + torch.empty(2, 1, dtype=torch.float8_e8m0fnu), requires_grad=False + ) + + +class HostModel(nn.Module): + def __init__(self): + super().__init__() + self.layer = HostLayer() + + +class TestHostEngramCache(CustomTestCase): + def test_memfd_shards_do_not_run_anonymous_hugepage_compaction(self): + from sglang.srt.layers.engram import _HostTable + + table = _HostTable.__new__(_HostTable) + table.layout, table.fd, table.dirty = "per_rank", 123, True + table.bytes = SimpleNamespace(data_ptr=lambda: 4096) + with ( + patch("sglang.srt.layers.engram._huge_pages_backing", return_value=(64, 0)), + patch("sglang.srt.layers.engram.drop_checkpoint_page_cache") as drop, + patch.object( + table, "_collapse", side_effect=AssertionError("anonymous-only") + ), + ): + table.finish_load("test") + drop.assert_not_called() + self.assertFalse(table.dirty) + + def test_shared_backing_survives_owner_fd_close_and_repeated_client_mapping(self): + + fd = os.memfd_create("engram-unit", 0) + os.ftruncate(fd, 66) + backing = mmap.mmap(fd, 66, flags=mmap.MAP_SHARED) + raw = torch.frombuffer(backing, dtype=torch.uint8) + raw.fill_(127) + source = HostModel() + source.layer.weight = nn.Parameter( + raw[:64].view(torch.float8_e4m3fn).view(2, 32), requires_grad=False + ) + source.layer.scale = nn.Parameter( + raw[64:].view(torch.float8_e8m0fnu).view(2, 1), requires_grad=False + ) + source.layer.host_table = SimpleNamespace( + fd=fd, bytes=raw, nbytes=66, layout="per_rank" + ) + tensors = {name: (t, True) for name, t in source.state_dict().items()} + host_entries, fds = host_table_exports(source, tensors) + manifest = export_model_state(source, tensors) + sender = TorchIpcTransportBackend() + sender.configure_host_table_exports(host_entries, fds) + entries = sender.prepare_export(tensors) + entries[STATE_ENTRY] = {"model_state": manifest} + register_calls = [] + cudart = SimpleNamespace( + cudaHostRegister=lambda ptr, size, flags: ( + register_calls.append((ptr, size)) or 0 + ), + cudaHostUnregister=lambda ptr: 0, + ) + try: + # Two clients receive independent FDs for the same physical backing. + receivers = [] + for _ in range(2): + receiver = TorchIpcTransportBackend() + a, b = socket.socketpair() + try: + sender.send_fetch_state_response( + a, config={}, entries=entries, pid=os.getpid() + ) + result = receiver.recv_fetch_state_response(b, recv_msg(b)) + self.assertEqual( + os.fstat(receiver._host_received_fds[0]).st_ino, + os.fstat(fd).st_ino, + ) + self.assertFalse(os.get_inheritable(receiver._host_received_fds[0])) + receivers.append((receiver, result)) + finally: + a.close() + b.close() + os.close(fd) + fd = None + for receiver, result in receivers: + loader = IpcModelLoader( + LoadConfig(model_loader_extra_config={"dsv41_weight_cache": True}) + ) + loader._transport_backend = receiver + cfg = SimpleNamespace( + hf_config=SimpleNamespace(model_type="deepseek_v41"), + dtype=torch.float32, + ) + with ( + patch( + "sglang.srt.weight_cache.ipc_loader._initialize_model", + side_effect=lambda *args: HostModel(), + ), + patch.object(torch.cuda, "cudart", return_value=cudart), + patch( + "sglang.srt.layers.engram._registered_device_ptr", + side_effect=lambda ptr: ptr + 4096, + ), + ): + model = loader._load_zero_copy_mode( + cfg, None, result["entries"], None + ) + self.assertEqual(model.layer.weight.device.type, "cpu") + self.assertNotEqual(model.layer.weight.data_ptr(), raw.data_ptr()) + self.assertEqual( + model.layer._host_table_ptrs[0], + model.layer.host_table.device_ptr, + ) + self.assertEqual( + model.layer._host_table_ptrs[1] + - model.layer._host_table_ptrs[0], + 64, + ) + raw[0] = 42 + self.assertEqual( + model.layer.weight.view(torch.uint8)[0, 0].item(), 42 + ) + table = model.layer.host_table + del model + gc.collect() + receiver._host_tables.clear() + table.close_import() + self.assertEqual( + len(register_calls), 2 + ) # once per table/client, not per tensor + finally: + if fd is not None: + os.close(fd) + del tensors, source, raw + gc.collect() + backing.close() + + def test_meta_embedding_does_not_allocate_or_register_host_memory(self): + from sglang.srt.environ import envs + from sglang.srt.layers.engram import EngramEmbedding, _HostTable + from sglang.srt.runtime_context import get_context, get_parallel + + with ( + get_context().override_server_args(model_path="dummy"), + get_parallel().override( + tp_size=2, + tp_rank=1, + attn_tp_size=2, + attn_tp_rank=1, + moe_tp_size=2, + moe_tp_rank=1, + ), + envs.SGLANG_ENABLE_DSV41_ENGRAM_HOST_TABLE.override(True), + envs.SGLANG_DSV41_ENGRAM_HOST_TABLE_LAYOUT.override("shared"), + patch.object( + _HostTable, "__init__", side_effect=AssertionError("must not allocate") + ), + torch.device("meta"), + ): + model = EngramEmbedding(8, 32, 1) + self.assertTrue(model.weight.is_meta) + self.assertEqual(tuple(model.weight.shape), (8, 32)) + self.assertIsNone(model.host_table) + + def test_fragmented_fd_header(self): + fd = os.memfd_create("fragmented", 0) + a, b = socket.socketpair() + try: + header = struct.pack("