Module fast_weights

Module fast_weights 

Source
Expand description

Fast safetensors loader (InstantTensor-style) — pure Rust.

Two wins over the mmap-based loader in crate::weights:

  1. O_DIRECT reads. Bypasses the OS page cache, so the bytes never compete with GPU allocations on GB10 unified memory. The mmap path already works around this with POSIX_FADV_DONTNEED post-load; here we avoid the pollution in the first place.
  2. Pipelined read/copy. One background reader thread fetches the next tensor while the main thread does copy_h2d for the current one. Overlaps disk I/O with the host→device memcpy.

Behavioural parity with crate::weights::SafetensorsLoader is preserved — same EP filtering, same OOM pre-flight, same UVM fallback on GPU allocation failure, same extra-weights handling.

Structs§

FastSafetensorsLoader
Pure-Rust InstantTensor-style loader. Same public shape as crate::weights::SafetensorsLoader.

Constants§

DEFAULT_DIRECT_IO_TENSOR_CAP
Default tensor-count cap for per-shard O_DIRECT. Above this, the fast loader uses buffered reads even when try_direct_io = true. See the field doc on FastSafetensorsLoader::direct_io_tensor_cap.

Functions§

is_vision_tensor
Is this tensor part of a multimodal checkpoint’s vision tower?