[ormandj]

Stacked on #42919 ("[Model Loader] Let InstantTensor load DeepSeek-V4.1 with the Engram host table"): this branch carries its commit, and the diff to review is the last commit, "[DeepSeek-V4.1] Fill host-resident Engram rows with parallel positioned reads". #40453 (InstantTensor loader by YitaoYuan) is now on main. I will rebase this PR when #42919 lands.

Motivation

With --engram-config '{"placement": "cpu"}', the two Engram tables (188.8 GiB of rows and scales) live in pinned host memory, and each TP rank loads only its own row range. With one shard per rank ("replication": "none"), a rank holds 47.2 GiB at TP4. Under --load-format instanttensor, #42919 loads these files through the CPU safetensors iterator. That iterator returns memory-mapped tensors, so EngramEmbedding._load_rows copies its rows out of the file mapping one page fault at a time. On the setup below, that phase took 68.3 s (2.8 GiB/s aggregate over 4 ranks). This was more than half of the target model's load time.

A rank's rows are one contiguous byte range of the file. Large parallel reads straight into the pinned shard are much faster.

Modifications

• weight_utils.SafetensorsRowSource(path, offset, shape, dtype): a row-major checkpoint tensor that can be read by row range. read_rows_into(dst, row_start) checks the destination: contiguous CPU tensor, same dtype and row shape, range in bounds. It then fills dst with parallel os.preadv calls (8 threads, 64 MiB chunks) at the right file offsets, retrying short reads; a premature end of file raises. With drop_page_cache=True, it calls posix_fadvise(POSIX_FADV_DONTNEED) after each chunk where available, so the kernel can drop the cached copy. That call is advisory, and its errors are logged at debug level.
• weight_utils.host_resident_weights_iterator(files, patterns): reads offsets from the safetensors header. It yields a SafetensorsRowSource for tensor names that match the model's host_resident_weight_patterns, and memory-mapped tensors from safe_open for all other names.
• Under --load-format instanttensor, files with host-resident tensors now use this iterator. The Source.host_resident_weight_patterns docstring states that the matching parameters' weight loaders must accept a row source. --weight-loader-disable-mmap is not applied to these files: the tables are read with preadv, and the few other tensors in the files come from safe_open. With that flag, the previous path could not load these files: safetensors.torch.load (safetensors 0.8.0) has no F8_E8M0 dtype.
• EngramEmbedding._load_rows reads only its own rows from a row source into the same destination as before: the rank's slice of the full table in the shared layout, or its whole shard in the per_rank layout. It requests page-cache eviction only in the per_rank layout, where this was measured; the shared layout keeps the page cache as before. Tensor inputs keep the existing copy. Row sources are only produced for host tables, so device tables are unchanged. The DSpark draft drops Engram names before it touches the value, so it reads no Engram rows.

Accuracy Tests

Load-path change only; the rows loaded are byte-identical. New test/registered/unit/model_loader/test_host_resident_rows.py (CPU) uses a checkpoint with an FP8 E4M3 table, its E8M0 scales, and an FP32 table:

• Random row ranges read through read_rows_into match the memory-mapped tensor byte for byte. Small chunks (1-77 bytes) and 1/3/8 threads exercise chunk boundaries inside rows and inside elements.
• A real EngramEmbedding is built for every rank at TP 1-4, with 37 rows (uneven splits). The per_rank and shared host layouts use a plain CPU buffer in place of the pinned host table, so the weight and scale are views into one buffer. The tensor path with CPU parameters is covered too. Loading from row sources gives the same bytes as loading from the memory-mapped tensors. In the shared layout, other ranks' rows are left untouched.
• Short reads are retried, a failing posix_fadvise does not fail the load, and a truncated file raises EOFError.
• EngramEmbedding requests page-cache eviction in the per_rank layout only, not in shared or for device tables.
• Destinations with a different dtype of the same size, a different row shape, an out-of-range row range, or non-contiguous memory are rejected.
• In test_instanttensor_loader.py, the host-resident names arrive as row sources with the full contents.

On this branch (current main plus #40453 and #42919), both files pass (21 tests) on macOS (CPU) and on Linux aarch64 (CPU, PyTorch 2.14.1); the Linux run covers the real posix_fadvise path. The rest of test/registered/unit/model_loader/ also passes on macOS (139 passed, 1 skipped in total; test_modelopt_loader.py was not collected locally because tvm_ffi is not installed). Mutating the read offset, the chunk ranges, the short-read retry, the rank's row_start, or the dtype check makes the new tests fail.

Earlier, on an older main with the same production diff, a separate check used the real InstantTensor library on one RTX PRO 6000 Blackwell GPU. Rows read into pinned buffers for every rank at TP 1-4 matched the memory map byte for byte and copied to the GPU unchanged. Pinned, cudaHostRegistered host tables were exercised by the serving runs below.

#42787 (column-aligned Engram shards) changes how row_start and rows are chosen; _load_rows only consumes them, and the two branches merge without conflicts.

2026-10-29: rebased onto main 974831643f, where #40453 and #42787 have landed (the branch now carries only the commit of #42919 and this one) and #46893 replaced the host-table variables with --engram-config, so the per_rank layout is now "placement": "cpu" with "replication": "none" (where page-cache eviction applies) and shared is "replication": "per_node"; the row test selects them through --engram-config and covers both the column-aligned and the equal-row split, both test files pass on the new head (21 tests, macOS CPU), and the SM120 GPU results in this description are from the previous head and will be re-run.

Speed Tests and Profiling

DeepSeek-V4.1-Flash, TP4/EP4, 4x RTX PRO 6000 Blackwell (SM120, PCIe), DSpark draft, host Engram table (per-rank layout), checkpoint on local NVMe (ZFS), on an older main with the same production diff. A/B of this change on the same build, one server start each, timed from the server log. The "Engram files" time runs from the end of the InstantTensor stream to the Engram host-table completion message on TP0:

	Engram files (188.8 GiB)	Target model load (TP0)	Process start to ready
before (#42919)	68.3 s (2.8 GiB/s)	119.9 s	3 m 43 s
this PR	13.5 s (14.0 GiB/s)	64.0 s	2 m 46 s

For reference, the default safetensors loader takes 264.2 s for the target model on the same setup. The time spent on the device files (about 40 s) and the DSpark draft load (40.0 s before, 39.5 s after) did not change.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

Developed with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #40961112832
Latest PR Test (Extra): ❌ Run #40961112286
Latest PR Test (AMD ROCm 10): ❌ Run #40961112526