[Perf] dsv41: Speed up Engram host-table loading with parallel positioned reads

python/sglang/srt/layers/engram.py

@@ -1132,11 +1132,19 @@ def _table_ptrs(self) -> tuple[int, int]:
         return self._host_table_ptrs
 
     def _load_rows(self, param: nn.Parameter, loaded_weight: torch.Tensor):
+        from sglang.srt.model_loader.weight_utils import SafetensorsRowSource
+
         rows = slice(self.row_start, self.row_start + self.rows)
-        if self._shared:
-            param.data[rows].copy_(loaded_weight[rows])
+        dst = param.data[rows] if self._shared else param.data
+        if isinstance(loaded_weight, SafetensorsRowSource):
+            # Page-cache eviction was measured with per-rank host shards only.
+            loaded_weight.read_rows_into(
+                dst,
+                self.row_start,
+                drop_page_cache=self.host_table is not None and not self._shared,
+            )
         else:
-            param.data.copy_(loaded_weight[rows])
+            dst.copy_(loaded_weight[rows])
         if self.host_table is not None:
             self.host_table.dirty = True
 

python/sglang/srt/model_loader/loader.py

@@ -12,6 +12,7 @@
 import gc
 import glob
 import hashlib
+import itertools
 import json
 import logging
 import math
@@ -122,6 +123,7 @@
     get_gguf_extra_tensor_names,
     get_quant_config,
     gguf_quant_weights_iterator,
+    host_resident_weights_iterator,
     initialize_capture_safe_weights,
     initialize_dummy_weights,
     instanttensor_weights_iterator,
@@ -359,6 +361,41 @@ def load_model(
         raise NotImplementedError
 
 
+def _split_host_resident_files(
+    files: List[str], patterns: Optional[Tuple[str, ...]]
+) -> Tuple[List[str], List[str]]:
+    """Split safetensors files into (device, host) groups, keeping their order.
+
+    A file goes to the host group when any tensor name in its header contains
+    one of ``patterns``. Only headers are read.
+    """
+    if not patterns:
+        return list(files), []
+    from safetensors.torch import safe_open
+
+    device_files, host_files = [], []
+    for path in files:
+        host = False
+        # Other files stay in the device group; InstantTensor rejects them.
+        if path.endswith(".safetensors"):
+            with safe_open(path, framework="pt") as f:
+                host = any(p in name for name in f.keys() for p in patterns)
+        (host_files if host else device_files).append(path)
+    return device_files, host_files
+
+
+def _to_device_unless_host_resident(
+    weights: Iterable[Tuple[str, torch.Tensor]], patterns: Tuple[str, ...]
+) -> Generator[Tuple[str, torch.Tensor], None, None]:
+    """Move tensors that are not host-resident to the current device, where
+    InstantTensor puts every other tensor, so tensors that are combined
+    during loading share a device."""
+    for name, tensor in weights:
+        if not any(p in name for p in patterns):
+            tensor = tensor.to(current_platform.get_device(torch.cuda.current_device()))
+        yield name, tensor
+
+
 def _validate_default_loader_extra_config(
     *, extra_config: dict, load_format: LoadFormat
 ) -> None:
@@ -419,6 +456,13 @@ class Source:
         model_config: Optional[ModelConfig] = None
         """The model configuration (for checking architecture, etc)."""
 
+        host_resident_weight_patterns: Optional[Tuple[str, ...]] = None
+        """Substrings of checkpoint tensor names the model keeps in host memory.
+
+        Device-direct loaders (InstantTensor) load the files containing them on
+        the host and yield matching tensors as ``SafetensorsRowSource``, which
+        the parameters' weight loaders must accept."""
+
         @classmethod
         def init_new(cls, model_config: ModelConfig, model):
             return cls(
@@ -430,6 +474,9 @@ def init_new(cls, model_config: ModelConfig, model):
                     model, "allow_patterns_overrides", None
                 ),
                 model_config=model_config,
+                host_resident_weight_patterns=getattr(
+                    model, "host_resident_weight_patterns", None
+                ),
             )
 
     @dataclasses.dataclass(frozen=True)
@@ -641,11 +688,33 @@ def _get_weights_iterator(
                 hf_weights_files,
             )
         elif self.load_config.load_format == LoadFormat.INSTANTTENSOR:
-            weights_iterator = instanttensor_weights_iterator(
-                hf_weights_files,
-                extra_config=extra_config,
-                load_group=self.load_config.load_group,
+            device_files, host_files = _split_host_resident_files(
+                hf_weights_files, source.host_resident_weight_patterns
+            )
+            weights_iterator = (
+                instanttensor_weights_iterator(
+                    device_files,
+                    extra_config=extra_config,
+                    load_group=self.load_config.load_group,
+                )
+                if device_files
+                else iter(())
             )
+            if host_files:
+                logger.info(
+                    "Loading %d checkpoint file(s) with host-resident weights "
+                    "through positioned reads.",
+                    len(host_files),
+                )
+                host_weights = host_resident_weights_iterator(
+                    host_files, source.host_resident_weight_patterns
+                )
+                weights_iterator = itertools.chain(
+                    weights_iterator,
+                    _to_device_unless_host_resident(
+                        host_weights, source.host_resident_weight_patterns
+                    ),
+                )
         elif use_safetensors:
             weight_loader_disable_mmap = get_model().weight_loader_disable_mmap
             configured_prefetch = get_model().weight_loader_prefetch_checkpoints

python/sglang/srt/model_loader/weight_utils.py

@@ -12,6 +12,7 @@
 import itertools
 import json
 import logging
+import math
 import os
 import re
 import struct
@@ -1192,6 +1193,100 @@ def get_pp_stage_load_group() -> LoadGroup:
     return parallel.tp_group if parallel.pp_size > 1 else _DEFAULT_LOAD_GROUP
 
 
+class SafetensorsRowSource:
+    """A row-major safetensors tensor that is read by row range on demand.
+
+    A rank's rows are one contiguous byte range of the file, so a row-sharded
+    host-resident weight can fill its own destination with positioned reads
+    instead of copying out of a memory map page by page.
+    """
+
+    def __init__(
+        self, path: str, offset: int, shape: Tuple[int, ...], dtype: torch.dtype
+    ):
+        self.path = path
+        self.offset = offset
+        self.shape = tuple(shape)
+        self.dtype = dtype
+        self.row_bytes = math.prod(self.shape[1:]) * dtype.itemsize
+
+    def read_rows_into(
+        self,
+        dst: torch.Tensor,
+        row_start: int,
+        drop_page_cache: bool = False,
+        num_threads: int = 8,
+        chunk_bytes: int = 64 << 20,
+    ) -> None:
+        """Read rows ``[row_start, row_start + len(dst))`` into ``dst``.
+
+        With ``drop_page_cache``, ask the kernel to evict each range once read.
+        """
+        rows = dst.shape[0]
+        if dst.device.type != "cpu" or not dst.is_contiguous():
+            raise ValueError("SafetensorsRowSource reads into contiguous CPU tensors")
+        if (
+            dst.dtype != self.dtype
+            or tuple(dst.shape[1:]) != self.shape[1:]
+            or row_start < 0
+            or row_start + rows > self.shape[0]
+        ):
+            raise ValueError(
+                f"rows {row_start}:{row_start + rows} of {self.dtype} {self.shape} "
+                f"do not fit {dst.dtype} {tuple(dst.shape)}"
+            )
+        total = rows * self.row_bytes
+        if total == 0:
+            return
+        view = memoryview(dst.view(-1).view(torch.uint8).numpy())
+        base = self.offset + row_start * self.row_bytes
+        fd = os.open(self.path, os.O_RDONLY)
+        try:
+
+            def read(start: int) -> None:
+                end = min(start + chunk_bytes, total)
+                done = start
+                while done < end:
+                    n = os.preadv(fd, [view[done:end]], base + done)
+                    if n <= 0:
+                        raise EOFError(f"short read from {self.path} at {base + done}")
+                    done += n
+                # The bytes now live in dst; ask the kernel to drop the cached copy.
+                if drop_page_cache and hasattr(os, "posix_fadvise"):
+                    try:
+                        os.posix_fadvise(
+                            fd, base + start, end - start, os.POSIX_FADV_DONTNEED
+                        )
+                    except OSError as e:
+                        logger.debug("posix_fadvise failed for %s: %s", self.path, e)
+
+            with concurrent.futures.ThreadPoolExecutor(num_threads) as pool:
+                list(pool.map(read, range(0, total, chunk_bytes)))
+        finally:
+            os.close(fd)
+
+
+def host_resident_weights_iterator(
+    hf_weights_files: List[str], patterns: Tuple[str, ...]
+) -> Generator[Tuple[str, Union[torch.Tensor, SafetensorsRowSource]], None, None]:
+    """Iterate over safetensors weights, yielding a SafetensorsRowSource
+    instead of a tensor for names that contain one of ``patterns``."""
+    for path in hf_weights_files:
+        with open(path, "rb") as f:
+            (header_len,) = struct.unpack("<Q", f.read(8))
+            header = json.loads(f.read(header_len))
+        with safetensors.safe_open(path, framework="pt", device="cpu") as f:
+            for name in f.keys():
+                # A memory-mapped view: no data is read here.
+                tensor = f.get_tensor(name)
+                if any(p in name for p in patterns):
+                    begin, _ = header[name]["data_offsets"]
+                    tensor = SafetensorsRowSource(
+                        path, 8 + header_len + begin, tensor.shape, tensor.dtype
+                    )
+                yield name, tensor
+
+
 def instanttensor_weights_iterator(
     hf_weights_files: List[str],
     extra_config: Optional[dict] = None,

python/sglang/srt/models/deepseek_v4.py

@@ -47,6 +47,7 @@
 from sglang.kernels.ops.quantization.fp8_kernel import (
     sglang_per_token_group_quant_fp8,
 )
+from sglang.srt.arg_groups.engram_config import EngramPlacement
 from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
@@ -4629,6 +4630,14 @@ def get_hf_to_sglang_mapper(cls, config) -> Optional[WeightsMapper]:
         # written against different names, so only Quark checkpoints get it.
         return _QUARK_CKPT_TO_SGLANG_MAPPER if _is_quark_checkpoint(config) else None
 
+    @property
+    def host_resident_weight_patterns(self) -> Optional[Tuple[str, ...]]:
+        """Checkpoint names of the Engram tables kept in host memory."""
+        config = get_exec().offload.engram_config
+        if config is not None and config.placement == EngramPlacement.CPU:
+            return (".engram.embed.",)
+        return None
+
     def __init__(
         self,
         config: DeepSeekV4Config,

python/sglang/srt/models/deepseek_v4_dspark.py

@@ -694,6 +694,8 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
     precompile_kernels_after_loading = (
         DeepseekV4ForCausalLM.precompile_kernels_after_loading
     )
+    # The draft reads the target checkpoint, Engram files included.
+    host_resident_weight_patterns = DeepseekV4ForCausalLM.host_resident_weight_patterns
 
     @classmethod
     def shared_experts_fusion_disable_reason(