[Feature] Minimal safetensors read path over raw LBA storage

[CyberSecurityErial]

Motivation

uGDS currently exposes raw, block-aligned NVMe I/O to GPU memory, while safetensors describes tensors using names, dtypes, shapes, and logical byte ranges inside files. In model inference, checkpoint weights are normally immutable, so a small read-only mapping layer is sufficient; a general-purpose filesystem is not required.

The goal of this issue is a minimal runnable safetensors path that preserves the existing uGDS raw I/O design and keeps the code and test surface small.

Proposed MVP

1. Primary mapping: tensor metadata to logical file range

Parse a standard safetensors header and, when present, model.safetensors.index.json:

tensor name -> shard -> [file offset, length) + dtype + shape

The tensor payload starts at 8 + header_length + BEGIN, and its byte length is END - BEGIN. Header length, ranges, shape multiplication, duplicate keys, overlaps, and bounds must be validated.

2. Secondary mapping: logical file range to raw LBA

Import each canonical safetensors shard byte-for-byte into one contiguous, aligned raw object. A sidecar manifest records:

shard name -> raw object base LBA + canonical size + identity/hash

Then:

physical byte = object_base + file_offset
LBA range = outward-aligned physical byte range

This MVP deliberately does not use FIEMAP and does not support arbitrary filesystem files.

3. Tensor mapping

Read an aligned envelope into registered GPU memory and expose only the valid payload as a tensor view using the parsed dtype and shape. The central primitive should be read_into(name, destination); a convenience get_tensor(name) can allocate storage and call it. If direct placement is unsafe because of alignment or guard-byte requirements, use a small registered staging arena and a GPU copy.

Initial support is limited to complete, contiguous, byte-addressable tensors. Arbitrary strided slices and sub-byte dtypes are out of scope.

4. Read-only rules and safety

• Raw objects are immutable after commit. Model updates create a new object and switch the manifest generation.
• Validate device/namespace identity, LBA size, capacity, object bounds, canonical file size, and manifest/header hash before reading.
• A request owns its GPU buffer registration until all I/O completes. Close or deregister while in flight must fail or wait safely.
• Partial I/O, timeout, or validation failure must never publish a ready tensor.
• The importer writes payload first, flushes it, and publishes the committed manifest last.

5. Focused test cases

Keep the matrix intentionally small:

1. One single-shard file with a few representative byte-addressable dtypes.
2. One two-shard Hugging Face index case.
3. One unaligned tensor offset/length case exercising the aligned envelope.
4. One invalid header/range case.
5. One stale manifest or device mismatch case.
6. Differential tensor bytes, dtype, and shape against the official safetensors loader.

Non-goals for this MVP

• A general-purpose filesystem or POSIX path support.
• FIEMAP/fragmented extent support.
• In-place checkpoint writes or mutation.
• Optimized per-tensor repacking.
• Multi-GPU, tensor-parallel-aware packing, quantization transforms, or arbitrary slices.
• A new asynchronous engine; the existing synchronous/batch path is sufficient initially.
• Transparent integration with every framework in the first PR.

Implementation constraints

• Keep the existing raw I/O API behavior unchanged.
• Prefer a small optional safetensors layer and minimal new files.
• Avoid broad refactors before the MVP is proven.
• The implementation PR should begin with an RFC document, then proceed incrementally through the five sections above, with review between stages.

Acceptance criteria

A standard safetensors checkpoint can be imported as a read-only contiguous raw object, a named tensor can be located and read through uGDS into GPU memory, and its dtype, shape, and bytes match the official loader. Existing uGDS tests remain unchanged and passing.