Summary

Two opt-in Python-layer additions that together remove the last two host-side copies-worth of latency between object storage and GPU weights:

1. Pinned-memory destination buffers — let the streamer land bytes in page-locked memory so the consumer's H2D copy takes the DMA fast path (~5.3× faster than the pageable path np.empty forces today).
2. Fused dtype casting — optionally convert tensors to a target dtype during the single copy out of the stream buffer, because a dtype-converting H2D copy from pinned memory falls off the DMA fast path entirely (~9.3× slower — as slow as pageable), silently forfeiting benefit for the very common fp32-checkpoint → bf16-model case.

Motivation: two measured cliffs

We ship a diffusion-model weight-loading pipeline in vllm-project/vllm-omni#5105 and measured the following on RTX PRO 6000 Blackwell (PCIe 5.0 x16), per 4 GiB fp32:

Cliff #1: pageable stream buffer

FilesRequestsIteratorWithBuffer currently allocates with np.empty(buffer_size, dtype=np.uint8), which is pageable. Every downstream param.copy_() pays the 5.3× pageable penalty regardless of how fast the storage side is. This may also be related to the 6× slower-than-disk S3 loading reported in #120.

Cliff #2: pinned + dtype-converting H2D

Even with a pinned buffer, letting the consumer do the dtype conversion during H2D throws the fast path away. The conversion must happen on the host, fused into the copy that fills the pinned buffer. This matters because fp32 checkpoints served as bf16/fp16 are the norm in open models — Wan 2.x, LTX-Video and SDXL all ship fp32-only weights (~2.5 M downloads/month combined) — and it also halves the bytes that cross PCIe for fp32→bf16.

Proposal

1. Pinned destinations (RUNAI_STREAMER_PIN_MEMORY=1 / pin_memory=True)

Replace the internal buffer allocation with page-locked memory when requested:

# FilesRequestsIteratorWithBuffer
if pin_memory:
    self._backing = torch.empty(buffer_size, dtype=torch.uint8, pin_memory=True)
    self.buffer = self._backing.numpy()
else:
    self.buffer = np.empty(buffer_size, dtype=np.uint8)  # today

Notes:

• Zero change to the request/response flow; libstreamer writes into the same memoryviews — the C++ engine is agnostic.
• cudaHostAlloc of a multi-GB buffer costs ~100 ms/GB once; amortized over a model load it is negligible, but the flag stays opt-in for short-lived or memory-constrained uses.
• Fallback: if pinned allocation fails (RLIMIT_MEMLOCK, containers), log and degrade to pageable — never fail the load.

2. Fused dtype casting (get_tensors(target_dtypes=...))

streamer.stream_file(path)
for name, tensor in streamer.get_tensors(target_dtypes={"blocks.0.ffn.weight": torch.bfloat16}):
    ...  # tensor is already bf16, in (pinned) host memory

• target_dtypes maps tensor names to destination dtypes; unmapped tensors keep today's zero-copy view behavior.
• For mapped tensors the streamer performs one dst.copy_(src_view) on the host — the same rounding the consumer's param.copy_() would apply later, so results are bit-identical to the status quo.
• Non-float and quantized dtypes (fp8, int8, etc.) would never be auto-cast: scale tensors must reach the consumer bit-exact.
• Casting also reduces buffer pressure and PCIe bytes for downcasts (fp32→bf16 halves both), which interacts favorably with RUNAI_STREAMER_MEMORY_LIMIT.

Where the changes go (Python-only)

No changes to libstreamer (C++) — runai_request already accepts arbitrary caller-provided destination memoryviews.

Compatibility

• Both features are opt-in; defaults preserve current behavior exactly.
• No libstreamer (C++) changes required.
• torch is already a dependency of the safetensors layer.

Related issues

• runai_streamer: S3 loading 6x slower (800MB/s) than local disk (4.9GiB/s) #120 — S3 loading 6× slower than disk: the pageable H2D cliff may be a contributing factor
• Distributed Model Streamer for Multi-GPU/Multi-Node Loading #84 — Distributed Model Streamer: pinned buffers compose well with distributed loading; fused casting halves the bytes per broadcast

Validation we can bring

• Microbenchmarks for the two cliffs (script included in the PR).
• End-to-end A/B in vllm-omni: S3 → pinned → TP2 GPU loads of Wan2.2 (118 GB fp32) and HunyuanImage-3.0 (157 GB bf16 MoE), byte-identical outputs across ~30 runs.
• We are happy to split this into two PRs (pinned buffers first, fused cast second) or land them together — maintainer preference.