NVMe-Aware Tensor Streaming for Local LLM Inference

Computation Graphs & ML Infrastructure advanced How AI Works ~35 min read

Overview

Local LLM inference usually assumes the model must fit in GPU memory, unified memory, or at least physical RAM. That assumption keeps the runtime simple: memory-map a GGUF file, offload as many layers as possible to the accelerator, and let the operating system page ordinary file-backed weights as needed. It breaks down when a developer tries to run a 31 GB or 40 GB quantized model on a 32 GB machine. A naive runtime can fall into swap thrash or fail before the first token.

NVMe-aware tensor streaming treats the model file as a cold weight store and makes residency an explicit runtime decision. The scheduler keeps latency-critical tensors on GPU, leaves some tensors in RAM, and streams the remaining weights from fast local storage into bounded pools just before they are needed. The important shift is that storage is not accidental swap. The inference engine knows which tensors are hot, which layers are upcoming, and which architecture-specific sparsity patterns allow it to avoid reading most of the file on each token.

Hypura is a concrete example of this pattern for Apple Silicon and GGUF models. Its README describes tensor placement across Metal GPU memory, RAM, and NVMe; direct read-only I/O with pread and F_NOCACHE; expert streaming for Mixtral-style MoE models; dense FFN streaming for oversized dense models; and hardware profiling that chooses pool sizes and prefetch depth. The transferable lesson is not the exact CLI. It is the runtime architecture required to make oversize local inference predictable instead of relying on the OS paging policy.

Core Concepts

Tensor Residency Tiers

The runtime divides tensors into placement tiers. GPU-resident tensors are the ones with high reuse or strong latency sensitivity, such as embeddings, normalization weights, attention weights, and small control tensors. RAM-resident tensors are file-backed or copied weights that do not fit in the accelerator working set but can be accessed without storage latency. NVMe-resident tensors are cold until the scheduler reads them into a pool buffer for one layer, one expert group, or one feed-forward block.

Architecture-Aware Weight Classification

Tensor streaming only works if the runtime understands the model structure. A transformer layer is not one opaque blob: attention, norms, embeddings, dense feed-forward matrices, router weights, and MoE expert tensors have different access patterns. Hypura classifies tensor roles from GGUF metadata and tensor names, then gives each role a different placement score. Small tensors that execute every token deserve permanent residency; very large FFN tensors may be streamed if their bandwidth cost is acceptable.

Storage as Cold Memory, Not Swap

Operating-system swap reacts after memory pressure has already happened. NVMe-aware inference schedules reads before compute reaches a missing tensor, uses bounded buffers instead of unbounded paging, and avoids polluting the filesystem cache when read-once tensor chunks would evict hotter data. That makes the storage tier part of the model execution plan rather than a failure mode.

Key Concept: Capacity Moves from RAM Size to I/O Schedule

An oversized model can become runnable when the engine reads only the tensors needed for the current execution window, but throughput becomes a function of NVMe bandwidth, prefetch accuracy, buffer sizing, and how much compute can overlap with storage reads.

How It Works

Build a Placement Plan from Model and Hardware Profiles

The runtime starts by reading model metadata and profiling local hardware. For Apple Silicon, the scheduler needs the Metal recommended working set size, available unified memory headroom, and measured NVMe throughput. It then assigns each tensor to GPU, RAM, or NVMe, trying to preserve accelerator speed for hot tensors while keeping total resident memory below conservative limits.

def plan_tensor_placement(model, hardware):
    budget = {
        "gpu": hardware.metal_working_set_bytes,
        "ram": hardware.safe_ram_bytes,
        "nvme": float("inf"),
    }

    plan = []
    for tensor in model.tensors:
        role = classify_tensor_role(tensor.name, model.architecture)
        score = residency_score(role, tensor.bytes, tensor.reuse_per_token)

        if score.pin_to_gpu and budget["gpu"] >= tensor.bytes:
            tier = "gpu"
            budget["gpu"] -= tensor.bytes
        elif score.keep_in_ram and budget["ram"] >= tensor.bytes:
            tier = "ram"
            budget["ram"] -= tensor.bytes
        else:
            tier = "nvme"

        plan.append({"tensor": tensor.name, "tier": tier, "role": role})

    return plan

Stream MoE Experts by Intercepting Router Decisions

Mixture-of-experts models create a special opportunity. The model may contain many expert tensors, but each token activates only a small subset. Hypura describes intercepting router decisions during evaluation, identifying the selected experts, and reading only the needed expert strides from NVMe. A neuron or expert cache keeps recently used slices in memory so temporal locality across adjacent tokens avoids repeated disk reads.

Stream Dense FFN Blocks with Prefetch Lookahead

Dense models do not have MoE sparsity, so the scheduler must stream predictable large blocks. A practical plan keeps attention and normalization weights resident on GPU, while gate, up, and down projection tensors for upcoming layers move through a fixed-size pool. Prefetch lookahead should scale with available memory: deeper lookahead hides more latency, but excessive lookahead can evict useful buffers and increase memory pressure.

Practical Applications

Running Oversized Models on Developer Workstations

The direct use case is local experimentation with models slightly larger than installed memory. A developer can run a quantized 70B-class model on hardware that would otherwise crash, accepting lower tokens per second in exchange for privacy, offline execution, and zero cloud inference cost. This is useful for evaluation, prompt iteration, and occasional tasks where latency is less important than model size.

MoE Inference on Memory-Constrained Machines

MoE models are attractive for tensor streaming because most experts are inactive for a given token. A runtime that pins shared weights and streams selected expert slices can exploit the same conditional-compute property that makes MoE efficient at cluster scale. The difference is that the limiting resource is local memory capacity rather than aggregate accelerator FLOPS.

Drop-In Local Serving Behind Existing Client APIs

Hypura exposes an Ollama-compatible HTTP API, which shows a practical deployment pattern: keep the unusual scheduling inside the local runtime, but preserve a familiar chat and generation protocol for clients. Tooling that already speaks the Ollama API can point at the new server without needing to understand tiered tensor placement.

Implementation Patterns

Use Explicit Buffer Pools for NVMe Reads

Allocate a bounded pool for streamed tensors rather than allocating on every layer. Each buffer has a lifecycle: free, reading, ready, computing, and reusable. The scheduler should issue asynchronous reads early enough that compute rarely waits, but it should also cap outstanding I/O so prefetches do not push the machine into memory pressure.

class StreamPool:
    def __init__(self, slots, slot_bytes):
        self.free = [bytearray(slot_bytes) for _ in range(slots)]
        self.ready = {}

    def prefetch(self, tensor, file_descriptor):
        if tensor.name in self.ready:
            return
        buf = self.free.pop()
        async_pread(file_descriptor, buf, tensor.file_offset, tensor.bytes)
        self.ready[tensor.name] = {"buffer": buf, "state": "reading"}

    def await_tensor(self, tensor):
        entry = self.ready[tensor.name]
        wait_until_read_complete(entry)
        entry["state"] = "computing"
        return entry["buffer"]

    def release(self, tensor):
        entry = self.ready.pop(tensor.name)
        self.free.append(entry["buffer"])

Separate Hot Tensor Pins from Streamed Tensor Windows

Do not let streamed blocks compete directly with permanently hot tensors. Keep separate budgets for pinned GPU tensors, RAM-resident overflow, and NVMe stream buffers. If a model fits in memory, the streaming path should add no overhead. If it does not fit, the engine should degrade into explicit streaming rather than allowing the OS to choose arbitrary victims.

Record Runtime Statistics for Future Scheduling

Co-activation, cache hit rate, tensor read latency, and per-layer stall time are scheduling inputs. Persisting small statistics files lets the next run choose better prefetch ordering and cache sizes. Keep this metadata tiny and separate from model files so normal inference remains read-only with respect to the large checkpoint.

Performance Considerations

Common Pitfalls

Further Reading