CERIN AMROTH · ML systems

How do I run a MoE whose experts do not fit in VRAM, streaming them from pinned host RAM or an NVMe arena?

Solved by grouped-nf4-gemm · this page in the repository (pinned e2af4cfb91b2, the source of this rendering; latest on main, unpinned; not the source of any fact rendered here)

Install routes

From docs/capabilities.json at the pinned commit.

pip install grouped-nf4-gemm

Primary route: kernel package.

Alternatives:

Environment (from the capability register): OS: Linux · Python: >=3.11 tested in CI (pyproject says >=3.9; 3.9/3.10 are not tested) · Accelerator: NVIDIA CUDA GPU for the serving tiers and for the default NF4 quantise bake (nvme_bake_nf4.bake_nf4 with quantize_fn=None); the relocation bake and verify (nvme_arena.bake / bake_expert_tensors / verify) need no GPU · Requires: a local NVMe or fast block device for the arena; O_DIRECT reads; bitsandbytes and CUDA for the NF4 quantise bake (nvme_bake_nf4.bake_nf4 with the default quantiser); the geometry/manifest path is pure torch when a quantize_fn is injected; pinned host RAM sized from measured free memory (capacity_for_bytes); torch>=2.8

Role of this page: kernel/storage primitives. The arena bake and verifier, the O_DIRECT reader, the pinned-DRAM row tier and the low-level residency engines. It is not the decision page and not the model-level integration; those live in the consumer. An NVMe primitive (arena, reader, row tier, residency engine) belongs here; model-level NVMe integration — binding those primitives to a loaded model and deciding which bytes live where — belongs to experts4bit-qlora.

Use this page when… you are building or debugging the storage layer itself: baking an expert-major arena, sizing a pinned row tier, reading rows with ArenaReader, or wiring ColdTier / ArenaExpertSource under your own forward. Start in experts4bit-qlora instead when the question starts from a model:

grouped-nf4-gemm ships the tiers under a fused-expert forward: a GPU-driven gather from pinned host memory over UVA (host_gather.gather_expert_rows), an expert-major on-disk arena you bake once (nvme_arena, nvme_bake_nf4), an O_DIRECT reader (nvme_reader.ArenaReader), a pinned-DRAM residency tier over that arena (nvme_residency.ColdTier), and the MXFP4 engines that consume them (mxfp4_pipelined.Mxfp4PipelinedGptOss, mxfp4_residency.Mxfp4NvmeResidency). Which bytes live where is the consumer's decision.

Symptoms

Why it happens

Top-k routing touches a small fraction of expert bytes per token, so residency, not capacity, is the binding constraint; but a serving loop needs those bytes to arrive as one aligned request landing where the kernel will read them. Safetensors is tensor-major; the page cache duplicates a DRAM tier and hands eviction to the OS; a pinned row costs more host memory than its stride (nvme_residency.PINNED_ROW_FACTOR). The tiers here fix the layout at bake time and never copy a row more than the link requires.

Which project solves it

grouped-nf4-gemm owns the primitives: bake, verify, read, residency, gather, the writable KV row pool, and the two MXFP4 residency engines. experts4bit-qlora (PyPI) owns residency integration into a model, the NF4 offload path that serves a checkpoint from pinned host RAM on a small-VRAM card, and serving. To run a MoE larger than VRAM end to end, install the consumer; this page is the primitives it stands on.

Install

Kernel package (the minimum route):

pip install grouped-nf4-gemm

The relocation bake and verifier are stdlib plus the byte-range hasher and run on any host; the reader and ColdTier(pinned=False) are CPU-only. Pinned buffers, the gather and the engines need Linux with an NVIDIA GPU (sm_80+), triton>=3.4 (Linux-only), torch>=2.8 (pre-releases accepted); CI tests Python 3.11. nvme_bake_nf4 needs bitsandbytes and CUDA to quantize. Through the model consumer:

pip install "experts4bit-qlora[fast]"

Smallest correct example

CPU-only: bake a toy gpt-oss-shaped checkpoint into an arena, verify it against the source, make two experts resident through the mmap tier, and check the bytes.

# CPU-only (stdlib + torch import chain; no GPU, no triton). O_DIRECT falls back to
# buffered reads with a one-time warning on filesystems that refuse it.
import json, os, struct, tempfile
from mxfp4_loader import EXPERT_SUFFIXES
from nvme_arena import bake, verify, load_index
from nvme_residency import ColdTier, capacity_for_bytes

E, L = 4, 2
shapes = {EXPERT_SUFFIXES[0]: (8, 16), EXPERT_SUFFIXES[1]: (8, 4),   # gate_up blocks / scales
          EXPERT_SUFFIXES[2]: (6, 8), EXPERT_SUFFIXES[3]: (6, 2)}    # down blocks / scales
snap = tempfile.mkdtemp()
src, hdr, blobs, off = {}, {}, [], 0
for layer in range(L):
    for suf, (n, w) in shapes.items():
        name, data = f"model.layers.{layer}.{suf}", os.urandom(E * n * w)
        src[name] = data
        hdr[name] = {"dtype": "U8", "shape": [E, n, w], "data_offsets": [off, off + len(data)]}
        blobs.append(data); off += len(data)
hj = json.dumps(hdr).encode()
with open(os.path.join(snap, "model.safetensors"), "wb") as f:
    f.write(struct.pack("<Q", len(hj))); f.write(hj); f.write(b"".join(blobs))

arena = os.path.join(snap, "toy.arena")
bake(snap, arena, log=lambda *a: None)             # + toy.arena.index.json, toy.arena.manifest.json
assert verify(arena, against_source=snap, log=lambda *a: None)["ok"]

index = load_index(arena)
stride = index["row_stride"]
tier = ColdTier(arena, hot_rows=capacity_for_bytes(8 * stride, stride, pinned=False), pinned=False)
slots = tier.ensure(1, [3, 0])                     # layer 1, experts 3 and 0 -> slot indices
assert len(slots) == 2 and tier.resident(1, 3)
row = tier.row(1, 3)                               # the row's bytes, one memoryview
seg = {g["suffix"]: g for g in index["segments"]}[EXPERT_SUFFIXES[0]]
n, w = shapes[EXPERT_SUFFIXES[0]]
assert bytes(row[seg["seg_off"]:seg["seg_off"] + seg["length"]]) == src["model.layers.1." + EXPERT_SUFFIXES[0]][3 * n * w:4 * n * w]
tier.close()

GPU: the host-RAM primitive, a device-side gather from a pinned expert stack.

# GPU (sm_80+) + triton
import torch
from host_gather import gather_expert_rows

E, ROW = 64, 4096                                          # 64 expert rows of 4 KiB
host = torch.randint(0, 256, (E, ROW), dtype=torch.uint8).pin_memory()   # pinned host RAM
ids = torch.tensor([9, 3, 9, 41], dtype=torch.int32, device="cuda")
dst = torch.empty(ids.numel(), ROW, dtype=torch.uint8, device="cuda")
gather_expert_rows(dst, host, ids)                          # the GPU reads host memory over UVA
torch.cuda.synchronize()
assert torch.equal(dst.cpu(), host[ids.cpu().long()])

Expected result

The CPU block bakes, verifies the full chain (source range hash, manifest, arena bytes), and hands back a row whose segment equals the source slice. The GPU block copies exactly the requested rows with no CPU involvement in choosing them.

Supported scope

Limitations

nvme-ceilings.md · K3-PROVENANCE-CHAIN.md · RESULTS-nvme-determinism.md · STATUS.md · claims.json · native-mxfp4-moe-inference.md · verify-quantized-checkpoint-provenance.md

Evidence

Register: claims.json. Confirmed: claim gnf4.flagship.235b-phaseA (synthetic-weight offload at the link's waterfall ceiling), claim gnf4.flagship.235b-phaseB (the real checkpoint served from pinned host RAM across five pods, with the additive per-box law), claim gnf4.flagship.prefetch-closed-negative. Measured: claim gnf4.nvme.tier-batch-only, claim gnf4.k3.oracle-exact (arena round trip on a byte-verified store), claim gnf4.cold-engine.phase0-premise-refuted. Receipts: bench/phase3/flagship/RESULTS-flagship-phaseB.md, bench/phase3/flagship/RESULTS-flagship-offload.md, RESULTS-k3-slice-roundtrip.md.

Common wrong approaches

Source and freshness

This page is a rendering of docs/solutions/stream-moe-experts-from-host-or-nvme.md at commit e2af4cfb91b2 (sha256 0eca5f08a1d9bbb4…). Numbers are never copied here: every measured statement cites a claim ID in docs/claims.json at that commit. Repository-relative links resolve to this site's pages where the document is published and to the pinned commit otherwise.

CURRENT · source: pjordanandrsn/grouped-nf4-gemm@e2af4cfb91b2 · rendered package: 0.30.2 · latest published package: 0.30.2