GPU Buffers

Note

This package is experimental.

GPU-backed raw storage and compression helpers (experimental).

Raw storage

Untyped GPU/CPU storage with a simple bump allocator.

class RawBuffer(size: int, device: str | device = 'cpu')

Bases: object

Linear byte storage with explicit allocation tracking.

Allocate size bytes on device.

Parameters:

• size – Total storage size in bytes.

• device – Torch device for the underlying storage.

malloc(length: int) → tuple[int, int]

Reserve length contiguous bytes and return (start, length).

write_bytes(malloc: tuple[int, int], tensor: Tensor)

Copy tensor bytes into the region described by malloc.

read_bytes(malloc: tuple[int, int], dtype: dtype)

View malloc bytes as a 1D tensor with dtype.

read_into(malloc: tuple[int, int], tensor: Tensor)

Copy bytes from malloc into the start of tensor.

Utilities

Torch dtype helpers for GPU buffer compression.

find_smallest_dtype(max_val: int, signed: bool = False, fallback: dtype = torch.float32) → dtype

Find smallest dtype for runs_type.

Compression methods

GPU compression method registry.