From 99e55fc5da5361b490620c64786b520575f01a60 Mon Sep 17 00:00:00 2001
From: David Orman <ormandj@corenode.com>
Date: Sat, 3 Oct 2026 20:24:46 -0500
Subject: [PATCH 1/2] [Model Loader] Keep host-resident weights off the
 InstantTensor GPU stream

A model can name the checkpoint tensors it keeps in host memory through
host_resident_weight_patterns. Under --load-format instanttensor, the loader
reads the safetensors headers and streams only files without such tensors to
the GPU. Files that hold them load through the CPU safetensors iterator;
their other tensors are moved to the device, where InstantTensor places
everything else. DeepSeek-V4 and its DSpark draft name the Engram tables
when --engram-config places them on the CPU ("placement": "cpu").
---
 python/sglang/srt/model_loader/loader.py      |  77 ++++++++-
 python/sglang/srt/models/deepseek_v4.py       |   9 ++
 .../sglang/srt/models/deepseek_v4_dspark.py   |   2 +
 .../model_loader/test_instanttensor_loader.py | 150 +++++++++++++++++-
 4 files changed, 232 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py
index edda2a761822..e0e2e6664e1e 100644
--- a/python/sglang/srt/model_loader/loader.py
+++ b/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
@@ -359,6 +360,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 +455,12 @@ 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
+        through the CPU safetensors path instead."""
+
         @classmethod
         def init_new(cls, model_config: ModelConfig, model):
             return cls(
@@ -430,6 +472,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 +686,35 @@ 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 the safetensors path.",
+                    len(host_files),
+                )
+                host_weights = buffered_multi_thread_safetensors_weights_iterator(
+                    host_files,
+                    max_workers=self.DEFAULT_NUM_THREADS,
+                    disable_mmap=get_model().weight_loader_disable_mmap,
+                )
+                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
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 63e5ba148197..68263b2f20a7 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/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,
diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py
index 9385ab8cf551..8135407ff59e 100644
--- a/python/sglang/srt/models/deepseek_v4_dspark.py
+++ b/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(
diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py
index 10214390cdf6..181b6514f5ab 100644
--- a/test/registered/unit/model_loader/test_instanttensor_loader.py
+++ b/test/registered/unit/model_loader/test_instanttensor_loader.py
@@ -1,4 +1,5 @@
 import json
+import os
 import sys
 import tempfile
 import unittest
@@ -8,14 +9,15 @@
 from unittest.mock import MagicMock, patch
 
 import torch
-from safetensors.torch import save_file
+from safetensors.torch import safe_open, save_file
 
 import sglang.srt.model_loader.loader as loader_mod
 import sglang.srt.model_loader.weight_utils as weight_utils
+from sglang.srt.arg_groups.engram_config import EngramConfig, EngramPlacement
 from sglang.srt.configs.load_config import _DEFAULT_LOAD_GROUP, LoadConfig
 from sglang.srt.platforms.cuda import CudaSRTPlatform
 from sglang.srt.platforms.rocm import RocmSRTPlatform
-from sglang.srt.runtime_context import get_context
+from sglang.srt.runtime_context import get_context, get_exec
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -297,5 +299,149 @@ def test_pp_embedding_can_reopen_instanttensor_checkpoint(self):
         self.safe_open.assert_not_called()
 
 
+class TestInstantTensorHostResidentWeights(CustomTestCase):
+    """Checkpoint files holding host-resident tensors bypass InstantTensor."""
+
+    # Stands in for the GPU: tensors placed on the device show up as "meta".
+    DEVICE = torch.device("meta")
+
+    def setUp(self):
+        super().setUp()
+        folder = tempfile.TemporaryDirectory()
+        self.addCleanup(folder.cleanup)
+        # DeepSeek-V4.1 checkpoint names; the Engram file also holds a tensor
+        # that is not host-resident.
+        layout = {
+            "model-00001.safetensors": ["layers.0.attn.wq_a.weight"],
+            "model-00002.safetensors": [
+                "layers.1.engram.embed.weight",
+                "layers.1.engram.embed.scale",
+                "layers.1.engram.wkv.weight",
+            ],
+            "model-00003.safetensors": ["layers.2.ffn.w1.weight"],
+        }
+        torch.manual_seed(0)
+        self.files, self.expected = [], {}
+        for file_name, names in layout.items():
+            tensors = {name: torch.randn(4, 8) for name in names}
+            path = os.path.join(folder.name, file_name)
+            save_file(tensors, path)
+            self.files.append(path)
+            self.expected.update(tensors)
+        self.engram_file = self.files[1]
+
+        self.instanttensor_calls = []
+        patches = ExitStack()
+        self.addCleanup(patches.close)
+        patches.enter_context(
+            patch.object(
+                loader_mod,
+                "instanttensor_weights_iterator",
+                side_effect=self._fake_instanttensor,
+            )
+        )
+        patches.enter_context(
+            patch.object(
+                loader_mod.current_platform, "get_device", return_value=self.DEVICE
+            )
+        )
+        patches.enter_context(
+            patch.object(loader_mod.torch.cuda, "current_device", return_value=0)
+        )
+
+    def _fake_instanttensor(self, files, **_kwargs):
+        self.instanttensor_calls.append(list(files))
+        for path in files:
+            with safe_open(path, framework="pt") as f:
+                for name in f.keys():
+                    yield name, f.get_tensor(name).to(self.DEVICE)
+
+    def _load(self, source):
+        model_loader = loader_mod.DefaultModelLoader(
+            LoadConfig(load_format="instanttensor")
+        )
+        resolved = loader_mod.DefaultModelLoader.ResolvedSource(
+            source=source,
+            hf_folder=os.path.dirname(self.files[0]),
+            weight_files=tuple(self.files),
+            use_safetensors=True,
+        )
+        with get_context().override_server_args():
+            return list(
+                model_loader._get_weights_iterator(source, resolved_source=resolved)
+            )
+
+    def _dsv4_source(self, host_table: bool, model_cls=None):
+        from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+
+        model_cls = model_cls or DeepseekV4ForCausalLM
+        model = model_cls.__new__(model_cls)
+        model_config = SimpleNamespace(model_path="model", revision=None)
+        placement = EngramPlacement.CPU if host_table else EngramPlacement.GPU
+        with (
+            get_context().override_server_args(),
+            get_exec().offload.override(
+                engram_config=EngramConfig(placement=placement)
+            ),
+        ):
+            return loader_mod.DefaultModelLoader.Source.init_new(model_config, model)
+
+    def _host_resident_names(self, result):
+        """Check each tensor arrives once and return the names left on the host."""
+        self.assertCountEqual([name for name, _ in result], self.expected)
+        on_host = set()
+        for name, tensor in result:
+            expected = self.expected[name]
+            self.assertEqual(
+                (tensor.shape, tensor.dtype), (expected.shape, expected.dtype)
+            )
+            if tensor.device != self.DEVICE:
+                self.assertEqual(tensor.device.type, "cpu")
+                torch.testing.assert_close(tensor, expected, rtol=0, atol=0)
+                on_host.add(name)
+        return on_host
+
+    def test_engram_tables_stay_on_host(self):
+        result = self._load(self._dsv4_source(host_table=True))
+
+        self.assertEqual(self.instanttensor_calls, [[self.files[0], self.files[2]]])
+        # The Engram file's other tensor is read on the host but placed on the
+        # device like every InstantTensor tensor.
+        self.assertEqual(
+            self._host_resident_names(result),
+            {"layers.1.engram.embed.weight", "layers.1.engram.embed.scale"},
+        )
+
+    def test_device_tables_keep_every_file_on_instanttensor(self):
+        result = self._load(self._dsv4_source(host_table=False))
+
+        self.assertEqual(self.instanttensor_calls, [self.files])
+        self.assertEqual(self._host_resident_names(result), set())
+
+    def test_only_host_resident_files_skip_instanttensor(self):
+        self.files = [self.engram_file]
+        self.expected = {k: v for k, v in self.expected.items() if ".engram." in k}
+        result = self._load(self._dsv4_source(host_table=True))
+
+        self.assertEqual(self.instanttensor_calls, [])
+        self.assertEqual(len(self._host_resident_names(result)), 2)
+
+    def test_non_safetensors_files_are_left_to_instanttensor(self):
+        # InstantTensor reports unsupported files; the split must not open them.
+        self.assertEqual(
+            loader_mod._split_host_resident_files(
+                ["model.bin", self.engram_file], (".engram.embed.",)
+            ),
+            (["model.bin"], [self.engram_file]),
+        )
+
+    def test_dspark_draft_declares_the_same_tables(self):
+        from sglang.srt.models.deepseek_v4_dspark import DeepseekV4ForCausalLMDSpark
+
+        for enabled, expected in ((False, None), (True, (".engram.embed.",))):
+            source = self._dsv4_source(enabled, DeepseekV4ForCausalLMDSpark)
+            self.assertEqual(source.host_resident_weight_patterns, expected)
+
+
 if __name__ == "__main__":
     unittest.main()

From f415e581e9e92eef014d16b7e240b9d93c25a564 Mon Sep 17 00:00:00 2001
From: David Orman <ormandj@corenode.com>
Date: Sat, 3 Oct 2026 20:31:19 -0500
Subject: [PATCH 2/2] [DeepSeek-V4.1] Fill host-resident Engram rows with
 parallel positioned reads

With the Engram host table, each rank owns a row range of the Engram tables
in pinned host memory. The safetensors path copied those rows out of a
memory map, faulting the file in page by page. For files with host-resident
tensors, yield a SafetensorsRowSource for the matching names instead;
EngramEmbedding reads only its own rows with parallel preadv straight into
its shard; for per-rank shards ("replication": "none") it also asks the
kernel to drop the read range from the page cache.
---
 python/sglang/srt/layers/engram.py            |  14 +-
 python/sglang/srt/model_loader/loader.py      |  14 +-
 .../sglang/srt/model_loader/weight_utils.py   |  95 +++++++
 .../model_loader/test_host_resident_rows.py   | 241 ++++++++++++++++++
 .../model_loader/test_instanttensor_loader.py |  12 +
 5 files changed, 366 insertions(+), 10 deletions(-)
 create mode 100644 test/registered/unit/model_loader/test_host_resident_rows.py

diff --git a/python/sglang/srt/layers/engram.py b/python/sglang/srt/layers/engram.py
index 7d8e7c4b724b..58688776b9a0 100644
--- a/python/sglang/srt/layers/engram.py
+++ b/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
 
diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py
index e0e2e6664e1e..e593e8f3a56a 100644
--- a/python/sglang/srt/model_loader/loader.py
+++ b/python/sglang/srt/model_loader/loader.py
@@ -123,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,
@@ -458,8 +459,9 @@ class Source:
         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
-        through the CPU safetensors path instead."""
+        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):
@@ -701,13 +703,11 @@ def _get_weights_iterator(
             if host_files:
                 logger.info(
                     "Loading %d checkpoint file(s) with host-resident weights "
-                    "through the safetensors path.",
+                    "through positioned reads.",
                     len(host_files),
                 )
-                host_weights = buffered_multi_thread_safetensors_weights_iterator(
-                    host_files,
-                    max_workers=self.DEFAULT_NUM_THREADS,
-                    disable_mmap=get_model().weight_loader_disable_mmap,
+                host_weights = host_resident_weights_iterator(
+                    host_files, source.host_resident_weight_patterns
                 )
                 weights_iterator = itertools.chain(
                     weights_iterator,
diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py
index 64d781ea7ff1..ea8c50bdea9a 100644
--- a/python/sglang/srt/model_loader/weight_utils.py
+++ b/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,
diff --git a/test/registered/unit/model_loader/test_host_resident_rows.py b/test/registered/unit/model_loader/test_host_resident_rows.py
new file mode 100644
index 000000000000..bc9e54e100a9
--- /dev/null
+++ b/test/registered/unit/model_loader/test_host_resident_rows.py
@@ -0,0 +1,241 @@
+"""Positioned row reads of host-resident safetensors weights."""
+
+import errno
+import itertools
+import os
+import random
+import shutil
+import tempfile
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+from safetensors.torch import safe_open, save_file
+
+import sglang.srt.layers.engram as engram
+import sglang.srt.model_loader.weight_utils as weight_utils
+from sglang.srt.arg_groups.engram_config import (
+    EngramConfig,
+    EngramLookup,
+    EngramPlacement,
+    EngramReplication,
+)
+from sglang.srt.model_loader.weight_utils import (
+    SafetensorsRowSource,
+    host_resident_weights_iterator,
+)
+from sglang.srt.runtime_context import get_context, get_exec
+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")
+
+NUM_ROWS, DIM = 37, 64  # 37 rows split unevenly over 2, 3 and 4 ranks.
+# Hash-column segments of the table (prefix sums): 12 columns shard on column
+# boundaries over 1-4 ranks; 5 columns fall back to the equal-row split.
+COLUMN_BOUNDS = {
+    "columns": (0, 2, 5, 7, 11, 13, 17, 19, 23, 29, 31, 34, 37),
+    "rows": (0, 7, 14, 22, 29, 37),
+}
+WEIGHT = "layers.1.engram.embed.weight"
+SCALE = "layers.1.engram.embed.scale"
+FP32_TABLE = "layers.1.engram.embed.fp32"
+OTHER = "layers.1.engram.wkv.weight"
+HOST_RESIDENT = (WEIGHT, SCALE, FP32_TABLE)
+SENTINEL = 0xA5
+# --engram-config per layout: device shards, a host shard per rank, or one host
+# table per node. An explicit lookup skips the auto choice, which probes GPUs.
+LAYOUTS = {
+    None: EngramConfig(lookup=EngramLookup.ALL_REDUCE),
+    "per_rank": EngramConfig(
+        placement=EngramPlacement.CPU, lookup=EngramLookup.ALL_REDUCE
+    ),
+    "shared": EngramConfig(
+        placement=EngramPlacement.CPU,
+        replication=EngramReplication.PER_NODE,
+        lookup=EngramLookup.ALL_REDUCE,
+    ),
+}
+
+
+class _FakeHostTable:
+    """_HostTable without the memfd, the huge pages and cudaHostRegister."""
+
+    def __init__(self, replication, nbytes, name, group):
+        self.replication = replication
+        self.dirty = False
+        self.bytes = torch.full((nbytes,), SENTINEL, dtype=torch.uint8)
+        self.device_ptr = self.bytes.data_ptr()
+
+
+def _random_bytes(shape, dtype, generator):
+    raw = torch.randint(0, 256, shape, dtype=torch.uint8, generator=generator)
+    return raw.view(dtype)
+
+
+class TestHostResidentRows(CustomTestCase):
+    def setUp(self):
+        super().setUp()
+        folder = tempfile.TemporaryDirectory()
+        self.addCleanup(folder.cleanup)
+        self.path = os.path.join(folder.name, "model-00002.safetensors")
+        g = torch.Generator().manual_seed(0)
+        self.tensors = {
+            WEIGHT: _random_bytes((NUM_ROWS, DIM), torch.float8_e4m3fn, g),
+            SCALE: _random_bytes(
+                (NUM_ROWS, DIM // engram.FP8_BLOCK_SIZE), torch.float8_e8m0fnu, g
+            ),
+            # Multi-byte elements, so small chunks split elements across reads.
+            FP32_TABLE: torch.randn(NUM_ROWS, 5, generator=g),
+            OTHER: torch.randn(8, 16, generator=g),
+        }
+        save_file(self.tensors, self.path)
+        self.items = dict(host_resident_weights_iterator([self.path], (".embed.",)))
+        with safe_open(self.path, framework="pt") as f:
+            self.mmap = {name: f.get_tensor(name) for name in self.tensors}
+
+    def assertBytesEqual(self, actual, expected, msg=None):
+        self.assertEqual(actual.dtype, expected.dtype, msg)
+        self.assertTrue(
+            torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), msg
+        )
+
+    def test_iterator_yields_row_sources_for_matching_names(self):
+        self.assertCountEqual(self.items, self.tensors)
+        for name in HOST_RESIDENT:
+            source = self.items[name]
+            self.assertIsInstance(source, SafetensorsRowSource)
+            self.assertEqual(source.shape, tuple(self.tensors[name].shape))
+            self.assertEqual(source.dtype, self.tensors[name].dtype)
+        self.assertBytesEqual(self.items[OTHER], self.tensors[OTHER])
+
+    def test_random_row_ranges_match_mmap(self):
+        rng = random.Random(0)
+        for name in HOST_RESIDENT:
+            full = self.mmap[name]
+            for _ in range(30):
+                start = rng.randrange(NUM_ROWS + 1)
+                rows = rng.randrange(NUM_ROWS - start + 1)
+                chunk_bytes = rng.choice([1, 3, 8, 77, 1 << 20])
+                num_threads = rng.choice([1, 3, 8])
+                dst = torch.empty((rows, *full.shape[1:]), dtype=full.dtype)
+                self.items[name].read_rows_into(
+                    dst, start, num_threads=num_threads, chunk_bytes=chunk_bytes
+                )
+                self.assertBytesEqual(
+                    dst,
+                    full[start : start + rows],
+                    f"{name} rows {start}:{start + rows}, chunk {chunk_bytes}",
+                )
+
+    def test_short_reads_and_failed_cache_advice(self):
+        real_preadv = os.preadv
+
+        def short_preadv(fd, buffers, offset):
+            return real_preadv(fd, [buffers[0][:5]], offset)
+
+        def failing_fadvise(*_args):
+            raise OSError(errno.ENOSYS, "posix_fadvise")
+
+        dst = torch.empty(10, 5)
+        with (
+            patch.object(weight_utils.os, "preadv", short_preadv),
+            patch.object(
+                weight_utils.os, "posix_fadvise", failing_fadvise, create=True
+            ),
+            patch.object(weight_utils.os, "POSIX_FADV_DONTNEED", 4, create=True),
+        ):
+            self.items[FP32_TABLE].read_rows_into(
+                dst, 3, drop_page_cache=True, num_threads=3, chunk_bytes=64
+            )
+        self.assertBytesEqual(dst, self.mmap[FP32_TABLE][3:13])
+
+    def test_page_cache_eviction_only_for_per_rank_tables(self):
+        for layout, expected in ((None, False), ("per_rank", True), ("shared", False)):
+            calls = []
+            with (
+                patch.object(
+                    weight_utils.os,
+                    "posix_fadvise",
+                    lambda *args: calls.append(args),
+                    create=True,
+                ),
+                patch.object(weight_utils.os, "POSIX_FADV_DONTNEED", 4, create=True),
+            ):
+                self._load_engram(2, 1, layout, self.items)
+            self.assertEqual(bool(calls), expected, layout)
+
+    def test_truncated_file_raises(self):
+        path = self.path + ".copy"
+        shutil.copyfile(self.path, path)
+        source = dict(host_resident_weights_iterator([path], (".embed.",)))[WEIGHT]
+        os.truncate(path, source.offset + 2 * source.row_bytes)
+
+        dst = torch.empty(2, DIM, dtype=torch.float8_e4m3fn)
+        source.read_rows_into(dst, 0)
+        self.assertBytesEqual(dst, self.tensors[WEIGHT][:2])
+        with self.assertRaises(EOFError):
+            source.read_rows_into(torch.empty_like(dst), 1)
+
+    def _load_engram(self, tp_size, tp_rank, layout, weights, split="columns"):
+        """Load WEIGHT and SCALE into one rank's EngramEmbedding."""
+        parallel = SimpleNamespace(tp_size=tp_size, tp_rank=tp_rank, tp_group=None)
+        with (
+            patch.object(engram, "get_parallel", return_value=parallel),
+            patch.object(engram, "_HostTable", _FakeHostTable),
+            patch.object(engram, "_sorted_host_reads", return_value=False),
+            patch.object(engram.torch.cuda, "current_device", return_value=0),
+            get_context().override_server_args(),
+            get_exec().offload.override(engram_config=LAYOUTS[layout]),
+        ):
+            embed = engram.EngramEmbedding(
+                NUM_ROWS, DIM, layer_id=1, column_bounds=COLUMN_BOUNDS[split]
+            )
+        for param, name in ((embed.weight, WEIGHT), (embed.scale, SCALE)):
+            param.weight_loader(param, weights[name])
+        return embed
+
+    def test_engram_shards_match_mmap_path(self):
+        # Device rows, and both host-table layouts with weight and scale
+        # carved out of one host buffer.
+        for layout, split, tp_size in itertools.product(
+            (None, "per_rank", "shared"), COLUMN_BOUNDS, (1, 2, 3, 4)
+        ):
+            for tp_rank in range(tp_size):
+                msg = f"{layout=} {split=} {tp_size=} {tp_rank=}"
+                via_mmap = self._load_engram(tp_size, tp_rank, layout, self.mmap, split)
+                via_rows = self._load_engram(
+                    tp_size, tp_rank, layout, self.items, split
+                )
+                rows = slice(via_rows.row_start, via_rows.row_start + via_rows.rows)
+                for name in ("weight", "scale"):
+                    actual = getattr(via_rows, name).data
+                    self.assertBytesEqual(actual, getattr(via_mmap, name).data, msg)
+                    own = actual[rows] if layout == "shared" else actual
+                    expected = self.tensors[WEIGHT if name == "weight" else SCALE]
+                    self.assertBytesEqual(own, expected[rows], msg)
+                    if layout == "shared":
+                        # Other ranks' rows are left untouched.
+                        others = torch.ones(NUM_ROWS, dtype=torch.bool)
+                        others[rows] = False
+                        untouched = actual[others].view(torch.uint8) == SENTINEL
+                        self.assertTrue(untouched.all(), msg)
+                if layout is not None:
+                    self.assertTrue(via_rows.host_table.dirty, msg)
+
+    def test_rejects_mismatched_destinations(self):
+        source = self.items[WEIGHT]
+        for dst, row_start in (
+            (torch.empty(3, DIM, dtype=torch.float8_e5m2), 0),  # same size, wrong dtype
+            (torch.empty(3, DIM - 1, dtype=torch.float8_e4m3fn), 0),
+            (torch.empty(4, DIM, dtype=torch.float8_e4m3fn), NUM_ROWS - 3),
+            (torch.empty(DIM, 3, dtype=torch.float8_e4m3fn).t(), 0),
+        ):
+            with self.subTest(shape=tuple(dst.shape), row_start=row_start):
+                with self.assertRaises(ValueError):
+                    source.read_rows_into(dst, row_start)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/registered/unit/model_loader/test_instanttensor_loader.py b/test/registered/unit/model_loader/test_instanttensor_loader.py
index 181b6514f5ab..8e2d0b03063f 100644
--- a/test/registered/unit/model_loader/test_instanttensor_loader.py
+++ b/test/registered/unit/model_loader/test_instanttensor_loader.py
@@ -391,6 +391,9 @@ def _host_resident_names(self, result):
         self.assertCountEqual([name for name, _ in result], self.expected)
         on_host = set()
         for name, tensor in result:
+            if isinstance(tensor, weight_utils.SafetensorsRowSource):
+                rows, tensor = tensor, torch.empty(tensor.shape, dtype=tensor.dtype)
+                rows.read_rows_into(tensor, 0)
             expected = self.expected[name]
             self.assertEqual(
                 (tensor.shape, tensor.dtype), (expected.shape, expected.dtype)
@@ -403,6 +406,15 @@ def _host_resident_names(self, result):
 
     def test_engram_tables_stay_on_host(self):
         result = self._load(self._dsv4_source(host_table=True))
+        row_sources = {
+            name
+            for name, tensor in result
+            if isinstance(tensor, weight_utils.SafetensorsRowSource)
+        }
+        self.assertEqual(
+            row_sources,
+            {"layers.1.engram.embed.weight", "layers.1.engram.embed.scale"},
+        )
 
         self.assertEqual(self.instanttensor_calls, [[self.files[0], self.files[2]]])
         # The Engram file's other tensor is read on the host but placed on the